From 3b058d490d1af6fed13df829230655192c4c83bb Mon Sep 17 00:00:00 2001 From: ihujun Date: Mon, 10 Aug 2026 20:06:07 +0800 Subject: [PATCH] feat: add NPU SAS attention and LI patch for DeepSeek-V4 without mindspeed --- pyproject.toml | 1 + src/twinkle/kernel/config.py | 6 + src/twinkle/kernel/ops/__init__.py | 2 +- .../kernel/ops/dsv4_sas_li/__init__.py | 65 ++ .../kernel/ops/dsv4_sas_li/aclnn/__init__.py | 9 + .../ops/dsv4_sas_li/aclnn/_aclnn_common.h | 707 ++++++++++++++++++ .../kernel/ops/dsv4_sas_li/aclnn/builder.py | 61 ++ .../aclnn/lightning_indexer/binding.cpp | 52 ++ .../aclnn/lightning_indexer_grad/binding.cpp | 51 ++ .../aclnn/sparse_attn_sharedkv/binding.cpp | 163 ++++ .../kernel/ops/dsv4_sas_li/aclnn_ops.py | 129 ++++ src/twinkle/kernel/ops/dsv4_sas_li/npu.py | 334 +++++++++ tests/kernel/ops/test_aclnn_ops.py | 55 ++ 13 files changed, 1634 insertions(+), 1 deletion(-) create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/__init__.py create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/__init__.py create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/_aclnn_common.h create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/builder.py create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer/binding.cpp create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer_grad/binding.cpp create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn/sparse_attn_sharedkv/binding.cpp create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/aclnn_ops.py create mode 100644 src/twinkle/kernel/ops/dsv4_sas_li/npu.py create mode 100644 tests/kernel/ops/test_aclnn_ops.py diff --git a/pyproject.toml b/pyproject.toml index 12fb2d929..27a720456 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,3 +80,4 @@ where = ["src"] [tool.setuptools.package-data] "twinkle_client.skills.bundled" = ["*.md"] +"twinkle.kernel.ops.dsv4_sas_li.aclnn" = ["*.h", "*.cpp", "**/*.cpp"] diff --git a/src/twinkle/kernel/config.py b/src/twinkle/kernel/config.py index 67d623e7a..dc113a231 100644 --- a/src/twinkle/kernel/config.py +++ b/src/twinkle/kernel/config.py @@ -111,6 +111,12 @@ def _build() -> dict[Any, Any]: # logical target: handled by a custom installer (never resolved by the generic replacer) cfg['sdpa'] = KernelChoice(op='sdpa_attention', backends=('npu', )) cfg['fla'] = KernelChoice(op='fla', backends=('npu', )) + + # DeepSeek-V4 SAS + LI (env-gated by TWINKLE_NPU_DSV4_SAS, npu only) + _dsv4 = 'transformers.models.deepseek_v4.modeling_deepseek_v4' + cfg[f'{_dsv4}.DeepseekV4Attention.forward'] = KernelChoice(op='dsv4_attention', backends=('npu', )) + cfg[f'{_dsv4}.DeepseekV4Indexer.forward'] = KernelChoice(op='dsv4_indexer', backends=('npu', )) + cfg[f'{_dsv4}.DeepseekV4CSACompressor.forward'] = KernelChoice(op='dsv4_csa_compressor', backends=('npu', )) return cfg diff --git a/src/twinkle/kernel/ops/__init__.py b/src/twinkle/kernel/ops/__init__.py index 14ce9cd1e..c05205e65 100644 --- a/src/twinkle/kernel/ops/__init__.py +++ b/src/twinkle/kernel/ops/__init__.py @@ -12,7 +12,7 @@ references + availability checks only, no optional-dependency imports). """ # Trigger built-in op registration (must happen before the first kernelize() call) -from . import fla, geglu, moe, rms_norm, rotary, sdpa_attention, swiglu # noqa: F401,E402 +from . import dsv4_sas_li, fla, geglu, moe, rms_norm, rotary, sdpa_attention, swiglu # noqa: F401,E402 from .ep import EpExpertsGmm, ep_forward __all__ = [ diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/__init__.py b/src/twinkle/kernel/ops/dsv4_sas_li/__init__.py new file mode 100644 index 000000000..850e32285 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4 SAS/LI op registration: three forward-level replacements gated +by the ``TWINKLE_NPU_DSV4_SAS`` env var. + +When enabled, the full patch set is applied: + + - ``DeepseekV4Attention.forward`` → NPU sparse attention (SAS) + - ``DeepseekV4Indexer.forward`` → Lightning Indexer (LI) + - ``DeepseekV4CSACompressor.forward`` → full replacement returning a + 3-tuple ``(compressed_kv, block_bias, top_k_indices)`` + +LI is always on under SAS — there is no use case for SAS without LI +(CSA would fall back to the slower stock indexer) or LI without SAS +(indices would go unused). The CSA compressor is a **full forward +replacement** rather than a wrapper (see ``npu.py`` docstring for details). +""" +from __future__ import annotations + +import os + +from ...registry import KernelImpl, is_npu_available, lazy_import, register_op + + +def _dsv4_sas_available() -> tuple[bool, str | None]: + env = os.environ.get('TWINKLE_NPU_DSV4_SAS', '').lower().strip() + if not env or env in ('0', 'false', 'off', 'no'): + return False, 'TWINKLE_NPU_DSV4_SAS not enabled' + ok, reason = is_npu_available() + if not ok: + return ok, reason + return True, None + + +_DSV4_BASE = 'twinkle.kernel.ops.dsv4_sas_li.npu' + +register_op( + 'dsv4_attention', + implementations={ + 'npu': KernelImpl( + load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_attention_forward'), + available=_dsv4_sas_available, + ), + }, +) + +register_op( + 'dsv4_indexer', + implementations={ + 'npu': KernelImpl( + load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_indexer_forward'), + available=_dsv4_sas_available, + ), + }, +) + +register_op( + 'dsv4_csa_compressor', + implementations={ + 'npu': + KernelImpl( + load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_csa_compressor_forward'), + available=_dsv4_sas_available, + ), + }, +) diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/__init__.py b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/__init__.py new file mode 100644 index 000000000..2a237b338 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Self-compiled ACLNN C++ extensions for Ascend NPU fusion operators. + +Provides JIT-compiled bindings for DeepSeek-V4 SAS (Sparse Attention with +Shared-KV) and LI (Lightning Indexer) without depending on mindspeed. +""" +from .builder import build_op + +__all__ = ['build_op'] diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/_aclnn_common.h b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/_aclnn_common.h new file mode 100644 index 000000000..e186b7ae5 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/_aclnn_common.h @@ -0,0 +1,707 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#ifndef _ACLNN_COMMON_H +#define _ACLNN_COMMON_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/framework/OpCommand.h" +#include "torch_npu/csrc/framework/interface/EnvVariables.h" +#include "torch_npu/csrc/aten/NPUNativeFunctions.h" +#include "torch_npu/csrc/core/npu/DeviceUtils.h" +#if __has_include("torch_npu/csrc/flopcount/FlopCount.h") + #include "torch_npu/csrc/flopcount/FlopCount.h" +#endif +#define NPU_NAME_SPACE at_npu::native + +using aclOpExecutor = struct aclOpExecutor; +using aclTensor = struct aclTensor; +using aclScalar = struct aclScalar; +using aclIntArray = struct aclIntArray; +using aclFloatArray = struct aclFloatArray; +using aclBoolArray = struct aclBoolArray; +using aclTensorList = struct aclTensorList; + +using _aclCreateTensor = aclTensor *(*)(const int64_t *view_dims, uint64_t view_dims_num, aclDataType data_type, + const int64_t *stride, int64_t offset, aclFormat format, const int64_t *storage_dims, uint64_t storage_dims_num, + void *tensor_data); +using _aclCreateScalar = aclScalar *(*)(void *value, aclDataType data_type); +using _aclCreateIntArray = aclIntArray *(*)(const int64_t *value, uint64_t size); +using _aclCreateFloatArray = aclFloatArray *(*)(const float *value, uint64_t size); +using _aclCreateBoolArray = aclBoolArray *(*)(const bool *value, uint64_t size); +using _aclCreateTensorList = aclTensorList *(*)(const aclTensor *const *value, uint64_t size); + +using _aclDestroyTensor = int (*)(const aclTensor *tensor); +using _aclDestroyScalar = int (*)(const aclScalar *scalar); +using _aclDestroyIntArray = int (*)(const aclIntArray *array); +using _aclDestroyFloatArray = int (*)(const aclFloatArray *array); +using _aclDestroyBoolArray = int (*)(const aclBoolArray *array); +using _aclDestroyTensorList = int (*)(const aclTensorList *array); + +constexpr int kHashBufSize = 8192; +constexpr int kHashBufMaxSize = kHashBufSize + 1024; +extern thread_local char g_hashBuf[kHashBufSize]; +extern thread_local int g_hashOffset; + +#define AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(_) \ + _(at::ScalarType::Byte, ACL_UINT8) \ + _(at::ScalarType::Char, ACL_INT8) \ + _(at::ScalarType::Short, ACL_INT16) \ + _(at::ScalarType::Int, ACL_INT32) \ + _(at::ScalarType::Long, ACL_INT64) \ + _(at::ScalarType::Half, ACL_FLOAT16) \ + _(at::ScalarType::Float, ACL_FLOAT) \ + _(at::ScalarType::Double, ACL_DOUBLE) \ + _(at::ScalarType::ComplexHalf, ACL_COMPLEX32) \ + _(at::ScalarType::ComplexFloat, ACL_COMPLEX64) \ + _(at::ScalarType::ComplexDouble, ACL_COMPLEX128) \ + _(at::ScalarType::Bool, ACL_BOOL) \ + _(at::ScalarType::QInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QInt32, ACL_DT_UNDEFINED) \ + _(at::ScalarType::BFloat16, ACL_BF16) \ + _(at::ScalarType::QUInt4x2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::QUInt2x4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Bits1x8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Bits2x4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Bits4x2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Bits8, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Bits16, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Float8_e5m2, ACL_FLOAT8_E5M2) \ + _(at::ScalarType::Float8_e4m3fn, ACL_FLOAT8_E4M3FN) \ + _(at::ScalarType::Float8_e5m2fnuz, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Float8_e4m3fnuz, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt16, ACL_UINT16) \ + _(at::ScalarType::UInt32, ACL_UINT32) \ + _(at::ScalarType::UInt64, ACL_UINT64) \ + _(at::ScalarType::UInt1, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt3, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt5, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt6, ACL_DT_UNDEFINED) \ + _(at::ScalarType::UInt7, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int1, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int2, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int3, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int4, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int5, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int6, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Int7, ACL_DT_UNDEFINED) \ + _(at::ScalarType::Float8_e8m0fnu, ACL_FLOAT8_E8M0) \ + _(at::ScalarType::Undefined, ACL_DT_UNDEFINED) \ + _(at::ScalarType::NumOptions, ACL_DT_UNDEFINED) + +constexpr aclDataType kATenScalarTypeToAclDataTypeTable[static_cast(at::ScalarType::NumOptions) + 1] = { +#define DEFINE_ENUM(_1, n) n, + AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(DEFINE_ENUM) +#undef DEFINE_ENUM +}; + +enum QuantMode { + QUANT_MODE_NO_QUANT = 0, + QUANT_MODE_STATIC = 1, + QUANT_MODE_PERTOKEN = 2, + QUANT_MODE_PERGROUP = 3, + QUANT_MODE_MX = 4, +}; + +#define GET_OP_API_FUNC(apiName) reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName)) + +#define MEMCPY_TO_BUF(data_expression, size_expression) \ + if (g_hashOffset + (size_expression) > kHashBufSize) { \ + g_hashOffset = kHashBufMaxSize; \ + return; \ + } \ + auto ret = memcpy_s(g_hashBuf + g_hashOffset, size_expression, data_expression, size_expression); \ + TORCH_CHECK(ret == 0, "memcpy_s failed, error:", ret); \ + g_hashOffset += size_expression; + +inline const char *GetOpApiLibName(void) +{ + return "libopapi.so"; +} + +inline const char *GetCustOpApiLibName(void) +{ + return "libcust_opapi.so"; +} + +inline void *GetOpApiFuncAddrInLib(void *handler, const char *libName, const char *apiName) +{ + auto funcAddr = dlsym(handler, apiName); + if (funcAddr == nullptr) { + ASCEND_LOGW("dlsym %s from %s failed, error:%s.", apiName, libName, dlerror()); + } + return funcAddr; +} + +inline void *GetOpApiLibHandler(const char *libName) +{ + auto handler = dlopen(libName, RTLD_LAZY); + if (handler == nullptr) { + ASCEND_LOGW("dlopen %s failed, error:%s.", libName, dlerror()); + } + return handler; +} + +inline std::vector GetCustOpApiHandlers() +{ + std::vector handlers; + const char *env = std::getenv("ASCEND_CUSTOM_OPP_PATH"); + if (env != nullptr) { + std::string envStr(env); + std::istringstream iss(envStr); + std::string path; + while (std::getline(iss, path, ':')) { + if (path.empty()) { + continue; + } + std::string soPath = path + "/op_api/lib/" + GetCustOpApiLibName(); + auto handler = dlopen(soPath.c_str(), RTLD_LAZY); + if (handler != nullptr) { + handlers.push_back(handler); + } else { + ASCEND_LOGW("dlopen %s failed, error:%s.", soPath.c_str(), dlerror()); + } + } + } + if (handlers.empty()) { + auto handler = GetOpApiLibHandler(GetCustOpApiLibName()); + if (handler != nullptr) { + handlers.push_back(handler); + } + } + return handlers; +} + +inline void *GetOpApiFuncAddr(const char *apiName) +{ + static auto custHandlers = GetCustOpApiHandlers(); + for (auto handler : custHandlers) { + auto funcAddr = GetOpApiFuncAddrInLib(handler, GetCustOpApiLibName(), apiName); + if (funcAddr != nullptr) { + return funcAddr; + } + } + + static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName()); + if (opApiHandler == nullptr) { + return nullptr; + } + return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName); +} + +inline c10::Scalar ConvertTensorToScalar(const at::Tensor &tensor) +{ + c10::Scalar expScalar; + const at::Tensor *aclInput = &tensor; + if (aclInput->scalar_type() == at::ScalarType::Double) { + double value = *(double *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Long) { + int64_t value = *(int64_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Float) { + float value = *(float *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Int) { + int value = *(int *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Half) { + c10::Half value = *(c10::Half *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::Bool) { + int8_t value = *(int8_t *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexDouble) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::ComplexFloat) { + c10::complex value = *(c10::complex *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else if (aclInput->scalar_type() == at::ScalarType::BFloat16) { + c10::BFloat16 value = *(c10::BFloat16 *)aclInput->data_ptr(); + c10::Scalar scalar(value); + expScalar = scalar; + } else { + ASCEND_LOGE("unsupported scalar type! "); + } + return expScalar; +} + +inline at::Tensor CopyTensorHostToDevice(const at::Tensor &cpu_tensor) +{ + at::Tensor cpuPinMemTensor = cpu_tensor.pin_memory(); + int deviceIndex = 0; + return cpuPinMemTensor.to( + c10::Device(torch_npu::utils::get_npu_device_type(), deviceIndex), cpuPinMemTensor.scalar_type(), true, true); +} + +inline at::Tensor CopyScalarToDevice(const c10::Scalar &cpu_scalar, at::ScalarType scalar_data_type) +{ + return CopyTensorHostToDevice(scalar_to_tensor(cpu_scalar).to(scalar_data_type)); +} + +inline aclTensor *ConvertType(const at::Tensor &at_tensor) +{ + static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor); + if (aclCreateTensor == nullptr) { + return nullptr; + } + + if (!at_tensor.defined()) { + return nullptr; + } + at::ScalarType scalar_data_type = at_tensor.scalar_type(); + aclDataType acl_data_type = kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, std::string(c10::toString(scalar_data_type)) + " has not been supported") + c10::SmallVector storageDims; + // if acl_data_type is ACL_STRING, storageDims is empty. + auto itemsize = at_tensor.itemsize(); + if (itemsize == 0) { + AT_ERROR("When ConvertType, tensor item size of cannot be zero."); + return nullptr; + } + if (acl_data_type != ACL_STRING) { + storageDims.push_back(at_tensor.storage().nbytes() / itemsize); + } + + const auto dimNum = at_tensor.sizes().size(); + aclFormat format = ACL_FORMAT_ND; + switch (dimNum) { + case 3: + format = ACL_FORMAT_NCL; + break; + case 4: + format = ACL_FORMAT_NCHW; + break; + case 5: + format = ACL_FORMAT_NCDHW; + break; + default: + format = ACL_FORMAT_ND; + } + + if (at_tensor.unsafeGetTensorImpl()->is_wrapped_number()) { + c10::Scalar expScalar = ConvertTensorToScalar(at_tensor); + at::Tensor aclInput = CopyScalarToDevice(expScalar, scalar_data_type); + return aclCreateTensor(aclInput.sizes().data(), + aclInput.sizes().size(), + acl_data_type, + aclInput.strides().data(), + aclInput.storage_offset(), + format, + storageDims.data(), + storageDims.size(), + const_cast(aclInput.storage().data())); + } + + auto acl_tensor = aclCreateTensor(at_tensor.sizes().data(), + at_tensor.sizes().size(), + acl_data_type, + at_tensor.strides().data(), + at_tensor.storage_offset(), + format, + storageDims.data(), + storageDims.size(), + const_cast(at_tensor.storage().data())); + return acl_tensor; +} + +inline aclScalar *ConvertType(const at::Scalar &at_scalar) +{ + static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar); + if (aclCreateScalar == nullptr) { + return nullptr; + } + + at::ScalarType scalar_data_type = at_scalar.type(); + aclDataType acl_data_type = kATenScalarTypeToAclDataTypeTable[static_cast(scalar_data_type)]; + TORCH_CHECK( + acl_data_type != ACL_DT_UNDEFINED, std::string(c10::toString(scalar_data_type)) + " has not been supported") + aclScalar *acl_scalar = nullptr; + switch (scalar_data_type) { + case at::ScalarType::Double: { + double value = at_scalar.toDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Long: { + int64_t value = at_scalar.toLong(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::Bool: { + bool value = at_scalar.toBool(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + case at::ScalarType::ComplexDouble: { + auto value = at_scalar.toComplexDouble(); + acl_scalar = aclCreateScalar(&value, acl_data_type); + break; + } + default: + acl_scalar = nullptr; + break; + } + return acl_scalar; +} + +inline aclIntArray *ConvertType(const at::IntArrayRef &at_array) +{ + static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray); + if (aclCreateIntArray == nullptr) { + return nullptr; + } + auto array = aclCreateIntArray(at_array.data(), at_array.size()); + return array; +} + +template +inline aclBoolArray *ConvertType(const std::array &value) +{ + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclBoolArray *ConvertType(const at::ArrayRef &value) +{ + static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray); + if (aclCreateBoolArray == nullptr) { + return nullptr; + } + + auto array = aclCreateBoolArray(value.data(), value.size()); + return array; +} + +inline aclTensorList *ConvertType(const at::TensorList &at_tensor_list) +{ + if (at_tensor_list.size() == 0) { + return nullptr; + } + static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList); + if (aclCreateTensorList == nullptr) { + return nullptr; + } + + std::vector tensor_list(at_tensor_list.size()); + for (size_t i = 0; i < at_tensor_list.size(); i++) { + tensor_list[i] = ConvertType(at_tensor_list[i]); + } + auto acl_tensor_list = aclCreateTensorList(tensor_list.data(), tensor_list.size()); + return acl_tensor_list; +} + +inline aclTensor *ConvertType(const c10::optional &opt_tensor) +{ + if (opt_tensor.has_value() && opt_tensor.value().defined()) { + return ConvertType(opt_tensor.value()); + } + return nullptr; +} + +inline aclIntArray *ConvertType(const c10::optional &opt_array) +{ + if (opt_array.has_value()) { + return ConvertType(opt_array.value()); + } + return nullptr; +} + +inline aclScalar *ConvertType(const c10::optional &opt_scalar) +{ + if (opt_scalar.has_value()) { + return ConvertType(opt_scalar.value()); + } + return nullptr; +} + +inline aclDataType ConvertType(const at::ScalarType scalarType) +{ + return kATenScalarTypeToAclDataTypeTable[static_cast(scalarType)]; +} + +template +T ConvertType(T value) +{ + return value; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr, std::index_sequence) +{ + using OpApiFunc = int (*)(typename std::decay(params))>::type...); + auto func = reinterpret_cast(opApiAddr); + return func; +} + +template +auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr) +{ + static constexpr auto size = std::tuple_size::value; + return ConvertToOpApiFunc(params, opApiAddr, std::make_index_sequence{}); +} + +inline void Release(aclTensor *p) +{ + static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor); + if (aclDestroyTensor == nullptr) { + return; + } + aclDestroyTensor(p); +} + +inline void Release(aclScalar *p) +{ + static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar); + if (aclDestroyScalar == nullptr) { + return; + } + aclDestroyScalar(p); +} + +inline void Release(aclIntArray *p) +{ + static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray); + if (aclDestroyIntArray == nullptr) { + return; + } + + aclDestroyIntArray(p); +} + +inline void Release(aclBoolArray *p) +{ + static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray); + if (aclDestroyBoolArray == nullptr) { + return; + } + + aclDestroyBoolArray(p); +} + +inline void Release(aclTensorList *p) +{ + static const auto aclDestroyTensorList = GET_OP_API_FUNC(aclDestroyTensorList); + if (aclDestroyTensorList == nullptr) { + return; + } + + aclDestroyTensorList(p); +} + +template +void Release(T value) +{ + (void)value; +} + +template +void CallRelease(Tuple t, std::index_sequence) +{ + (void)std::initializer_list{(Release(std::get(t)), 0)...}; +} + +template +void ReleaseConvertTypes(Tuple &t) +{ + static constexpr auto size = std::tuple_size::value; + CallRelease(t, std::make_index_sequence{}); +} + +template +constexpr auto ConvertTypes(Ts &...args) +{ + return std::make_tuple(ConvertType(args)...); +} + +template +auto call(Function f, Tuple t, std::index_sequence) +{ + return f(std::get(t)...); +} + +template +auto call(Function f, Tuple t) +{ + static constexpr auto size = std::tuple_size::value; + return call(f, t, std::make_index_sequence{}); +} + +template +void AddParamToBuf(const std::array &value) +{ + MEMCPY_TO_BUF(value.data(), value.size() * sizeof(bool)); +} + +template +void AddParamToBuf(const T &value) +{ + MEMCPY_TO_BUF(&value, sizeof(T)); +} + +void AddParamToBuf(const at::Tensor &); +void AddParamToBuf(const at::Scalar &); +void AddParamToBuf(const at::IntArrayRef &); +void AddParamToBuf(const at::ArrayRef &); +void AddParamToBuf(const at::TensorList &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const c10::optional &); +void AddParamToBuf(const at::ScalarType); +void AddParamToBuf(const string &); +void AddParamToBuf(); + +template +void AddParamToBuf(const T &arg, Args &...args) +{ + AddParamToBuf(arg); + AddParamToBuf(args...); +} + +uint64_t CalcHashId(); +using InitHugeMemThreadLocal = int (*)(void *, bool); +using UnInitHugeMemThreadLocal = void (*)(void *, bool); +using ReleaseHugeMem = void (*)(void *, bool); + +/** + * check arg is at::Tensor ? + */ +template +struct is_at_tensor : std::false_type {}; + +template<> +struct is_at_tensor : std::true_type {}; + +/** + * check arg is at::TensorList ? + */ +template +struct is_at_tensor_list : std::false_type {}; + +template<> +struct is_at_tensor_list : std::true_type {}; + +/** + * find first at::Tensor + */ +template +typename std::enable_if::type GetFirstTensor(const std::tuple& t, at::Tensor& res) {} + +template +typename std::enable_if < I::type GetFirstTensor(const std::tuple &t, at::Tensor &res) +{ + if constexpr (is_at_tensor>::type>::value) { + res = std::get(t); + return; + } else if constexpr (is_at_tensor_list>::type>::value) { + res = std::get(t)[0]; + return; + } + return GetFirstTensor(t, res); +} + +/** + * get the device + */ +template +auto DecodeDevice(Ts&... args) -> at::Device +{ + auto tp = std::make_tuple(args...); + at::Tensor ft; + GetFirstTensor(tp, ft); + return ft.device(); +} + +#define ACLNN_CMD(aclnn_api, ...) \ + do { \ + auto device = DecodeDevice(__VA_ARGS__); \ + const c10::OptionalDeviceGuard device_guard(device); \ + static const auto getWorkspaceSizeFuncAddr = GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \ + static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \ + static const auto initMemAddr = GetOpApiFuncAddr("InitHugeMemThreadLocal"); \ + static const auto unInitMemAddr = GetOpApiFuncAddr("UnInitHugeMemThreadLocal"); \ + static const auto releaseMemAddr = GetOpApiFuncAddr("ReleaseHugeMem"); \ + TORCH_CHECK(getWorkspaceSizeFuncAddr != nullptr && opApiFuncAddr != nullptr, \ + #aclnn_api, \ + " or ", \ + #aclnn_api "GetWorkspaceSize", \ + " not in ", \ + GetOpApiLibName(), \ + ", or ", \ + GetOpApiLibName(), \ + "not found."); \ + auto acl_stream = c10_npu::getCurrentNPUStream().stream(false); \ + uint64_t workspace_size = 0; \ + uint64_t *workspace_size_addr = &workspace_size; \ + aclOpExecutor *executor = nullptr; \ + aclOpExecutor **executor_addr = &executor; \ + InitHugeMemThreadLocal initMemFunc = reinterpret_cast(initMemAddr); \ + UnInitHugeMemThreadLocal unInitMemFunc = reinterpret_cast(unInitMemAddr); \ + if (initMemFunc) { \ + initMemFunc(nullptr, false); \ + } \ + auto converted_params = ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \ + static auto getWorkspaceSizeFunc = ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \ + auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \ + TORCH_CHECK(workspace_status == 0, "call " #aclnn_api " failed, detail:", aclGetRecentErrMsg()); \ + at::Tensor workspace_tensor; \ + void *workspace_addr = nullptr; \ + if (workspace_size != 0) { \ + at::TensorOptions options = at::TensorOptions(torch_npu::utils::get_npu_device_type()); \ + workspace_tensor = at::empty({workspace_size}, options.dtype(at::kByte)); \ + workspace_addr = const_cast(workspace_tensor.storage().data()); \ + } \ + auto acl_call = [converted_params, workspace_addr, workspace_size, acl_stream, executor]() -> int { \ + typedef int (*OpApiFunc)(void *, uint64_t, aclOpExecutor *, const aclrtStream); \ + OpApiFunc opApiFunc = reinterpret_cast(opApiFuncAddr); \ + auto api_ret = opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \ + TORCH_CHECK(api_ret == 0, "call " #aclnn_api " failed, detail:", aclGetRecentErrMsg()); \ + ReleaseConvertTypes(converted_params); \ + ReleaseHugeMem releaseMemFunc = reinterpret_cast(releaseMemAddr); \ + if (releaseMemFunc) { \ + releaseMemFunc(nullptr, false); \ + } \ + return api_ret; \ + }; \ + at_npu::native::OpCommand cmd; \ + cmd.Name(#aclnn_api); \ + cmd.SetCustomHandler(acl_call); \ + cmd.Run(); \ + if (unInitMemFunc) { \ + unInitMemFunc(nullptr, false); \ + } \ + } while (false) + +#endif // _ACLNN_COMMON_H diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/builder.py b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/builder.py new file mode 100644 index 000000000..8679625b9 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/builder.py @@ -0,0 +1,61 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""JIT-compile ACLNN C++ binding extensions on first use. + +Adapted from ``torchtitan_npu.ops.aclnn.builder``. Each ``build_op`` call +loads (and caches) a pybind11 extension whose ``binding.cpp`` invokes the +ACLNN kernel macro ``ACLNN_CMD``. +""" +import os +import torch + +_loaded_ops = {} + + +def build_op(name, sources, verbose=False): + if name in _loaded_ops: + return _loaded_ops[name] + + cwd = os.path.dirname(os.path.abspath(__file__)) + ascend_home = os.environ.get('ASCEND_HOME_PATH', '/usr/local/Ascend/cann') + + try: + import torch_npu + tnpu_path = os.path.dirname(os.path.abspath(torch_npu.__file__)) + except ImportError: + raise ImportError('torch_npu is required to compile ACLNN extensions. ' + 'Install it for your Ascend NPU environment.') + + torch_path = os.path.dirname(os.path.abspath(torch.__file__)) + + cflags = [ + '-D_FORTIFY_SOURCE=2', + '-O2', + f'-I{cwd}', + f'-I{ascend_home}/include', + f'-I{torch_path}/include', + f'-I{torch_path}/include/torch/csrc/api/include', + f'-I{tnpu_path}/include', + f'-I{tnpu_path}/include/torch_npu/csrc/framework/utils', + f'-I{tnpu_path}/include/torch_npu/csrc/aten', + f'-I{tnpu_path}/include/third_party/hccl/inc', + f'-I{tnpu_path}/include/third_party/acl/inc', + ] + + ldflags = [ + f'-L{ascend_home}/lib64', + f'-L{tnpu_path}/lib', + '-ltorch_npu', + '-lascendcl', + ] + + from torch.utils.cpp_extension import load + + op_module = load( + name=name, + sources=[os.path.join(cwd, s) for s in sources], + extra_cflags=cflags, + extra_ldflags=ldflags, + verbose=verbose, + ) + _loaded_ops[name] = op_module + return op_module diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer/binding.cpp b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer/binding.cpp new file mode 100644 index 000000000..cb99540b3 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer/binding.cpp @@ -0,0 +1,52 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include "../_aclnn_common.h" + +std::tuple npu_lightning_indexer( + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &weights, + const c10::optional &actual_seq_q, + const c10::optional &actual_seq_k, + const c10::optional &block_table, + std::string layout_q, + std::string layout_k, + int64_t sparse_count, + int64_t sparse_mode, + int64_t pre_tokens, + int64_t next_tokens, + int64_t cmp_ratio, + bool return_values) +{ + TORCH_CHECK(query.numel() > 0, "query is empty"); + + const char *layout_q_ptr = layout_q.c_str(); + const char *layout_k_ptr = layout_k.c_str(); + + auto q_sizes = query.sizes(); + int64_t B = q_sizes[0]; + int64_t S1 = q_sizes[1]; + int64_t N2 = key.size(2); + + auto opts_int = at::TensorOptions().dtype(at::kInt).device(query.device()); + at::Tensor sparse_indices = at::empty({B, S1, N2, sparse_count}, opts_int); + at::Tensor sparse_values = at::empty({B, S1, N2, sparse_count}, query.options()); + + ACLNN_CMD(aclnnLightningIndexer, + query, key, weights, + actual_seq_q, actual_seq_k, block_table, + layout_q_ptr, layout_k_ptr, + sparse_count, sparse_mode, pre_tokens, next_tokens, + cmp_ratio, return_values, + sparse_indices, sparse_values); + + return {sparse_indices, sparse_values}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("npu_lightning_indexer", &npu_lightning_indexer, "Lightning Indexer"); +} diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer_grad/binding.cpp b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer_grad/binding.cpp new file mode 100644 index 000000000..3fa26f6cc --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/lightning_indexer_grad/binding.cpp @@ -0,0 +1,51 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include "../_aclnn_common.h" + +std::tuple npu_lightning_indexer_grad( + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &dy, + const at::Tensor &sparse_indices, + const at::Tensor &weights, + const c10::optional &actual_seq_lengths_query, + const c10::optional &actual_seq_lengths_key, + const c10::optional layout, + c10::optional sparse_mode, + c10::optional pre_tokens, + c10::optional next_tokens, + c10::optional cmp_ratio, + c10::optional head_num) +{ + at::Tensor d_query = at::zeros(query.sizes(), query.options()); + at::Tensor d_key = at::zeros(key.sizes(), key.options()); + at::Tensor d_weights = at::zeros(weights.sizes(), weights.options()); + + std::string layout_str = layout.value_or("BSND"); + char *layout_ptr = const_cast(layout_str.c_str()); + const int64_t sparse_mode_val = sparse_mode.value_or(0); + const int64_t pre_tokens_val = pre_tokens.value_or(9223372036854775807LL); + const int64_t next_tokens_val = next_tokens.value_or(9223372036854775807LL); + const int64_t cmp_ratio_val = cmp_ratio.value_or(1); + const int64_t head_num_val = head_num.value_or(64); + const bool deterministic = false; + + ACLNN_CMD(aclnnLightningIndexerGrad, + query, key, dy, sparse_indices, weights, + actual_seq_lengths_query, actual_seq_lengths_key, + head_num_val, layout_ptr, + sparse_mode_val, pre_tokens_val, next_tokens_val, + deterministic, cmp_ratio_val, + d_query, d_key, d_weights); + + return std::make_tuple(d_query, d_key, d_weights); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("npu_lightning_indexer_grad", &npu_lightning_indexer_grad, + "Lightning Indexer Backward"); +} diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/sparse_attn_sharedkv/binding.cpp b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/sparse_attn_sharedkv/binding.cpp new file mode 100644 index 000000000..e1796841e --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn/sparse_attn_sharedkv/binding.cpp @@ -0,0 +1,163 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include +#include "../_aclnn_common.h" + +at::Tensor npu_sparse_attn_sharedkv_metadata( + const c10::optional &cuSeqLensQ, + const c10::optional &sequsedOriKv, + const c10::optional &sequsedCmpKv, + const c10::optional &sequsedQ, + const c10::optional &sequsedKv, + int64_t numHeadsQ, + int64_t numHeadsKv, + int64_t headDim, + int64_t batchSize, + int64_t maxSeqLenQ, + int64_t maxSeqLenKv, + int64_t oriTopk, + int64_t cmpTopk, + int64_t cmpRatio, + int64_t oriMaskMode, + int64_t cmpMaskMode, + int64_t oriWinLeft, + int64_t oriWinRight, + const c10::optional layoutQ, + const c10::optional layoutKv, + bool hasOriKv, + bool hasCmpKv) +{ + std::string layoutQStr = layoutQ.value_or("SBH"); + std::string layoutKvStr = layoutKv.value_or("SBH"); + const char *layoutQPtr = layoutQStr.c_str(); + const char *layoutKvPtr = layoutKvStr.c_str(); + at::Tensor metadata = at::empty(1024, at::TensorOptions(torch_npu::utils::get_npu_device_type()).dtype(at::kInt)); + ACLNN_CMD(aclnnSparseAttnSharedkvMetadata, + cuSeqLensQ, sequsedOriKv, sequsedCmpKv, + sequsedQ, sequsedKv, + numHeadsQ, numHeadsKv, headDim, batchSize, + maxSeqLenQ, maxSeqLenKv, oriTopk, cmpTopk, cmpRatio, + oriMaskMode, cmpMaskMode, oriWinLeft, oriWinRight, + layoutQPtr, layoutKvPtr, + hasOriKv, hasCmpKv, + metadata); + return metadata; +} + +std::tuple npu_sparse_attn_sharedkv( + const at::Tensor &query, + const c10::optional &oriKv, + const c10::optional &cmpKv, + const c10::optional &oriSparseIndices, + const c10::optional &cmpSparseIndices, + const c10::optional &oriBlockTable, + const c10::optional &cmpBlockTable, + const c10::optional &cuSeqLensQ, + const c10::optional &cuSeqLensOriKv, + const c10::optional &cuSeqLensCmpKv, + const c10::optional &sequsedQ, + const c10::optional &sequsedKv, + const c10::optional &sinks, + const c10::optional &metadata, + double softmaxScale, + int64_t cmpRatio, + int64_t oriMaskMode, + int64_t cmpMaskMode, + int64_t oriKvStride, + int64_t cmpKvStride, + int64_t oriWinLeft, + int64_t oriWinRight, + const c10::optional layoutQ, + const c10::optional layoutKv, + bool returnSoftmaxLse) +{ + std::string layoutq = layoutQ.value_or("SBH"); + std::string layoutkv = layoutKv.value_or("SBH"); + const char *layoutQPtr = layoutq.c_str(); + const char *layoutKvPtr = layoutkv.c_str(); + + at::Tensor attnOutput = at::empty(query.sizes(), query.options()); + at::Tensor softmaxLseOut; + if (returnSoftmaxLse) { + std::vector lse_sizes(query.sizes().begin(), query.sizes().end()); + lse_sizes.back() = 1; + softmaxLseOut = at::empty(lse_sizes, query.options().dtype(c10::ScalarType::Float)); + } else { + softmaxLseOut = at::Tensor(); + } + + ACLNN_CMD(aclnnSparseAttnSharedkv, + query, oriKv, cmpKv, + oriSparseIndices, cmpSparseIndices, + oriBlockTable, cmpBlockTable, + cuSeqLensQ, cuSeqLensOriKv, cuSeqLensCmpKv, + sequsedQ, sequsedKv, + sinks, metadata, + softmaxScale, cmpRatio, oriMaskMode, cmpMaskMode, + oriKvStride, cmpKvStride, + oriWinLeft, oriWinRight, + layoutQPtr, layoutKvPtr, + returnSoftmaxLse, + attnOutput, softmaxLseOut); + return std::make_tuple(attnOutput, softmaxLseOut); +} + +std::tuple npu_sparse_attn_sharedkv_grad( + const at::Tensor &query, + const at::Tensor &oriKv, + const c10::optional &cmpKv, + const c10::optional &dOut, + const c10::optional &out, + const c10::optional &lse, + const c10::optional &oriSparseIndices, + const c10::optional &cmpSparseIndices, + const c10::optional &cuSeqlensQ, + const c10::optional &cuSeqlensOriKv, + const c10::optional &cuSeqlensCmpKv, + const at::Tensor &sinks, + double scaleValue, + int64_t cmpRatio, + int64_t oriMaskMode, + int64_t cmpMaskMode, + int64_t oriWinLeft, + int64_t oriWinRight, + const c10::optional layout) +{ + std::string layoutValue = layout.value_or("SBH"); + const char *layoutPtr = layoutValue.c_str(); + + at::Tensor dQuery = at::empty(query.sizes(), query.options()); + at::Tensor dOriKv = at::empty(oriKv.sizes(), oriKv.options()); + at::Tensor dSinks = at::empty(sinks.sizes(), sinks.options()); + + at::Tensor dCmpKv; + if (cmpRatio > 1 && cmpKv.has_value() && cmpKv.value().defined()) { + dCmpKv = at::empty(cmpKv.value().sizes(), cmpKv.value().options()); + } else { + dCmpKv = at::Tensor(); + } + + ACLNN_CMD(aclnnSparseAttnSharedkvGrad, + query, oriKv, cmpKv, + dOut, out, lse, + oriSparseIndices, cmpSparseIndices, + cuSeqlensQ, cuSeqlensOriKv, cuSeqlensCmpKv, + sinks, + scaleValue, cmpRatio, oriMaskMode, cmpMaskMode, + oriWinLeft, oriWinRight, + layoutPtr, + dQuery, dOriKv, dCmpKv, dSinks); + return std::make_tuple(dQuery, dOriKv, dCmpKv, dSinks); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("npu_sparse_attn_sharedkv_metadata", &npu_sparse_attn_sharedkv_metadata, + "Shared-KV Sparse Attention Metadata"); + m.def("npu_sparse_attn_sharedkv", &npu_sparse_attn_sharedkv, + "Shared-KV Sparse Attention Forward"); + m.def("npu_sparse_attn_sharedkv_grad", &npu_sparse_attn_sharedkv_grad, + "Shared-KV Sparse Attention Backward"); +} diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/aclnn_ops.py b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn_ops.py new file mode 100644 index 000000000..da5f7c1f1 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/aclnn_ops.py @@ -0,0 +1,129 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4 NPU fusion operators (SAS + LI) backed by self-compiled ACLNN. + +JIT-compiles thin C++ bindings that invoke the ACLNN kernels directly +(``aclnnSparseAttnSharedkv``, ``aclnnLightningIndexer``, etc.), no mindspeed. +""" +from __future__ import annotations + +import torch + +from twinkle import get_logger +from .aclnn.builder import build_op + +logger = get_logger() + +TORCH_MAX_INT = 9223372036854775807 + +_sas_op = _li_op = _li_grad_op = None + + +def _ensure_ops(): + global _sas_op, _li_op, _li_grad_op + if _sas_op is None: + _sas_op = build_op('sparse_attn_sharedkv', ['sparse_attn_sharedkv/binding.cpp']) + _li_op = build_op('lightning_indexer', ['lightning_indexer/binding.cpp']) + _li_grad_op = build_op('lightning_indexer_grad', ['lightning_indexer_grad/binding.cpp']) + logger.info('[NPU] [ACLNN] SAS + LI ops compiled successfully') + + +class SparseAttnSharedKV(torch.autograd.Function): + """SAS: Shared-KV sparse attention (forward + backward via ACLNN). + + Metadata is computed OUTSIDE this Function (in the convenience wrapper) + so it is NOT recomputed by gradient checkpointing's backward recompute. + """ + + @staticmethod + def forward(ctx, query, ori_kv, cmp_kv, cmp_sparse_indices, sinks, metadata, softmax_scale, cmp_ratio, + ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, layout): + _ensure_ops() + ori_stride = ori_kv.stride(0) + cmp_stride = cmp_kv.stride(0) if cmp_kv is not None else 0 + + result, lse = _sas_op.npu_sparse_attn_sharedkv(query, ori_kv, cmp_kv, None, cmp_sparse_indices, None, None, + None, None, None, None, None, sinks, metadata, softmax_scale, + cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_stride, cmp_stride, + ori_win_left, ori_win_right, layout, layout, True) + + ctx.save_for_backward(query, ori_kv, cmp_kv, result, lse, cmp_sparse_indices, sinks) + ctx.scale, ctx.cmp_ratio = softmax_scale, cmp_ratio + ctx.ori_mm, ctx.cmp_mm = ori_mask_mode, cmp_mask_mode + ctx.ori_wl, ctx.ori_wr = ori_win_left, ori_win_right + ctx.layout = layout + return result + + @staticmethod + def backward(ctx, grad_output): + q, ori_kv, cmp_kv, result, lse, cmp_si, sinks = ctx.saved_tensors + q_g, kv_g, cmp_g, sinks_g = _sas_op.npu_sparse_attn_sharedkv_grad(q, ori_kv, cmp_kv, grad_output, result, lse, + None, cmp_si, None, None, None, sinks, + ctx.scale, ctx.cmp_ratio, ctx.ori_mm, + ctx.cmp_mm, ctx.ori_wl, ctx.ori_wr, + ctx.layout) + return (q_g, kv_g, cmp_g, None, sinks_g, None) + (None, ) * 8 + + +def npu_sparse_attn_shared_kv(query, + ori_kv, + cmp_kv, + cmp_sparse_indices, + sinks, + softmax_scale, + cmp_ratio, + ori_mask_mode=4, + cmp_mask_mode=3, + ori_win_left=127, + ori_win_right=0, + layout='BSND'): + """Convenience wrapper: ``[B,S,N,D]`` query, ``[B,S,D]`` shared KV. + + Metadata is pre-computed here (outside the autograd Function) so gradient + checkpointing's backward recomputation does NOT re-invoke the metadata kernel. + """ + _ensure_ops() + b, s_q, n_h, h_d = query.shape + s_kv = ori_kv.size(1) + topk = 0 if cmp_sparse_indices is None else cmp_sparse_indices.size(-1) + has_cmp_kv = cmp_kv is not None + + e = torch.tensor([]).npu() + metadata = _sas_op.npu_sparse_attn_sharedkv_metadata(e, e, e, e, e, n_h, 1, h_d, b, s_q, s_kv, 0, topk, cmp_ratio, + ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, + layout, layout, True, has_cmp_kv) + + query = query.contiguous() + ori_kv = ori_kv.unsqueeze(2).contiguous() + cmp_kv = cmp_kv if cmp_kv is None else cmp_kv.unsqueeze(2).contiguous() + if cmp_sparse_indices is not None: + cmp_sparse_indices = cmp_sparse_indices.unsqueeze(2).contiguous() + return SparseAttnSharedKV.apply(query, ori_kv, cmp_kv, cmp_sparse_indices, sinks, metadata, softmax_scale, + cmp_ratio, ori_mask_mode, cmp_mask_mode, ori_win_left, ori_win_right, + layout).contiguous() + + +class LightningIndexer(torch.autograd.Function): + """LI: Lightning Indexer (forward + backward via ACLNN).""" + + @staticmethod + def forward(ctx, query, key, weights, sparse_count, sparse_mode, cmp_ratio): + _ensure_ops() + indices, values = _li_op.npu_lightning_indexer(query, key, weights, None, None, None, 'BSND', 'BSND', + sparse_count, sparse_mode, TORCH_MAX_INT, TORCH_MAX_INT, + cmp_ratio, True) + ctx.save_for_backward(query, key, weights, indices) + ctx.sparse_mode, ctx.cmp_ratio = sparse_mode, cmp_ratio + return indices, values + + @staticmethod + def backward(ctx, grad_indices, grad_values): + q, key, weights, indices = ctx.saved_tensors + q_g, k_g, w_g = _li_grad_op.npu_lightning_indexer_grad(q, key, grad_values, indices, weights, None, None, + 'BSND', ctx.sparse_mode, TORCH_MAX_INT, TORCH_MAX_INT, + ctx.cmp_ratio, None) + return q_g, k_g, w_g, None, None, None + + +def npu_lightning_indexer(query, key, weights, sparse_count=2048, sparse_mode=3, cmp_ratio=1): + """Convenience wrapper for the Lightning Indexer.""" + return LightningIndexer.apply(query, key, weights, sparse_count, sparse_mode, cmp_ratio) diff --git a/src/twinkle/kernel/ops/dsv4_sas_li/npu.py b/src/twinkle/kernel/ops/dsv4_sas_li/npu.py new file mode 100644 index 000000000..cf2dc5ec9 --- /dev/null +++ b/src/twinkle/kernel/ops/dsv4_sas_li/npu.py @@ -0,0 +1,334 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4 NPU attention (SAS) / Lightning Indexer (LI) forwards. + +Used as class-attribute replacements on the corresponding HF classes via +the ``DEFAULT_KERNEL_CONFIG`` mapping (see ``config`` and ``registry``). +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from twinkle import get_logger + +logger = get_logger() + +_sas_logged = False +_li_logged = False + + +def npu_dsv4_attention_forward( + self, + hidden_states, + position_embeddings, + position_ids, + attention_mask, + past_key_values=None, + **kwargs, +): + """Drop-in ``DeepseekV4Attention.forward`` using NPU sparse attention (SAS). + + Falls back to the standard HF attention interface when the ACLNN + extension is unavailable. Expects ``self.compressor(...)`` to return a + 3-tuple ``(compressed_kv, block_bias, top_k_indices)`` (see + :func:`npu_dsv4_make_compressor_wrapper`). + """ + from transformers.models.deepseek_v4.modeling_deepseek_v4 import (ALL_ATTENTION_FUNCTIONS, apply_rotary_pos_emb, + eager_attention_forward) + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + cos, sin = position_embeddings[self.rope_layer_type] + + q_residual = self.q_a_norm(self.q_a_proj(hidden_states)) + q = self.q_b_proj(q_residual).view(*hidden_shape).transpose(1, 2) + q = self.q_b_norm(q) + q = apply_rotary_pos_emb(q, cos, sin) + + kv = self.kv_norm(self.kv_proj(hidden_states)).view(*hidden_shape).transpose(1, 2) + kv = apply_rotary_pos_emb(kv, cos, sin) + + if past_key_values is not None: + kv = past_key_values.update(kv, kv, self.layer_idx)[0] + + ori_kv = kv + compressed_kv = None + block_bias = None + top_k_indices = None + if self.compressor is not None: + compressor_out = self.compressor(hidden_states, q_residual, position_ids, past_key_values, self.layer_idx) + if len(compressor_out) == 3: + compressed_kv, block_bias, top_k_indices = compressor_out + else: + compressed_kv, block_bias = compressor_out + + use_sas = True + if self.layer_type == 'sliding_attention': + cmp_ratio = 1 + cmp_kv_arg = None + cmp_sparse_indices = None + elif self.layer_type == 'compressed_sparse_attention': + cmp_ratio = self.config.compress_rates['compressed_sparse_attention'] + if compressed_kv is not None and compressed_kv.shape[2] > 0: + cmp_kv_arg = compressed_kv.squeeze(1).contiguous() + if top_k_indices is not None: + cmp_sparse_indices = top_k_indices.clamp(min=0).to(torch.int32) + else: + cmp_sparse_indices = None + else: + use_sas = False + cmp_kv_arg = None + cmp_sparse_indices = None + else: + cmp_ratio = self.config.compress_rates['heavily_compressed_attention'] + if compressed_kv is not None and compressed_kv.shape[2] > 0: + cmp_kv_arg = compressed_kv.squeeze(1).contiguous() + else: + use_sas = False + cmp_kv_arg = None + cmp_sparse_indices = None + + try: + from .aclnn_ops import npu_sparse_attn_shared_kv + attn_output = npu_sparse_attn_shared_kv( + query=q.transpose(1, 2).contiguous(), + ori_kv=ori_kv.squeeze(1).contiguous(), + cmp_kv=cmp_kv_arg, + cmp_sparse_indices=cmp_sparse_indices, + sinks=self.sinks.float(), + softmax_scale=self.scaling, + cmp_ratio=cmp_ratio, + ori_win_left=self.sliding_window - 1, + ) + global _sas_logged + if not _sas_logged: + logger.info( + '[NPU] [DSV4-SAS] Twinkle sparse attention active ' + '(layer_type=%s, cmp_ratio=%s, topk=%s)', + self.layer_type, + cmp_ratio, + 0 if cmp_sparse_indices is None else cmp_sparse_indices.shape[-1], + ) + _sas_logged = True + attn_weights = None + except (ImportError, RuntimeError): + use_sas = False + + if not use_sas: + if compressed_kv is not None: + kv = torch.cat([kv, compressed_kv], dim=2) + if isinstance(attention_mask, torch.Tensor) and kv.shape[2] > attention_mask.shape[-1]: + if block_bias is not None: + attention_mask = torch.cat([attention_mask, block_bias.to(attention_mask.dtype)], dim=-1) + else: + attention_mask = F.pad(attention_mask, (0, kv.shape[2] - attention_mask.shape[-1]), value=0.0) + + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, + eager_attention_forward) + attn_output, attn_weights = attention_interface( + self, + q, + kv, + kv, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + s_aux=self.sinks, + **kwargs, + ) + + attn_output = apply_rotary_pos_emb(attn_output.transpose(1, 2), cos, -sin).transpose(1, 2) + grouped = attn_output.reshape(*input_shape, self.config.o_groups, -1) + grouped = self.o_a_proj(grouped).flatten(2) + output = self.o_b_proj(grouped) + return output, attn_weights + + +def npu_dsv4_indexer_forward( + self, + hidden_states, + q_residual, + position_ids, + past_key_values, + layer_idx, +): + """Drop-in ``DeepseekV4Indexer.forward`` using NPU Lightning Indexer (LI). + + Falls back to a pure-torch top-k implementation when the ACLNN + extension is unavailable. + """ + from transformers.models.deepseek_v4.modeling_deepseek_v4 import apply_rotary_pos_emb + + batch, seq_len, _ = hidden_states.shape + cache_layer = past_key_values.layers[layer_idx] if past_key_values is not None else None + kv = self.kv_proj(hidden_states) + gate = self.gate_proj(hidden_states) + + if cache_layer is None: + usable = (kv.shape[1] // self.compress_rate) * self.compress_rate + chunk_kv, chunk_gate, first_window_position = kv[:, :usable], gate[:, :usable], 0 + else: + chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights('indexer', kv, gate) + + if chunk_kv.shape[1] > 0: + n_windows = chunk_kv.shape[1] // self.compress_rate + ratio = self.compress_rate + chunk_kv = chunk_kv.view(batch, n_windows, ratio, -1) + chunk_gate = chunk_gate.view(batch, n_windows, ratio, -1) + self.position_bias.to(chunk_gate.dtype) + + new_kv = chunk_kv.new_zeros((batch, n_windows, 2 * ratio, self.head_dim)) + new_gate = chunk_gate.new_full((batch, n_windows, 2 * ratio, self.head_dim), float('-inf')) + new_kv[:, :, ratio:] = chunk_kv[..., self.head_dim:] + new_gate[:, :, ratio:] = chunk_gate[..., self.head_dim:] + if n_windows > 1: + new_kv[:, 1:, :ratio] = chunk_kv[:, :-1, :, :self.head_dim] + new_gate[:, 1:, :ratio] = chunk_gate[:, :-1, :, :self.head_dim] + if cache_layer is not None: + prior_kv, prior_gate = cache_layer.update_overlap_state('indexer', chunk_kv, chunk_gate, self.head_dim) + if prior_kv is not None: + new_kv[:, 0, :ratio] = prior_kv.to(new_kv.dtype) + new_gate[:, 0, :ratio] = prior_gate.to(new_gate.dtype) + + compressed = self.kv_norm((new_kv * new_gate.softmax(dim=2, dtype=torch.float32).to(new_kv.dtype)).sum(dim=2)) + positions = torch.arange(n_windows, device=compressed.device) + positions = positions * self.compress_rate + first_window_position + positions = positions.unsqueeze(0).expand(batch, -1) + cos, sin = self.rotary_emb(compressed, position_ids=positions, layer_type=self.rope_layer_type) + compressed = apply_rotary_pos_emb(compressed.unsqueeze(1), cos, sin).squeeze(1) + else: + compressed = chunk_kv.new_zeros((batch, 0, self.head_dim)) + + compressed_kv = (compressed if cache_layer is None else cache_layer.update_compressor_states('indexer', compressed)) + + cos_q, sin_q = self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=self.rope_layer_type) + q = self.q_b_proj(q_residual).view(batch, seq_len, -1, self.head_dim).transpose(1, 2) + q = apply_rotary_pos_emb(q, cos_q, sin_q).transpose(1, 2) + + def torch_indexer_top_k_indices(): + index_scores = self.scorer(q, compressed_kv, hidden_states) + compressed_len = compressed_kv.shape[1] + top_k = min(self.index_topk, compressed_len) + if compressed_len > 0: + causal_threshold = (position_ids + 1) // self.compress_rate + entry_indices = torch.arange(compressed_len, device=index_scores.device) + future_mask = entry_indices.view(1, 1, -1) >= causal_threshold.unsqueeze(-1) + index_scores = index_scores.masked_fill(future_mask, float('-inf')) + top_k_indices = index_scores.topk(top_k, dim=-1).indices + invalid = top_k_indices >= causal_threshold.unsqueeze(-1) + top_k_indices = torch.where(invalid, torch.full_like(top_k_indices, -1), top_k_indices) + if top_k < self.index_topk: + padding = top_k_indices.new_full((batch, seq_len, self.index_topk - top_k), -1) + top_k_indices = torch.cat([top_k_indices, padding], dim=-1) + return top_k_indices + return index_scores.new_full((batch, seq_len, self.index_topk), -1, dtype=torch.long) + + if compressed_kv.shape[1] > 0: + try: + from .aclnn_ops import npu_lightning_indexer + + scorer = self.scorer + weights = scorer.weights_proj(hidden_states).to(torch.bfloat16) * scorer.weights_scaling + q_indexer = q.to(torch.bfloat16) + k_indexer = compressed_kv.to(torch.bfloat16).unsqueeze(2) + top_k_indices, _ = npu_lightning_indexer( + q_indexer, + k_indexer, + weights, + sparse_count=self.index_topk, + sparse_mode=3, + cmp_ratio=self.compress_rate, + ) + top_k_indices = top_k_indices.squeeze(2) + global _li_logged + if not _li_logged: + logger.info( + '[NPU] [DSV4-LI] Twinkle lightning indexer active ' + '(sparse_count=%s, cmp_ratio=%s)', + self.index_topk, + self.compress_rate, + ) + _li_logged = True + return top_k_indices + except (ImportError, RuntimeError, NameError): + pass + + return torch_indexer_top_k_indices() + + +def npu_dsv4_csa_compressor_forward( + self, + hidden_states, + q_residual, + position_ids, + past_key_values, + layer_idx, +): + """Drop-in ``DeepseekV4CSACompressor.forward`` that returns a 3-tuple. + + Identical to the stock CSA compressor forward, but returns + ``(compressed_kv, block_bias, top_k_indices)`` so the SAS attention forward + can use ``top_k_indices`` directly — **without re-invoking the indexer**. + + The stock forward already calls ``self.indexer(...)`` internally to build + ``block_bias``; a wrapper that called the indexer a second time to fetch + ``top_k_indices`` would mutate ``DeepseekV4CSACache`` twice + (``store_compression_weights`` appends on every call). Under + gradient checkpointing the recomputed forward sees a cache already mutated + by the first forward, producing a different compressed length + (e.g. 714 vs 712 tokens) and triggering ``CheckpointError``. + """ + from transformers.models.deepseek_v4.modeling_deepseek_v4 import apply_rotary_pos_emb + + batch, seq_len, _ = hidden_states.shape + cache_layer = past_key_values.layers[layer_idx] if past_key_values is not None else None + kv = self.kv_proj(hidden_states) + gate = self.gate_proj(hidden_states) + + if cache_layer is None: + usable = (kv.shape[1] // self.compress_rate) * self.compress_rate + chunk_kv, chunk_gate, first_window_position = kv[:, :usable], gate[:, :usable], 0 + else: + chunk_kv, chunk_gate, first_window_position = cache_layer.store_compression_weights('compressor', kv, gate) + + if chunk_kv.shape[1] > 0: + n_windows = chunk_kv.shape[1] // self.compress_rate + ratio = self.compress_rate + chunk_kv = chunk_kv.view(batch, n_windows, ratio, -1) + chunk_gate = chunk_gate.view(batch, n_windows, ratio, -1) + self.position_bias + + new_kv = chunk_kv.new_zeros((batch, n_windows, 2 * ratio, self.head_dim)) + new_gate = chunk_gate.new_full((batch, n_windows, 2 * ratio, self.head_dim), float('-inf')) + new_kv[:, :, ratio:] = chunk_kv[..., self.head_dim:] + new_gate[:, :, ratio:] = chunk_gate[..., self.head_dim:] + if n_windows > 1: + new_kv[:, 1:, :ratio] = chunk_kv[:, :-1, :, :self.head_dim] + new_gate[:, 1:, :ratio] = chunk_gate[:, :-1, :, :self.head_dim] + if cache_layer is not None: + prior_kv, prior_gate = cache_layer.update_overlap_state('compressor', chunk_kv, chunk_gate, self.head_dim) + if prior_kv is not None: + new_kv[:, 0, :ratio] = prior_kv.to(new_kv.dtype) + new_gate[:, 0, :ratio] = prior_gate.to(new_gate.dtype) + + compressed = self.kv_norm((new_kv * new_gate.softmax(dim=2, dtype=torch.float32).to(new_kv.dtype)).sum(dim=2)) + positions = torch.arange(n_windows, device=compressed.device) + positions = positions * self.compress_rate + first_window_position + positions = positions.unsqueeze(0).expand(batch, -1) + cos, sin = self.rotary_emb(compressed, position_ids=positions, layer_type=self.rope_layer_type) + compressed = apply_rotary_pos_emb(compressed.unsqueeze(1), cos, sin).squeeze(1) + else: + compressed = chunk_kv.new_zeros((batch, 0, self.head_dim)) + + if cache_layer is not None: + compressed = cache_layer.update_compressor_states('compressor', compressed) + compressed_kv = compressed.unsqueeze(1) + + # Lightning Indexer — called ONCE here; the result is returned alongside + # block_bias so the SAS attention forward need not re-invoke the indexer. + top_k_indices = self.indexer(hidden_states, q_residual, position_ids, past_key_values, layer_idx) + compressed_len = compressed_kv.shape[2] + valid = top_k_indices >= 0 + safe_indices = torch.where(valid, top_k_indices, torch.full_like(top_k_indices, compressed_len)) + block_bias = compressed_kv.new_full((batch, 1, seq_len, compressed_len + 1), float('-inf')) + block_bias.scatter_(-1, safe_indices.unsqueeze(1).to(torch.int64), 0.0) + return compressed_kv, block_bias[..., :compressed_len], top_k_indices diff --git a/tests/kernel/ops/test_aclnn_ops.py b/tests/kernel/ops/test_aclnn_ops.py new file mode 100644 index 000000000..67f6df998 --- /dev/null +++ b/tests/kernel/ops/test_aclnn_ops.py @@ -0,0 +1,55 @@ +"""Tests for the self-compiled ACLNN SAS/LI wrapper module.""" +import inspect + +import torch + + +def test_aclnn_ops_imports(): + from twinkle.kernel.ops.dsv4_sas_li.aclnn_ops import ( + LightningIndexer, SparseAttnSharedKV, npu_lightning_indexer, npu_sparse_attn_shared_kv, + ) + assert callable(npu_sparse_attn_shared_kv) + assert callable(npu_lightning_indexer) + assert issubclass(SparseAttnSharedKV, torch.autograd.Function) + assert issubclass(LightningIndexer, torch.autograd.Function) + + +def test_sas_convenience_signature(): + from twinkle.kernel.ops.dsv4_sas_li.aclnn_ops import npu_sparse_attn_shared_kv + + sig = inspect.signature(npu_sparse_attn_shared_kv) + params = list(sig.parameters) + expected = ['query', 'ori_kv', 'cmp_kv', 'cmp_sparse_indices', 'sinks', + 'softmax_scale', 'cmp_ratio'] + assert params[:7] == expected + assert sig.parameters['ori_mask_mode'].default == 4 + assert sig.parameters['cmp_mask_mode'].default == 3 + assert sig.parameters['ori_win_left'].default == 127 + assert sig.parameters['ori_win_right'].default == 0 + assert sig.parameters['layout'].default == 'BSND' + + +def test_li_convenience_signature(): + from twinkle.kernel.ops.dsv4_sas_li.aclnn_ops import npu_lightning_indexer + + sig = inspect.signature(npu_lightning_indexer) + params = list(sig.parameters) + assert params == ['query', 'key', 'weights', 'sparse_count', 'sparse_mode', 'cmp_ratio'] + assert sig.parameters['sparse_count'].default == 2048 + assert sig.parameters['sparse_mode'].default == 3 + assert sig.parameters['cmp_ratio'].default == 1 + + +def test_dsv4_npu_does_not_import_mindspeed(): + """The SAS/LI forward functions must not reference mindspeed at module level.""" + import twinkle.kernel.ops.dsv4_sas_li.npu as att_mod + + source = inspect.getsource(att_mod) + assert 'import mindspeed' not in source + assert 'mindspeed.ops' not in source + + +def test_builder_import(): + from twinkle.kernel.ops.dsv4_sas_li.aclnn.builder import build_op + + assert callable(build_op)