From a773f3053dc16860676d01c2eb08c3b6e533a81c Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:55 +0800 Subject: [PATCH 01/11] refactor(ascend): centralize ACL status checks and tensor ownership --- src/native/ascend/error_.h | 42 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 src/native/ascend/error_.h diff --git a/src/native/ascend/error_.h b/src/native/ascend/error_.h new file mode 100644 index 000000000..993079bc8 --- /dev/null +++ b/src/native/ascend/error_.h @@ -0,0 +1,42 @@ +#ifndef INFINI_OPS_ASCEND_ERROR__H_ +#define INFINI_OPS_ASCEND_ERROR__H_ + +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/acl_meta.h" + +namespace infini::ops::ascend { + +inline bool IsAscend310P() { + const auto* soc = aclrtGetSocName(); + return soc != nullptr && std::strncmp(soc, "Ascend310P", 10) == 0; +} + +inline void CheckAcl(int64_t status, const char* operation) { +#ifndef NDEBUG + if (status != ACL_SUCCESS) { + std::fprintf(stderr, "%s:%d: %s: `%s` failed with ACL status %lld.\n", + __FILE__, __LINE__, __func__, operation, + static_cast(status)); + } + assert(status == ACL_SUCCESS && "`ACL` operation failed."); +#endif + (void)status; + (void)operation; +} + +using AclTensorOwner = std::unique_ptr; + +inline AclTensorOwner OwnAclTensor(aclTensor* tensor) { + assert(tensor != nullptr && "`aclCreateTensor` failed."); + return AclTensorOwner(tensor, &aclDestroyTensor); +} + +} // namespace infini::ops::ascend + +#endif From 56426c3e23478c8b6b9a450dc6ba0d2760b43111 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:55 +0800 Subject: [PATCH 02/11] fix(ascend): correct fused RMSNorm reuse and short rows --- .../ascend/ops/fused_add_rms_norm/kernel.h | 126 +++++++++++++++--- tests/test_fused_add_rms_norm.py | 23 ++++ 2 files changed, 130 insertions(+), 19 deletions(-) diff --git a/src/native/ascend/ops/fused_add_rms_norm/kernel.h b/src/native/ascend/ops/fused_add_rms_norm/kernel.h index 6dc4d2526..01ef9d8ce 100644 --- a/src/native/ascend/ops/fused_add_rms_norm/kernel.h +++ b/src/native/ascend/ops/fused_add_rms_norm/kernel.h @@ -14,6 +14,7 @@ #include "base/rms_norm.h" #include "data_type.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -25,8 +26,7 @@ inline void* AllocateOnes(Tensor::Size dim, DataType dtype) { auto bytes = dim * element_size; void* data = nullptr; auto ret = aclrtMalloc(&data, bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "`FusedAddRmsNorm` Ascend path failed to allocate unit weight"); + ascend::CheckAcl(ret, "FusedAddRmsNorm unit weight allocation"); if (dtype == DataType::kFloat32) { std::vector host(dim, 1.0f); @@ -42,8 +42,7 @@ inline void* AllocateOnes(Tensor::Size dim, DataType dtype) { aclrtMemcpy(data, bytes, host.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE); } - assert(ret == ACL_SUCCESS && - "`FusedAddRmsNorm` Ascend path failed to initialize unit weight"); + ascend::CheckAcl(ret, "FusedAddRmsNorm unit weight initialization"); return data; } @@ -65,6 +64,8 @@ class Operator tensor_bytes_(num_tokens_ * row_bytes_), needs_input_staging_(!input.IsContiguous()), needs_residual_staging_(!residual.IsContiguous()) { + use_small_fallback_ = ascend::IsAscend310P() && + dtype_ == DataType::kFloat16 && row_bytes_ < 32; if (!weight.has_value()) { unit_weight_data_ = ascend_fused_add_rms_norm_detail::AllocateOnes(dim_, input.dtype()); @@ -130,6 +131,18 @@ class Operator auto t_residual = residual_cache_.get(residual_data); auto t_weight = weight_cache_.get(const_cast(weight_data)); + if (use_small_fallback_) { + RunSmallFallback(input, input_data, residual_data, t_weight, epsilon, + stream); + if (needs_input_staging_) { + UnpackRows(input.data(), input_data, input_strides_, stream); + } + if (needs_residual_staging_) { + UnpackRows(residual.data(), residual_data, residual_strides_, stream); + } + return; + } + auto& rstd_arena = ascend::GetWorkspacePool().Ensure( stream, rstd_size_, "fused_add_rms_norm_rstd"); if (!rstd_tensor_) { @@ -138,24 +151,44 @@ class Operator ACL_FLOAT, /*strides=*/nullptr, 0, ACL_FORMAT_ND, rstd_shape_.data(), static_cast(rstd_shape_.size()), rstd_arena.buf); } else { - aclSetRawTensorAddr(rstd_tensor_, rstd_arena.buf); + ascend::CheckAcl(aclSetRawTensorAddr(rstd_tensor_, rstd_arena.buf), + "FusedAddRmsNorm rstd address"); } if (!executor_) { - aclnnInplaceAddRmsNormGetWorkspaceSize( - t_input, t_residual, t_weight, static_cast(epsilon), - rstd_tensor_, &ws_size_, &executor_); - aclSetAclOpExecutorRepeatable(executor_); + ascend::CheckAcl( + aclnnInplaceAddRmsNormGetWorkspaceSize( + t_input, t_residual, t_weight, static_cast(epsilon), + rstd_tensor_, &ws_size_, &executor_), + "FusedAddRmsNorm workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(executor_), + "FusedAddRmsNorm repeatable executor"); } else { - aclSetInputTensorAddr(executor_, 0, t_input, input_data); - aclSetInputTensorAddr(executor_, 1, t_residual, residual_data); - aclSetInputTensorAddr(executor_, 2, t_weight, - const_cast(weight_data)); - aclSetOutputTensorAddr(executor_, 0, rstd_tensor_, rstd_arena.buf); + ascend::CheckAcl(aclSetInputTensorAddr(executor_, 0, t_input, input_data), + "FusedAddRmsNorm input address"); + ascend::CheckAcl( + aclSetInputTensorAddr(executor_, 1, t_residual, residual_data), + "FusedAddRmsNorm residual address"); + ascend::CheckAcl(aclSetInputTensorAddr(executor_, 2, t_weight, + const_cast(weight_data)), + "FusedAddRmsNorm weight address"); + // InplaceAddRmsNorm has outputs [x1Ref, rstdOut, x2Ref]. Rebind + // the in-place outputs too when a cached executor sees new buffers. + ascend::CheckAcl( + aclSetOutputTensorAddr(executor_, 0, t_input, input_data), + "FusedAddRmsNorm input output address"); + ascend::CheckAcl( + aclSetOutputTensorAddr(executor_, 1, rstd_tensor_, rstd_arena.buf), + "FusedAddRmsNorm rstd output address"); + ascend::CheckAcl( + aclSetOutputTensorAddr(executor_, 2, t_residual, residual_data), + "FusedAddRmsNorm residual output address"); } auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_); - aclnnInplaceAddRmsNorm(arena.buf, ws_size_, executor_, stream); + ascend::CheckAcl( + aclnnInplaceAddRmsNorm(arena.buf, ws_size_, executor_, stream), + "FusedAddRmsNorm execution"); if (needs_input_staging_) { UnpackRows(input.data(), input_data, input_strides_, stream); @@ -166,6 +199,61 @@ class Operator } private: + void RunSmallFallback(const Tensor input, void* input_data, + void* residual_data, const aclTensor* t_weight, + float epsilon, aclrtStream stream) const { + // A single short row works on 310P, whereas packed short rows do not. + // Separate 32-byte arenas avoid adjacent-row reads/writes in that path. + auto& input_arena = ascend::GetWorkspacePool().Ensure( + stream, 32, "fused_add_rms_norm_short_input"); + auto& residual_arena = ascend::GetWorkspacePool().Ensure( + stream, 32, "fused_add_rms_norm_short_residual"); + auto& rstd_arena = ascend::GetWorkspacePool().Ensure( + stream, 32, "fused_add_rms_norm_short_rstd"); + Tensor row_input{input_arena.buf, Tensor::Shape{1, dim_}, dtype_, + input.device()}; + Tensor row_residual{residual_arena.buf, Tensor::Shape{1, dim_}, dtype_, + input.device()}; + Tensor rstd{rstd_arena.buf, Tensor::Shape{1, 1}, DataType::kFloat32, + input.device()}; + auto t_input = ascend::OwnAclTensor(ascend::BuildAclTensor(row_input)); + auto t_residual = + ascend::OwnAclTensor(ascend::BuildAclTensor(row_residual)); + auto t_rstd = ascend::OwnAclTensor(ascend::BuildAclTensor(rstd)); + for (Tensor::Size row = 0; row < num_tokens_; ++row) { + auto* input_row = static_cast(input_data) + row * row_bytes_; + auto* residual_row = + static_cast(residual_data) + row * row_bytes_; + ascend::CheckAcl( + aclrtMemcpyAsync(input_arena.buf, 32, input_row, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "FusedAddRmsNorm short input copy"); + ascend::CheckAcl( + aclrtMemcpyAsync(residual_arena.buf, 32, residual_row, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "FusedAddRmsNorm short residual copy"); + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + ascend::CheckAcl(aclnnInplaceAddRmsNormGetWorkspaceSize( + t_input.get(), t_residual.get(), t_weight, + static_cast(epsilon), t_rstd.get(), + &workspace_size, &executor), + "FusedAddRmsNorm short-row workspace query"); + auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size); + ascend::CheckAcl( + aclnnInplaceAddRmsNorm(arena.buf, workspace_size, executor, stream), + "FusedAddRmsNorm short-row execution"); + ascend::CheckAcl( + aclrtMemcpyAsync(input_row, row_bytes_, input_arena.buf, row_bytes_, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "FusedAddRmsNorm short input writeback"); + ascend::CheckAcl( + aclrtMemcpyAsync(residual_row, row_bytes_, residual_arena.buf, + row_bytes_, ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "FusedAddRmsNorm short residual writeback"); + } + } + int64_t RowOffset(int64_t row, const Tensor::Strides& strides) const { int64_t remaining = row; int64_t offset = 0; @@ -186,8 +274,7 @@ class Operator static_cast(dst) + row * row_bytes_, row_bytes_, static_cast(src) + offset * element_size_, row_bytes_, ACL_MEMCPY_DEVICE_TO_DEVICE, stream); - assert(ret == ACL_SUCCESS && - "`FusedAddRmsNorm` Ascend input pack failed"); + ascend::CheckAcl(ret, "FusedAddRmsNorm input pack"); } } @@ -199,8 +286,7 @@ class Operator static_cast(dst) + offset * element_size_, row_bytes_, static_cast(src) + row * row_bytes_, row_bytes_, ACL_MEMCPY_DEVICE_TO_DEVICE, stream); - assert(ret == ACL_SUCCESS && - "`FusedAddRmsNorm` Ascend output unpack failed"); + ascend::CheckAcl(ret, "FusedAddRmsNorm output unpack"); } } @@ -218,6 +304,8 @@ class Operator bool needs_residual_staging_{false}; + bool use_small_fallback_{false}; + void* unit_weight_data_{nullptr}; std::optional unit_weight_; diff --git a/tests/test_fused_add_rms_norm.py b/tests/test_fused_add_rms_norm.py index 9cbd9e765..033900ee8 100644 --- a/tests/test_fused_add_rms_norm.py +++ b/tests/test_fused_add_rms_norm.py @@ -70,6 +70,29 @@ def test_fused_add_rms_norm( torch.testing.assert_close(residual, expected_residual, rtol=rtol, atol=atol) +@pytest.mark.parametrize("width", (4, 8, 16, 128)) +@pytest.mark.parametrize("has_weight", (False, True)) +def test_fused_add_rms_norm_repeated_shape( + device, implementation_index, width, has_weight, +): + if device != "npu" or implementation_index != 0: + pytest.skip("Ascend native repeated-call regression") + + for value in (0.25, 0.75): + input = torch.full((2, 2, width), value, dtype=torch.float16, device=device) + residual = torch.full_like(input, 0.5) + weight = torch.full((width,), 0.75, dtype=torch.float16, device=device) if has_weight else None + expected_input = input.clone() + expected_residual = residual.clone() + _torch_fused_add_rms_norm(expected_input, expected_residual, weight, 1e-6) + infini.ops.fused_add_rms_norm( + input, residual, weight, 1e-6, implementation_index=implementation_index, + stream=get_stream(input.device), + ) + torch.testing.assert_close(input.cpu(), expected_input.cpu(), rtol=1e-2, atol=1e-2) + torch.testing.assert_close(residual.cpu(), expected_residual.cpu(), rtol=0, atol=0) + + def _torch_fused_add_rms_norm(input, residual, weight, epsilon): summed = (input.to(torch.float32) + residual.to(torch.float32)).to(input.dtype) variance = summed.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) From da7d350e7f72950c597eae8f816adcfaf4c7c87b Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:55 +0800 Subject: [PATCH 03/11] fix(ascend): support 310P strided KV cache updates --- .../op_kernel/reshape_and_cache_flash.cpp | 19 +++++ .../ops/reshape_and_cache_flash/kernel.h | 84 ++++++++++++++----- tests/test_reshape_and_cache_flash.py | 28 +++++++ 3 files changed, 111 insertions(+), 20 deletions(-) diff --git a/src/native/ascend/custom/reshape_and_cache_flash/op_kernel/reshape_and_cache_flash.cpp b/src/native/ascend/custom/reshape_and_cache_flash/op_kernel/reshape_and_cache_flash.cpp index aeb4922b0..7d1069856 100644 --- a/src/native/ascend/custom/reshape_and_cache_flash/op_kernel/reshape_and_cache_flash.cpp +++ b/src/native/ascend/custom/reshape_and_cache_flash/op_kernel/reshape_and_cache_flash.cpp @@ -43,6 +43,11 @@ class KernelReshapeAndCacheFlashStrided { } __aicore__ inline void Process() { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 + AscendC::DataCacheCleanAndInvalid( + slot_mapping_); +#endif auto local = copy_buffer_.Get(); AscendC::DataCopyExtParams copy_params{ 1, static_cast(head_bytes_), 0, 0, 0}; @@ -79,6 +84,11 @@ class KernelReshapeAndCacheFlashStrided { copy_params, pad_params); } } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 + AscendC::DataCacheCleanAndInvalid( + key_cache_); +#endif } private: @@ -88,10 +98,19 @@ class KernelReshapeAndCacheFlashStrided { AscendC::LocalTensor& local, const AscendC::DataCopyExtParams& copy_params, const AscendC::DataCopyPadExtParams& pad_params) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200 + // DataCopyPad is unsupported on 310P's m200 cores. Scalar GM accesses + // preserve exact bytes even when padded head addresses are unaligned. + for (int64_t byte = 0; byte < head_bytes_; ++byte) { + destination.SetValue(destination_offset + byte, + source.GetValue(source_offset + byte)); + } +#else AscendC::DataCopyPad(local, source[source_offset], copy_params, pad_params); AscendC::PipeBarrier(); AscendC::DataCopyPad(destination[destination_offset], local, copy_params); AscendC::PipeBarrier(); +#endif } AscendC::TPipe pipe_; diff --git a/src/native/ascend/ops/reshape_and_cache_flash/kernel.h b/src/native/ascend/ops/reshape_and_cache_flash/kernel.h index 7d3ace28c..eb05890c9 100644 --- a/src/native/ascend/ops/reshape_and_cache_flash/kernel.h +++ b/src/native/ascend/ops/reshape_and_cache_flash/kernel.h @@ -14,6 +14,7 @@ #include "base/reshape_and_cache_flash.h" #include "data_type.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -33,7 +34,8 @@ class Operator Tensor value_cache) : ReshapeAndCacheFlash(key, value, slot_mapping, k_scale, v_scale, kv_cache_dtype, key_cache, value_cache), - use_fast_path_(dtype_ != DataType::kFloat32 && key.IsContiguous() && + use_fast_path_(!ascend::IsAscend310P() && + dtype_ != DataType::kFloat32 && key.IsContiguous() && value.IsContiguous() && key_cache.IsContiguous() && value_cache.IsContiguous()), num_blocks_(key_cache.size(0)), @@ -48,12 +50,12 @@ class Operator value_cache_cache_(value_cache), slot_mapping_host_(num_tokens_), indices_host_(num_tokens_ * 2) { - if (!use_fast_path_ && num_tokens_ > 0) { + use_copy_fallback_ = ascend::IsAscend310P(); + if (!use_fast_path_ && !use_copy_fallback_ && num_tokens_ > 0) { auto ret = aclrtMalloc(&indices_device_, indices_host_.size() * sizeof(indices_host_[0]), ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` failed to allocate indices"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash index allocation"); } } @@ -71,6 +73,12 @@ class Operator if (num_tokens_ == 0) return; auto stream = static_cast(stream_); +#ifndef INFINI_HAS_CUSTOM_KERNELS + if (use_copy_fallback_) { + RunFallback(key, value, slot_mapping, key_cache, value_cache, stream); + return; + } +#endif if (use_fast_path_) { RunFastPath(key, value, slot_mapping, key_cache, value_cache, stream); } else { @@ -107,13 +115,11 @@ class Operator t_key, t_key_cache, t_slot_mapping, t_value, t_value_cache, nullptr, nullptr, nullptr, cache_mode_.data(), nullptr, nullptr, nullptr, &workspace_size, &executor); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` workspace query failed"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash fast workspace query"); auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size); ret = aclnnScatterPaKvCache(arena.buf, workspace_size, executor, stream); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` execution failed"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash fast execution"); } #ifdef INFINI_HAS_CUSTOM_KERNELS @@ -137,8 +143,7 @@ class Operator static_cast(value_cache_page_stride_), static_cast(value_cache_head_stride_), static_cast(element_size_)); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` strided execution failed"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash strided execution"); } #endif @@ -149,11 +154,52 @@ class Operator auto ret = aclrtMemcpyAsync(slot_mapping_host_.data(), mapping_bytes, slot_mapping.data(), mapping_bytes, ACL_MEMCPY_DEVICE_TO_HOST, stream); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` failed to copy slot mapping"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash slot mapping copy"); ret = aclrtSynchronizeStream(stream); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` failed to synchronize slot mapping"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash slot mapping synchronization"); + + if (use_copy_fallback_) { + // Runtime copies support arbitrary NHD/HND strides without relying on + // ScatterPaKvCache or a strided ScatterNdUpdate kernel on 310P. + const auto bytes = head_size_ * element_size_; + for (std::size_t token = 0; token < num_tokens_; ++token) { + const auto slot = slot_mapping_host_[token]; + if (slot < 0) continue; + assert(static_cast(slot) < num_blocks_ * block_size_ && + "`ReshapeAndCacheFlash` slot is outside the cache."); + const auto block = slot / block_size_; + const auto page = slot % block_size_; + for (std::size_t head = 0; head < num_heads_; ++head) { + const auto key_offset = + token * key_token_stride_ + head * key_head_stride_; + const auto value_offset = + token * value_token_stride_ + head * value_head_stride_; + const auto key_cache_offset = block * key_cache_block_stride_ + + page * key_cache_page_stride_ + + head * key_cache_head_stride_; + const auto value_cache_offset = block * value_cache_block_stride_ + + page * value_cache_page_stride_ + + head * value_cache_head_stride_; + ascend::CheckAcl( + aclrtMemcpyAsync(static_cast(key_cache.data()) + + key_cache_offset * element_size_, + bytes, + static_cast(key.data()) + + key_offset * element_size_, + bytes, ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "ReshapeAndCacheFlash key copy"); + ascend::CheckAcl( + aclrtMemcpyAsync(static_cast(value_cache.data()) + + value_cache_offset * element_size_, + bytes, + static_cast(value.data()) + + value_offset * element_size_, + bytes, ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "ReshapeAndCacheFlash value copy"); + } + } + return; + } for (std::size_t token = 0; token < num_tokens_; ++token) { auto slot = slot_mapping_host_[token]; @@ -168,8 +214,7 @@ class Operator auto indices_bytes = indices_host_.size() * sizeof(indices_host_[0]); ret = aclrtMemcpyAsync(indices_device_, indices_bytes, indices_host_.data(), indices_bytes, ACL_MEMCPY_HOST_TO_DEVICE, stream); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` failed to copy cache indices"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash index copy"); auto* key_data = static_cast(key.data()); auto* value_data = static_cast(value.data()); @@ -200,16 +245,15 @@ class Operator uint64_t workspace_size = 0; auto ret = aclnnScatterNdUpdateGetWorkspaceSize(cache, indices, update, &workspace_size, &executor); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` fallback query failed"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash fallback workspace query"); auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size); ret = aclnnScatterNdUpdate(arena.buf, workspace_size, executor, stream); - assert(ret == ACL_SUCCESS && - "Ascend `ReshapeAndCacheFlash` fallback execution failed"); + ascend::CheckAcl(ret, "ReshapeAndCacheFlash fallback execution"); } bool use_fast_path_{false}; + bool use_copy_fallback_{false}; std::size_t num_blocks_{0}; std::size_t element_size_{0}; diff --git a/tests/test_reshape_and_cache_flash.py b/tests/test_reshape_and_cache_flash.py index 3191a5d0d..1001bf7a4 100644 --- a/tests/test_reshape_and_cache_flash.py +++ b/tests/test_reshape_and_cache_flash.py @@ -154,6 +154,34 @@ def invoke(): torch.testing.assert_close(value_cache, expected_value_cache) +@pytest.mark.parametrize("cache_layout", ("NHD", "HND")) +def test_reshape_and_cache_flash_repeated_slots(device, implementation_index, cache_layout): + if device != "npu" or implementation_index != 0: + pytest.skip("Ascend native cache update regression") + + key_cache = _make_cache(3, 4, 2, 16, cache_layout, torch.float16, device).fill_(-8) + value_cache = _make_cache(3, 4, 2, 16, cache_layout, torch.float16, device).fill_(-9) + expected_key = key_cache.cpu().clone() + expected_value = value_cache.cpu().clone() + scale = torch.ones((), dtype=torch.float32, device=device) + for iteration, slots in enumerate(((0, -1, 5), (1, 8, -1))): + key_cpu = (torch.arange(96).reshape(3, 2, 16).float() / 32 + iteration).half() + value_cpu = -key_cpu + key, value = key_cpu.to(device), value_cpu.to(device) + mapping = torch.tensor(slots, dtype=torch.int64, device=device) + for token, slot in enumerate(slots): + if slot >= 0: + block, page = divmod(slot, 4) + expected_key[block, page] = key_cpu[token] + expected_value[block, page] = value_cpu[token] + infini.ops.reshape_and_cache_flash( + key, value, mapping, scale, scale, "auto", key_cache, value_cache, + implementation_index=implementation_index, stream=get_stream(key.device), + ) + torch.testing.assert_close(key_cache.cpu(), expected_key, rtol=0, atol=0) + torch.testing.assert_close(value_cache.cpu(), expected_value, rtol=0, atol=0) + + def _make_cache(num_blocks, block_size, num_heads, head_size, layout, dtype, device): if layout == "NHD": shape = (num_blocks, block_size, num_heads, head_size) From 4582f040bdf22ad6c1672d4bec08dfe7a2550a10 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:56 +0800 Subject: [PATCH 04/11] fix(ascend): fall back for unsupported 310P SwiGLU shapes --- src/native/ascend/ops/silu_and_mul/kernel.h | 93 +++++++++++++++++---- tests/test_silu_and_mul.py | 19 +++++ 2 files changed, 97 insertions(+), 15 deletions(-) diff --git a/src/native/ascend/ops/silu_and_mul/kernel.h b/src/native/ascend/ops/silu_and_mul/kernel.h index 84cb39471..f8715cb4b 100644 --- a/src/native/ascend/ops/silu_and_mul/kernel.h +++ b/src/native/ascend/ops/silu_and_mul/kernel.h @@ -7,9 +7,12 @@ #include "acl/acl.h" #include "aclnn/aclnn_base.h" #include "aclnn_copy.h" +#include "aclnnop/aclnn_mul.h" +#include "aclnnop/aclnn_silu.h" #include "aclnnop/aclnn_swi_glu.h" #include "base/silu_and_mul.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -60,34 +63,92 @@ class Operator : public SiluAndMul { } if (!swiglu_exec_) { - aclnnSwiGluGetWorkspaceSize(t_input, -1, t_swiglu_out, &swiglu_ws_, - &swiglu_exec_); - aclSetAclOpExecutorRepeatable(swiglu_exec_); - } else { - aclSetInputTensorAddr(swiglu_exec_, 0, t_input, - const_cast(input.data())); - aclSetOutputTensorAddr(swiglu_exec_, 0, t_swiglu_out, swiglu_out_data); + if (!use_fallback_) { + const auto status = aclnnSwiGluGetWorkspaceSize( + t_input, -1, t_swiglu_out, &swiglu_ws_, &swiglu_exec_); + // CANN rejects some small SwiGlu shapes on 310P at tiling time. + if (status == 561002 && ascend::IsAscend310P()) { + use_fallback_ = true; + swiglu_exec_ = nullptr; + } else { + ascend::CheckAcl(status, "SiluAndMul SwiGlu workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(swiglu_exec_), + "SiluAndMul repeatable executor"); + } + } + } else if (!use_fallback_) { + ascend::CheckAcl(aclSetInputTensorAddr(swiglu_exec_, 0, t_input, + const_cast(input.data())), + "SiluAndMul input address"); + ascend::CheckAcl(aclSetOutputTensorAddr(swiglu_exec_, 0, t_swiglu_out, + swiglu_out_data), + "SiluAndMul output address"); } - auto& arena = ascend::GetWorkspacePool().Ensure(stream, swiglu_ws_); - aclnnSwiGlu(arena.buf, swiglu_ws_, swiglu_exec_, stream); + if (use_fallback_) { + RunFallback(input, t_swiglu_out, stream); + } else { + auto& arena = ascend::GetWorkspacePool().Ensure(stream, swiglu_ws_); + ascend::CheckAcl(aclnnSwiGlu(arena.buf, swiglu_ws_, swiglu_exec_, stream), + "SiluAndMul SwiGlu execution"); + } if (needs_copy_) { if (!copy_exec_) { - aclnnInplaceCopyGetWorkspaceSize(t_out, t_swiglu_out, ©_ws_, - ©_exec_); - aclSetAclOpExecutorRepeatable(copy_exec_); + ascend::CheckAcl(aclnnInplaceCopyGetWorkspaceSize( + t_out, t_swiglu_out, ©_ws_, ©_exec_), + "SiluAndMul output copy workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(copy_exec_), + "SiluAndMul repeatable copy executor"); } else { - aclSetInputTensorAddr(copy_exec_, 0, t_out, out.data()); - aclSetInputTensorAddr(copy_exec_, 1, t_swiglu_out, swiglu_out_data); + ascend::CheckAcl( + aclSetInputTensorAddr(copy_exec_, 0, t_out, out.data()), + "SiluAndMul copy destination address"); + ascend::CheckAcl( + aclSetInputTensorAddr(copy_exec_, 1, t_swiglu_out, swiglu_out_data), + "SiluAndMul copy source address"); } auto& copy_arena = ascend::GetWorkspacePool().Ensure(stream, copy_ws_); - aclnnInplaceCopy(copy_arena.buf, copy_ws_, copy_exec_, stream); + ascend::CheckAcl( + aclnnInplaceCopy(copy_arena.buf, copy_ws_, copy_exec_, stream), + "SiluAndMul output copy execution"); } } private: + void RunFallback(const Tensor input, aclTensor* out, + aclrtStream stream) const { + const auto element_size = kDataTypeToSize.at(input_type_); + Tensor gate{const_cast(input.data()), out_shape_, input_type_, + input.device(), input_strides_}; + Tensor up{static_cast(const_cast(input.data())) + + hidden_size_ * element_size, + out_shape_, input_type_, input.device(), input_strides_}; + auto& staging = ascend::GetWorkspacePool().Ensure( + stream, output_size_ * element_size, "silu_and_mul_activation"); + Tensor activated{staging.buf, out_shape_, out_type_, input.device()}; + auto t_gate = ascend::OwnAclTensor(ascend::BuildAclTensor(gate)); + auto t_up = ascend::OwnAclTensor(ascend::BuildAclTensor(up)); + auto t_activated = ascend::OwnAclTensor(ascend::BuildAclTensor(activated)); + uint64_t workspace_size = 0; + aclOpExecutor* executor = nullptr; + ascend::CheckAcl(aclnnSiluGetWorkspaceSize(t_gate.get(), t_activated.get(), + &workspace_size, &executor), + "SiluAndMul fallback SiLU workspace query"); + auto& silu_arena = + ascend::GetWorkspacePool().Ensure(stream, workspace_size); + ascend::CheckAcl( + aclnnSilu(silu_arena.buf, workspace_size, executor, stream), + "SiluAndMul fallback SiLU execution"); + ascend::CheckAcl(aclnnMulGetWorkspaceSize(t_activated.get(), t_up.get(), + out, &workspace_size, &executor), + "SiluAndMul fallback Mul workspace query"); + auto& mul_arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size); + ascend::CheckAcl(aclnnMul(mul_arena.buf, workspace_size, executor, stream), + "SiluAndMul fallback Mul execution"); + } + mutable ascend::AclTensorCache input_cache_; mutable ascend::AclTensorCache out_cache_; @@ -105,6 +166,8 @@ class Operator : public SiluAndMul { mutable aclOpExecutor* copy_exec_{nullptr}; mutable uint64_t copy_ws_{0}; + + mutable bool use_fallback_{false}; }; } // namespace infini::ops diff --git a/tests/test_silu_and_mul.py b/tests/test_silu_and_mul.py index c5d23e784..cd1eb95f0 100644 --- a/tests/test_silu_and_mul.py +++ b/tests/test_silu_and_mul.py @@ -94,6 +94,25 @@ def test_silu_and_mul_non_default_stream(device, implementation_index): accelerator.synchronize() +@pytest.mark.parametrize("padded", (False, True)) +def test_silu_and_mul_repeated_small_shape(device, implementation_index, padded): + if device != "npu" or implementation_index != 0: + pytest.skip("Ascend native small-shape regression") + + for value in (0.25, 0.75): + input = torch.empty_strided( + (13, 8), (12, 1) if padded else (8, 1), + dtype=torch.float16, device=device, + ).fill_(value) + out = torch.empty_strided( + (13, 4), (6, 1) if padded else (4, 1), + dtype=torch.float16, device=device, + ).fill_(-9) + expected = _torch_silu_and_mul(input, torch.empty_like(out)).cpu() + actual = _silu_and_mul(input, out, implementation_index).cpu() + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + def _silu_and_mul(input, out, implementation_index=0): infini.ops.silu_and_mul( input, From 34b151dc1fadd53a88e1571dc47549d23b63add5 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:56 +0800 Subject: [PATCH 05/11] fix(ascend): chunk single-token NeoX RoPE on 310P --- .../ascend/ops/rotary_embedding/kernel.h | 118 +++++++++++------- tests/test_rotary_embedding.py | 66 ++++++++++ 2 files changed, 138 insertions(+), 46 deletions(-) diff --git a/src/native/ascend/ops/rotary_embedding/kernel.h b/src/native/ascend/ops/rotary_embedding/kernel.h index b8e922e5b..d82a111b4 100644 --- a/src/native/ascend/ops/rotary_embedding/kernel.h +++ b/src/native/ascend/ops/rotary_embedding/kernel.h @@ -15,6 +15,7 @@ #include "aclnnop/aclnn_rotary_position_embedding.h" #include "base/rotary_embedding.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -85,36 +86,21 @@ class Operator ascend::AclTensorCache({num_tokens, 1, head_dim}, acl_dtype, cos_data_); sin_rotary_cache_ = ascend::AclTensorCache({num_tokens, 1, head_dim}, acl_dtype, sin_data_); - query_cache_ = ascend::AclTensorCache( - Tensor{query.data(), - {num_tokens, static_cast(num_heads_), head_dim}, - query.dtype(), - query.device(), - {query_token_stride_, query_head_stride_, 1}}); - query_bytes_ = static_cast(query.numel()) * element_size_; ret = aclrtMalloc(&query_out_data_, query_bytes_, ACL_MEM_MALLOC_NORMAL_ONLY); assert(ret == ACL_SUCCESS && "Ascend `RotaryEmbedding` failed to allocate query output"); - query_out_cache_ = ascend::AclTensorCache( - {num_tokens, static_cast(num_heads_), head_dim}, acl_dtype, - query_out_data_); + MakeChunks(query, num_heads_, query_token_stride_, query_head_stride_, + query_out_data_, query_chunks_); if (key.has_value()) { - key_cache_ = ascend::AclTensorCache( - Tensor{key->data(), - {num_tokens, static_cast(num_kv_heads_), head_dim}, - key->dtype(), - key->device(), - {key_token_stride_, key_head_stride_, 1}}); key_bytes_ = static_cast(key->numel()) * element_size_; ret = aclrtMalloc(&key_out_data_, key_bytes_, ACL_MEM_MALLOC_NORMAL_ONLY); assert(ret == ACL_SUCCESS && "Ascend `RotaryEmbedding` failed to allocate key output"); - key_out_cache_ = ascend::AclTensorCache( - {num_tokens, static_cast(num_kv_heads_), head_dim}, - acl_dtype, key_out_data_); + MakeChunks(*key, num_kv_heads_, key_token_stride_, key_head_stride_, + key_out_data_, key_chunks_); } } @@ -128,10 +114,12 @@ class Operator sin_out_cache_.release(); cos_rotary_cache_.release(); sin_rotary_cache_.release(); - query_cache_.release(); - key_cache_.release(); - query_out_cache_.release(); - key_out_cache_.release(); + for (auto* chunks : {&query_chunks_, &key_chunks_}) { + for (auto& chunk : *chunks) { + chunk.input.release(); + chunk.output.release(); + } + } if (cos_table_data_) aclrtFree(cos_table_data_); if (sin_table_data_) aclrtFree(sin_table_data_); if (cos_data_) aclrtFree(cos_data_); @@ -190,19 +178,65 @@ class Operator auto t_cos = cos_rotary_cache_.get(cos_data_); auto t_sin = sin_rotary_cache_.get(sin_data_); - auto t_query = query_cache_.get(query.data()); - RunOne(t_query, t_cos, t_sin, query.data(), query_out_cache_, - query_out_data_, query_rotary_executor_, query_rotary_ws_size_, - stream); + RunChunks(query.data(), t_cos, t_sin, query_out_data_, query_chunks_, + stream); if (key.has_value()) { - auto t_key = key_cache_.get(key->data()); - RunOne(t_key, t_cos, t_sin, key->data(), key_out_cache_, key_out_data_, - key_rotary_executor_, key_rotary_ws_size_, stream); + RunChunks(key->data(), t_cos, t_sin, key_out_data_, key_chunks_, stream); } } private: + struct RotaryChunk { + ascend::AclTensorCache input; + ascend::AclTensorCache output; + size_t input_offset; + size_t output_offset; + aclOpExecutor* executor{nullptr}; + uint64_t workspace_size{0}; + }; + + void MakeChunks(Tensor tensor, int64_t heads, int64_t token_stride, + int64_t head_stride, void* output_data, + std::vector& chunks) { + // CANN 9.1 on 310P produces incomplete output for FP16 single-token + // NeoX RoPE with D=64/128 above eight heads. Limit each ACLNN launch. + const auto chunk_heads = ascend::IsAscend310P() && + query_type_ == DataType::kFloat16 && + num_tokens_ == 1 && is_neox_ && + (head_size_ == 64 || head_size_ == 128) + ? std::min(heads, 8) + : heads; + for (int64_t first = 0; first < heads; first += chunk_heads) { + const auto count = std::min(chunk_heads, heads - first); + const auto input_offset = first * head_stride * element_size_; + const auto output_offset = first * head_size_ * element_size_; + chunks.push_back( + {ascend::AclTensorCache( + Tensor{static_cast(tensor.data()) + input_offset, + {static_cast(num_tokens_), count, head_size_}, + tensor.dtype(), + tensor.device(), + {token_stride, head_stride, 1}}), + ascend::AclTensorCache( + {static_cast(num_tokens_), count, head_size_}, + ascend::ToAclDtype(tensor.dtype()), + static_cast(output_data) + output_offset), + input_offset, output_offset}); + } + } + + void RunChunks(void* input_data, aclTensor* cos, aclTensor* sin, + void* output_data, std::vector& chunks, + aclrtStream stream) const { + for (auto& chunk : chunks) { + auto input = static_cast(input_data) + chunk.input_offset; + auto output = static_cast(output_data) + chunk.output_offset; + RunOne(chunk.input.get(input), cos, sin, input, chunk.output, output, + chunk.executor, chunk.workspace_size, stream); + } + } + void RunOne(aclTensor* input, aclTensor* cos, aclTensor* sin, void* input_data, ascend::AclTensorCache& output_cache, void* output_data, aclOpExecutor*& executor, @@ -211,9 +245,9 @@ class Operator if (!executor) { auto ret = aclnnRotaryPositionEmbeddingGetWorkspaceSize( input, cos, sin, rotary_mode_, output, &workspace_size, &executor); - assert(ret == ACL_SUCCESS && - "`aclnnRotaryPositionEmbeddingGetWorkspaceSize` failed"); - aclSetAclOpExecutorRepeatable(executor); + ascend::CheckAcl(ret, "aclnnRotaryPositionEmbeddingGetWorkspaceSize"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(executor), + "aclSetAclOpExecutorRepeatable RoPE"); } else { aclSetInputTensorAddr(executor, 0, input, input_data); aclSetInputTensorAddr(executor, 1, cos, cos_data_); @@ -225,7 +259,7 @@ class Operator ascend::GetWorkspacePool().Ensure(stream, workspace_size); auto ret = aclnnRotaryPositionEmbedding(rotary_arena.buf, workspace_size, executor, stream); - assert(ret == ACL_SUCCESS && "`aclnnRotaryPositionEmbedding` failed"); + ascend::CheckAcl(ret, "aclnnRotaryPositionEmbedding"); // Q/K can be narrow views into a fused QKV projection. Preserve their // real token strides when copying the contiguous ACLNN output back. @@ -233,14 +267,12 @@ class Operator uint64_t copy_workspace_size = 0; ret = aclnnInplaceCopyGetWorkspaceSize(input, output, ©_workspace_size, ©_executor); - assert(ret == ACL_SUCCESS && - "Planning Ascend `RotaryEmbedding` output copy failed"); + ascend::CheckAcl(ret, "aclnnInplaceCopyGetWorkspaceSize RoPE"); auto& copy_arena = ascend::GetWorkspacePool().Ensure(stream, copy_workspace_size); ret = aclnnInplaceCopy(copy_arena.buf, copy_workspace_size, copy_executor, stream); - assert(ret == ACL_SUCCESS && - "Copying Ascend `RotaryEmbedding` output failed"); + ascend::CheckAcl(ret, "aclnnInplaceCopy RoPE"); } void UploadCosSinCache(const Tensor cos_sin_cache) const { @@ -321,18 +353,12 @@ class Operator mutable ascend::AclTensorCache sin_out_cache_; mutable ascend::AclTensorCache cos_rotary_cache_; mutable ascend::AclTensorCache sin_rotary_cache_; - mutable ascend::AclTensorCache query_cache_; - mutable ascend::AclTensorCache key_cache_; - mutable ascend::AclTensorCache query_out_cache_; - mutable ascend::AclTensorCache key_out_cache_; + mutable std::vector query_chunks_; + mutable std::vector key_chunks_; mutable aclOpExecutor* cos_index_executor_{nullptr}; mutable uint64_t cos_index_ws_size_{0}; mutable aclOpExecutor* sin_index_executor_{nullptr}; mutable uint64_t sin_index_ws_size_{0}; - mutable aclOpExecutor* query_rotary_executor_{nullptr}; - mutable uint64_t query_rotary_ws_size_{0}; - mutable aclOpExecutor* key_rotary_executor_{nullptr}; - mutable uint64_t key_rotary_ws_size_{0}; }; } // namespace infini::ops diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index c306fa5e7..1cfbf39ef 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -260,6 +260,72 @@ def test_rotary_embedding_ascend_strided_fused_qkv_view( torch.testing.assert_close(key, expected_key, rtol=rtol, atol=atol) +@pytest.mark.parametrize("num_tokens", (1, 2, 24)) +@pytest.mark.parametrize( + "head_size, num_heads", + ((64, 8), (64, 9), (64, 36), (64, 40), + (128, 8), (128, 9), (128, 16), (128, 32)), +) +@pytest.mark.parametrize("has_key", (False, True)) +@pytest.mark.parametrize("strided", (False, True)) +def test_rotary_embedding_ascend_neox_decode( + num_tokens, head_size, num_heads, has_key, strided, implementation_index, device +): + if torch.device(device).type != "npu" or implementation_index != 0: + pytest.skip("The decode regression targets native Ascend RoPE") + + num_kv_heads = num_heads // 2 if num_heads % 2 == 0 else 1 + half_dim = head_size // 2 + frequencies = torch.arange(256).float()[:, None] * ( + 1000000 ** (-torch.arange(half_dim).float() / half_dim) + )[None, :] + cache_cpu = torch.cat((frequencies.cos(), frequencies.sin()), -1).half() + cache = cache_cpu.to(device) + # Reuse the cached operator with new positions and Q/K allocations. + for position in (35, 36): + positions_cpu = torch.arange(position, position + num_tokens) + query_cpu = torch.randn(num_tokens, num_heads * head_size).half() + key_cpu = ( + torch.randn(num_tokens, num_kv_heads * head_size).half() + if has_key + else None + ) + padding = [] + + def to_device(value, heads): + if not strided: + return value.to(device) + parent = torch.full( + (num_tokens, heads, 2 * head_size), -17, + dtype=torch.float16, device=device, + ) + view = parent[..., :head_size] + view.copy_(value.view(num_tokens, heads, head_size)) + padding.append(parent[..., head_size:]) + return view + + query = to_device(query_cpu, num_heads) + key = to_device(key_cpu, num_kv_heads) if has_key else None + _torch_rotary_embedding( + positions_cpu, query_cpu, key_cpu, cache_cpu, + head_size, True, 0, False, + ) + infini.ops.rotary_embedding( + positions_cpu.to(device), query, key, cache, head_size, True, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + torch.testing.assert_close( + query.cpu().reshape(query_cpu.shape), query_cpu, rtol=1e-3, atol=1e-3 + ) + if has_key: + torch.testing.assert_close( + key.cpu().reshape(key_cpu.shape), key_cpu, rtol=1e-3, atol=1e-3 + ) + for untouched in padding: + assert torch.all(untouched.cpu() == -17) + + def _torch_rotary_embedding( positions, query, From 030b78c76a91feacc15c9ba0585bf173d4e34e3a Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:56 +0800 Subject: [PATCH 06/11] perf(ascend): batch strided RMSNorm staging copies --- src/native/ascend/ops/rms_norm/kernel.h | 63 +++++++----------- tests/test_rms_norm.py | 87 +++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 40 deletions(-) diff --git a/src/native/ascend/ops/rms_norm/kernel.h b/src/native/ascend/ops/rms_norm/kernel.h index 734d8ed42..d460613e4 100644 --- a/src/native/ascend/ops/rms_norm/kernel.h +++ b/src/native/ascend/ops/rms_norm/kernel.h @@ -6,6 +6,7 @@ #include "acl/acl.h" #include "aclnn/aclnn_base.h" +#include "aclnn_copy.h" #include "aclnn_rms_norm.h" #include "aclnnop/aclnn_cast.h" #include "base/rms_norm.h" @@ -27,20 +28,20 @@ class Operator : public RmsNorm { needs_input_staging_ = !input.IsContiguous(); needs_out_staging_ = !out.IsContiguous(); logical_rows_ = input.numel() / dim_; - element_size_ = kDataTypeToSize.at(input.dtype()); - row_bytes_ = dim_ * element_size_; - tensor_bytes_ = logical_rows_ * row_bytes_; + tensor_bytes_ = input.numel() * kDataTypeToSize.at(input.dtype()); std::vector input_shape(input.shape().begin(), input.shape().end()); std::vector out_shape(out.shape().begin(), out.shape().end()); if (needs_input_staging_) { + strided_in_cache_ = ascend::AclTensorCache(input); norm_in_cache_ = ascend::AclTensorCache( input_shape, ascend::ToAclDtype(input.dtype()), nullptr); } else { norm_in_cache_ = ascend::AclTensorCache(input); } if (needs_out_staging_) { + strided_out_cache_ = ascend::AclTensorCache(out); norm_out_cache_ = ascend::AclTensorCache( out_shape, ascend::ToAclDtype(out.dtype()), nullptr); } else { @@ -115,7 +116,9 @@ class Operator : public RmsNorm { if (needs_input_staging_) { auto& staging = ascend::GetWorkspacePool().Ensure(stream, tensor_bytes_, "rms_norm_input"); - PackRows(staging.buf, input.data(), input_shape_, input_strides_, stream); + CopyTensor(norm_in_cache_.get(staging.buf), + strided_in_cache_.get(const_cast(input.data())), + stream); input_data = staging.buf; } if (needs_out_staging_) { @@ -156,47 +159,29 @@ class Operator : public RmsNorm { aclnnRmsNorm(arena.buf, ws_size_, executor_, stream); if (needs_out_staging_) { - UnpackRows(out.data(), out_data, out_shape_, out_strides_, stream); + CopyTensor(strided_out_cache_.get(out.data()), t_out, stream); } } private: - int64_t RowOffset(int64_t row, const Tensor::Shape& shape, - const Tensor::Strides& strides) const { - int64_t remaining = row; - int64_t offset = 0; - for (int64_t axis = static_cast(shape.size()) - 2; axis >= 0; - --axis) { - auto coordinate = remaining % static_cast(shape[axis]); - remaining /= static_cast(shape[axis]); - offset += coordinate * static_cast(strides[axis]); - } - return offset; + void CopyTensor(aclTensor* dst, const aclTensor* src, + aclrtStream stream) const { + // CANN consumes copy executors even when marked repeatable. + aclOpExecutor* executor = nullptr; + uint64_t workspace_size = 0; + auto status = + aclnnInplaceCopyGetWorkspaceSize(dst, src, &workspace_size, &executor); + assert(status == ACL_SUCCESS && "`RmsNorm` copy workspace query failed."); + auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size, + "rms_norm_copy"); + status = aclnnInplaceCopy(arena.buf, workspace_size, executor, stream); + assert(status == ACL_SUCCESS && "`RmsNorm` staging copy failed."); + (void)status; } - void PackRows(void* dst, const void* src, const Tensor::Shape& shape, - const Tensor::Strides& strides, aclrtStream stream) const { - for (int64_t row = 0; row < logical_rows_; ++row) { - auto offset = RowOffset(row, shape, strides); - auto ret = aclrtMemcpyAsync( - static_cast(dst) + row * row_bytes_, row_bytes_, - static_cast(src) + offset * element_size_, row_bytes_, - ACL_MEMCPY_DEVICE_TO_DEVICE, stream); - assert(ret == ACL_SUCCESS && "`RmsNorm` input pack failed"); - } - } + mutable ascend::AclTensorCache strided_in_cache_; - void UnpackRows(void* dst, const void* src, const Tensor::Shape& shape, - const Tensor::Strides& strides, aclrtStream stream) const { - for (int64_t row = 0; row < logical_rows_; ++row) { - auto offset = RowOffset(row, shape, strides); - auto ret = aclrtMemcpyAsync( - static_cast(dst) + offset * element_size_, row_bytes_, - static_cast(src) + row * row_bytes_, row_bytes_, - ACL_MEMCPY_DEVICE_TO_DEVICE, stream); - assert(ret == ACL_SUCCESS && "`RmsNorm` output unpack failed"); - } - } + mutable ascend::AclTensorCache strided_out_cache_; mutable ascend::AclTensorCache norm_in_cache_; mutable ascend::AclTensorCache weight_cache_; @@ -206,8 +191,6 @@ class Operator : public RmsNorm { bool needs_out_staging_{false}; bool needs_weight_cast_{false}; int64_t logical_rows_{0}; - uint64_t element_size_{0}; - uint64_t row_bytes_{0}; uint64_t tensor_bytes_{0}; void* weight_fp32_data_{nullptr}; mutable aclOpExecutor* cast_exec_{nullptr}; diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 707f9a85f..8e746e6d0 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -96,6 +96,93 @@ def test_rms_norm_non_default_stream(device, implementation_index): accelerator.synchronize() +@pytest.mark.parametrize("seq_len", (1, 36, 333)) +@pytest.mark.parametrize("heads, offset", ((16, 0), (8, 2048))) +@pytest.mark.parametrize("strided_out", (False, True)) +@pytest.mark.parametrize("transpose_rows", (False, True)) +def test_rms_norm_qkv_slice( + seq_len, heads, offset, strided_out, transpose_rows, device, implementation_index +): + if device != "npu": + pytest.skip("regression for Ascend RMSNorm staging") + + width = heads * 128 + weight_cpu = torch.randn(128, dtype=torch.float16) + weight = weight_cpu.to(device) + + # Reuse cached operators with different allocations and values. + for _ in range(3): + source_cpu = torch.randn(seq_len, 4096, dtype=torch.float16) + source = source_cpu.to(device) + input = source[:, offset : offset + width].view(seq_len, heads, 128) + expected = _torch_rms_norm( + source_cpu[:, offset : offset + width].reshape(seq_len, heads, 128), + weight_cpu, + ) + + if strided_out: + backing = torch.full( + (seq_len, width + 256), -123.0, dtype=torch.float16, device=device + ) + out = backing[:, 128 : 128 + width].view(seq_len, heads, 128) + else: + out = torch.empty( + (seq_len, heads, 128), dtype=torch.float16, device=device + ) + + if transpose_rows: + input = input.transpose(0, 1) + out = out.transpose(0, 1) + expected = expected.transpose(0, 1) + + _rms_norm(input, weight, out=out, implementation_index=implementation_index) + torch.testing.assert_close(out.cpu(), expected, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(source.cpu(), source_cpu, rtol=0, atol=0) + + if strided_out: + actual_backing = backing.cpu() + expected_backing = torch.full( + (seq_len, width + 256), -123.0, dtype=torch.float16 + ) + + if transpose_rows: + expected = expected.transpose(0, 1) + + expected_backing[:, 128 : 128 + width] = expected.reshape(seq_len, width) + torch.testing.assert_close( + actual_backing, expected_backing, rtol=1e-2, atol=1e-2 + ) + torch.testing.assert_close( + actual_backing[:, :128], expected_backing[:, :128], rtol=0, atol=0 + ) + torch.testing.assert_close( + actual_backing[:, -128:], expected_backing[:, -128:], rtol=0, atol=0 + ) + + +def test_rms_norm_strided_npu_stream(device, implementation_index): + if device != "npu": + pytest.skip("regression for Ascend RMSNorm staging streams") + + source_cpu = torch.randn(36, 4096, dtype=torch.float16) + weight_cpu = torch.randn(128, dtype=torch.float16) + expected = _torch_rms_norm( + source_cpu[:, 2048:3072].reshape(36, 8, 128), weight_cpu + ) + stream = torch.npu.Stream() + + with torch.npu.stream(stream): + source = source_cpu.to(device) + weight = weight_cpu.to(device) + backing = torch.full((36, 1280), -123.0, dtype=torch.float16, device=device) + input = source[:, 2048:3072].view(36, 8, 128) + out = backing[:, 128:1152].view(36, 8, 128) + _rms_norm(input, weight, out=out, implementation_index=implementation_index) + + stream.synchronize() + torch.testing.assert_close(out.cpu(), expected, rtol=1e-2, atol=1e-2) + + def _rms_norm(input, weight, *, eps=1e-6, out=None, implementation_index=0): infini.ops.rms_norm( input, From a3e476aae35ea956c62f9d22a030c841d8ee5e2f Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 22:57:57 +0800 Subject: [PATCH 07/11] fix(ascend): pad 310P dense prefill for prompt attention --- .../ops/flash_attn_varlen_func/kernel.h | 101 +++++++++++++++++- tests/test_flash_attn_varlen_func.py | 33 ++++++ 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h index 52a054b7a..3ac4009d4 100644 --- a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h +++ b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h @@ -14,8 +14,10 @@ #include "acl/acl.h" #include "aclnn/aclnn_base.h" #include "aclnnop/aclnn_fused_infer_attention_score_v4.h" +#include "aclnnop/aclnn_prompt_flash_attention_v3.h" #include "base/flash_attn_varlen_func.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -58,6 +60,33 @@ class Operator q_lengths_i64_.resize(sequence_count_); k_lengths_i64_.resize(sequence_count_); + use_prompt_310p_ = ascend::IsAscend310P() && !block_table.has_value() && + sequence_count_ == 1 && q.size(0) > 0 && + q.dtype() == DataType::kFloat16 && + q.size(2) == v.size(2) && q.size(0) == k.size(0) && + q.size(0) == v.size(0) && causal && + window_size == std::vector({-1, -1}); + if (use_prompt_310p_) { + // The 310P prompt kernel's causal mask tiles require aligned sequences. + constexpr int64_t kPromptSequenceAlignment = 256; + padded_sequence_ = ((q.size(0) + kPromptSequenceAlignment - 1) / + kPromptSequenceAlignment) * + kPromptSequenceAlignment; + prompt_q_cache_ = ascend::AclTensorCache( + {1, padded_sequence_, static_cast(q.size(1) * q.size(2))}, + ACL_FLOAT16, nullptr); + prompt_k_cache_ = ascend::AclTensorCache( + {1, padded_sequence_, static_cast(k.size(1) * k.size(2))}, + ACL_FLOAT16, nullptr); + prompt_v_cache_ = ascend::AclTensorCache( + {1, padded_sequence_, static_cast(v.size(1) * v.size(2))}, + ACL_FLOAT16, nullptr); + prompt_out_cache_ = ascend::AclTensorCache( + {1, padded_sequence_, + static_cast(out.size(1) * out.size(2))}, + ACL_FLOAT16, nullptr); + } + if (causal || window_size != std::vector({-1, -1})) { InitializeAttentionMask(); } @@ -119,6 +148,15 @@ class Operator auto actual_seq_lengths_kv = MakeSequenceLengths(cu_seqlens_k, k_lengths_i64_, stream, /*cumulative=*/!block_table.has_value()); + if (use_prompt_310p_) { + assert(q_lengths_i64_[0] == static_cast(q.size(0)) && + k_lengths_i64_[0] == static_cast(k.size(0)) && + "310P prompt attention requires one complete sequence"); + RunPrompt310P(q, k, v, softmax_scale, out, stream); + aclDestroyIntArray(actual_seq_lengths); + aclDestroyIntArray(actual_seq_lengths_kv); + return; + } auto t_q = q_cache_.get(const_cast(q.data())); auto t_k = paged ? BuildPagedCacheAclTensor(k) : ascend::BuildAclTensor(k); auto t_v = paged ? BuildPagedCacheAclTensor(v) : ascend::BuildAclTensor(v); @@ -246,7 +284,7 @@ class Operator } void InitializeAttentionMask() { - constexpr int64_t mask_size = 2048; + const int64_t mask_size = use_prompt_310p_ ? padded_sequence_ : 2048; std::vector mask(mask_size * mask_size); for (int64_t row = 0; row < mask_size; ++row) { @@ -268,6 +306,61 @@ class Operator {mask_size, mask_size}, ACL_BOOL, attention_mask_data_); } + void RunPrompt310P(const Tensor q, const Tensor k, const Tensor v, + const std::optional scale, Tensor out, + aclrtStream stream) const { + const size_t q_row_bytes = q.size(1) * q.size(2) * sizeof(uint16_t); + const size_t k_row_bytes = k.size(1) * k.size(2) * sizeof(uint16_t); + const size_t v_row_bytes = v.size(1) * v.size(2) * sizeof(uint16_t); + const size_t q_bytes = padded_sequence_ * q_row_bytes; + const size_t k_bytes = padded_sequence_ * k_row_bytes; + const size_t v_bytes = padded_sequence_ * v_row_bytes; + const size_t total_bytes = q_bytes * 2 + k_bytes + v_bytes; + auto& padded = ascend::GetWorkspacePool().Ensure( + stream, total_bytes, "prompt_310p_padded_qkv_out"); + auto* q_data = static_cast(padded.buf); + auto* k_data = q_data + q_bytes; + auto* v_data = k_data + k_bytes; + auto* out_data = v_data + v_bytes; + ascend::CheckAcl( + aclrtMemsetAsync(padded.buf, total_bytes, 0, total_bytes, stream), + "zero padded prompt QKV"); + ascend::CheckAcl( + aclrtMemcpyAsync(q_data, q_bytes, q.data(), q.size(0) * q_row_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "pad prompt Q"); + ascend::CheckAcl( + aclrtMemcpyAsync(k_data, k_bytes, k.data(), k.size(0) * k_row_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "pad prompt K"); + ascend::CheckAcl( + aclrtMemcpyAsync(v_data, v_bytes, v.data(), v.size(0) * v_row_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "pad prompt V"); + char layout[] = "BSH"; + uint64_t workspace_bytes = 0; + aclOpExecutor* executor = nullptr; + const auto status = aclnnPromptFlashAttentionV3GetWorkspaceSize( + prompt_q_cache_.get(q_data), prompt_k_cache_.get(k_data), + prompt_v_cache_.get(v_data), nullptr, + attention_mask_cache_.get(attention_mask_data_), nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, nullptr, q.size(1), + scale.value_or(1.0 / std::sqrt(q.size(2))), + std::numeric_limits::max(), 0, layout, k.size(1), 0, 0, + prompt_out_cache_.get(out_data), &workspace_bytes, &executor); + ascend::CheckAcl(status, "query 310P PromptFlashAttentionV3 workspace"); + auto& workspace = ascend::GetWorkspacePool().Ensure( + stream, workspace_bytes, "prompt_310p_attention_workspace"); + ascend::CheckAcl(aclnnPromptFlashAttentionV3(workspace.buf, workspace_bytes, + executor, stream), + "run 310P PromptFlashAttentionV3"); + const auto output_bytes = q.size(0) * q_row_bytes; + ascend::CheckAcl( + aclrtMemcpyAsync(out.data(), output_bytes, out_data, output_bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "unpad prompt output"); + } + // Dense TND K/V uses cumulative endpoints, while paged K/V uses per-batch // lengths. Query is always dense TND and therefore cumulative. aclIntArray* MakeSequenceLengths(const Tensor cu_seqlens, @@ -299,6 +392,12 @@ class Operator mutable std::vector q_lengths_i64_; mutable std::vector k_lengths_i64_; void* attention_mask_data_{nullptr}; + bool use_prompt_310p_{false}; + int64_t padded_sequence_{0}; + mutable ascend::AclTensorCache prompt_q_cache_; + mutable ascend::AclTensorCache prompt_k_cache_; + mutable ascend::AclTensorCache prompt_v_cache_; + mutable ascend::AclTensorCache prompt_out_cache_; mutable std::array input_layout_{'T', 'N', 'D', '\0'}; }; diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 98ad63f98..d3f742a41 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -749,6 +749,39 @@ def _cumulative_lengths(lengths, device): return torch.tensor(values, dtype=torch.int32, device=device) +@pytest.mark.parametrize("sequence_length", (36, 128, 256, 333, 512)) +@pytest.mark.parametrize("seed", (0, 1)) +def test_ascend_310p_dense_prefill_padding( + device, implementation_index, sequence_length, seed +): + if device != "npu" or implementation_index != 0: + pytest.skip("310P native prefill regression") + if not torch.npu.get_device_name(0).startswith("Ascend310P"): + pytest.skip("310P-specific prompt kernel") + + generator = torch.Generator().manual_seed(seed) + q_cpu = torch.randn((sequence_length, 16, 128), generator=generator).half() + k_cpu = torch.randn((sequence_length, 8, 128), generator=generator).half() + v_cpu = torch.randn((sequence_length, 8, 128), generator=generator).half() + q, k, v = (tensor.to(device) for tensor in (q_cpu, k_cpu, v_cpu)) + lengths = torch.tensor([0, sequence_length], dtype=torch.int32, device=device) + out = torch.empty_like(q) + expected = _reference_varlen_attention( + q_cpu, k_cpu, v_cpu, (sequence_length,), (sequence_length,), + 128 ** -0.5, True, (-1, -1) + ) + for _ in range(2): + infini.ops.flash_attn_varlen_func( + q, k, v, lengths, lengths, None, None, + sequence_length, sequence_length, 0.0, 128 ** -0.5, + True, (-1, -1), 0.0, False, False, out, None, None, + stream=get_stream(q.device), implementation_index=0, + ) + actual = out.cpu() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + + def _pack_varlen_softmax_lse(softmax_lse, q_lens): if softmax_lse.ndim == 2: return softmax_lse From 1290ae217668dc009f7fd0e071329b9117bc89c3 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 23:00:01 +0800 Subject: [PATCH 08/11] refactor(ascend): check ACL calls in RMSNorm and RoPE --- src/native/ascend/ops/rms_norm/kernel.h | 65 +++++++----- .../ascend/ops/rotary_embedding/kernel.h | 99 ++++++++++--------- 2 files changed, 97 insertions(+), 67 deletions(-) diff --git a/src/native/ascend/ops/rms_norm/kernel.h b/src/native/ascend/ops/rms_norm/kernel.h index d460613e4..0396e2105 100644 --- a/src/native/ascend/ops/rms_norm/kernel.h +++ b/src/native/ascend/ops/rms_norm/kernel.h @@ -11,6 +11,7 @@ #include "aclnnop/aclnn_cast.h" #include "base/rms_norm.h" #include "native/ascend/common.h" +#include "native/ascend/error_.h" #include "native/ascend/workspace_pool_.h" #include "operator.h" @@ -54,8 +55,7 @@ class Operator : public RmsNorm { auto fp32_bytes = static_cast(dim_) * sizeof(float); auto ret = aclrtMalloc(&weight_fp32_data_, fp32_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "`RmsNorm` Ascend path failed to allocate cast weight"); + ascend::CheckAcl(ret, "RmsNorm cast weight allocation"); weight_fp32_cache_ = ascend::AclTensorCache({static_cast(dim_)}, ACL_FLOAT, weight_fp32_data_); } @@ -78,9 +78,12 @@ class Operator : public RmsNorm { if (needs_weight_cast_) { auto t_weight_src = weight_cache_.get(nullptr); auto t_weight_dst = weight_fp32_cache_.get(weight_fp32_data_); - aclnnCastGetWorkspaceSize(t_weight_src, ACL_FLOAT, t_weight_dst, - &cast_ws_, &cast_exec_); - aclSetAclOpExecutorRepeatable(cast_exec_); + ascend::CheckAcl( + aclnnCastGetWorkspaceSize(t_weight_src, ACL_FLOAT, t_weight_dst, + &cast_ws_, &cast_exec_), + "RmsNorm weight cast workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(cast_exec_), + "RmsNorm repeatable cast executor"); t_weight = t_weight_dst; } else { t_weight = weight_cache_.get(nullptr); @@ -90,9 +93,13 @@ class Operator : public RmsNorm { rstd_shape_.data(), static_cast(rstd_shape_.size()), ACL_FLOAT, rstd_strides_.data(), 0, ACL_FORMAT_ND, rstd_shape_.data(), static_cast(rstd_shape_.size()), nullptr); - aclnnRmsNormGetWorkspaceSize(t_in, t_weight, eps, t_out, rstd_tensor_, - &ws_size_, &executor_); - aclSetAclOpExecutorRepeatable(executor_); + assert(rstd_tensor_ != nullptr && "`RmsNorm` rstd tensor creation failed."); + ascend::CheckAcl( + aclnnRmsNormGetWorkspaceSize(t_in, t_weight, eps, t_out, rstd_tensor_, + &ws_size_, &executor_), + "RmsNorm workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(executor_), + "RmsNorm repeatable executor"); total_workspace_size_ = ws_size_ + rstd_size_; } @@ -103,7 +110,10 @@ class Operator : public RmsNorm { weight_cache_.release(); norm_out_cache_.release(); weight_fp32_cache_.release(); - if (weight_fp32_data_) aclrtFree(weight_fp32_data_); + if (weight_fp32_data_) { + ascend::CheckAcl(aclrtFree(weight_fp32_data_), + "RmsNorm cast weight release"); + } // `rstd_tensor_` leaks with the executor at shutdown (see `64c367c`). } @@ -135,11 +145,15 @@ class Operator : public RmsNorm { if (needs_weight_cast_) { auto t_weight_src = weight_cache_.get(const_cast(weight.data())); auto t_weight_dst = weight_fp32_cache_.get(weight_fp32_data_); - AclSetTensorAddr(cast_exec_, 0, t_weight_src, - const_cast(weight.data())); - AclSetTensorAddr(cast_exec_, 1, t_weight_dst, weight_fp32_data_); + ascend::CheckAcl(AclSetTensorAddr(cast_exec_, 0, t_weight_src, + const_cast(weight.data())), + "RmsNorm cast input address"); + ascend::CheckAcl( + AclSetTensorAddr(cast_exec_, 1, t_weight_dst, weight_fp32_data_), + "RmsNorm cast output address"); auto& cast_arena = ascend::GetWorkspacePool().Ensure(stream, cast_ws_); - aclnnCast(cast_arena.buf, cast_ws_, cast_exec_, stream); + ascend::CheckAcl(aclnnCast(cast_arena.buf, cast_ws_, cast_exec_, stream), + "RmsNorm weight cast execution"); t_weight = t_weight_dst; weight_data = weight_fp32_data_; } else { @@ -150,13 +164,19 @@ class Operator : public RmsNorm { auto& arena = ascend::GetWorkspacePool().Ensure( stream, total_workspace_size_, "rms_norm"); auto* rstd_data = static_cast(arena.buf) + ws_size_; - aclSetRawTensorAddr(rstd_tensor_, rstd_data); - - AclSetTensorAddr(executor_, 0, t_in, input_data); - AclSetTensorAddr(executor_, 1, t_weight, weight_data); - AclSetTensorAddr(executor_, 2, t_out, out_data); - AclSetTensorAddr(executor_, 3, rstd_tensor_, rstd_data); - aclnnRmsNorm(arena.buf, ws_size_, executor_, stream); + ascend::CheckAcl(aclSetRawTensorAddr(rstd_tensor_, rstd_data), + "RmsNorm rstd address"); + + ascend::CheckAcl(AclSetTensorAddr(executor_, 0, t_in, input_data), + "RmsNorm input address"); + ascend::CheckAcl(AclSetTensorAddr(executor_, 1, t_weight, weight_data), + "RmsNorm weight address"); + ascend::CheckAcl(AclSetTensorAddr(executor_, 2, t_out, out_data), + "RmsNorm output address"); + ascend::CheckAcl(AclSetTensorAddr(executor_, 3, rstd_tensor_, rstd_data), + "RmsNorm rstd output address"); + ascend::CheckAcl(aclnnRmsNorm(arena.buf, ws_size_, executor_, stream), + "RmsNorm execution"); if (needs_out_staging_) { CopyTensor(strided_out_cache_.get(out.data()), t_out, stream); @@ -171,12 +191,11 @@ class Operator : public RmsNorm { uint64_t workspace_size = 0; auto status = aclnnInplaceCopyGetWorkspaceSize(dst, src, &workspace_size, &executor); - assert(status == ACL_SUCCESS && "`RmsNorm` copy workspace query failed."); + ascend::CheckAcl(status, "RmsNorm copy workspace query"); auto& arena = ascend::GetWorkspacePool().Ensure(stream, workspace_size, "rms_norm_copy"); status = aclnnInplaceCopy(arena.buf, workspace_size, executor, stream); - assert(status == ACL_SUCCESS && "`RmsNorm` staging copy failed."); - (void)status; + ascend::CheckAcl(status, "RmsNorm staging copy"); } mutable ascend::AclTensorCache strided_in_cache_; diff --git a/src/native/ascend/ops/rotary_embedding/kernel.h b/src/native/ascend/ops/rotary_embedding/kernel.h index d82a111b4..9c1cf6adf 100644 --- a/src/native/ascend/ops/rotary_embedding/kernel.h +++ b/src/native/ascend/ops/rotary_embedding/kernel.h @@ -53,12 +53,10 @@ class Operator static_cast(max_seq_len_ * head_dim) * element_size_; auto ret = aclrtMalloc(&cos_table_data_, table_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate cosine table"); + ascend::CheckAcl(ret, "RoPE cosine table allocation"); ret = aclrtMalloc(&sin_table_data_, table_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate sine table"); + ascend::CheckAcl(ret, "RoPE sine table allocation"); UploadCosSinCache(cos_sin_cache); cos_sin_cache_data_ = cos_sin_cache.data(); @@ -66,11 +64,9 @@ class Operator const auto gathered_bytes = static_cast(num_tokens * head_dim) * element_size_; ret = aclrtMalloc(&cos_data_, gathered_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate gathered cosine"); + ascend::CheckAcl(ret, "RoPE gathered cosine allocation"); ret = aclrtMalloc(&sin_data_, gathered_bytes, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate gathered sine"); + ascend::CheckAcl(ret, "RoPE gathered sine allocation"); cos_table_cache_ = ascend::AclTensorCache({max_seq_len_, head_dim}, acl_dtype, cos_table_data_); @@ -89,16 +85,14 @@ class Operator query_bytes_ = static_cast(query.numel()) * element_size_; ret = aclrtMalloc(&query_out_data_, query_bytes_, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate query output"); + ascend::CheckAcl(ret, "RoPE query output allocation"); MakeChunks(query, num_heads_, query_token_stride_, query_head_stride_, query_out_data_, query_chunks_); if (key.has_value()) { key_bytes_ = static_cast(key->numel()) * element_size_; ret = aclrtMalloc(&key_out_data_, key_bytes_, ACL_MEM_MALLOC_NORMAL_ONLY); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to allocate key output"); + ascend::CheckAcl(ret, "RoPE key output allocation"); MakeChunks(*key, num_kv_heads_, key_token_stride_, key_head_stride_, key_out_data_, key_chunks_); } @@ -120,12 +114,18 @@ class Operator chunk.output.release(); } } - if (cos_table_data_) aclrtFree(cos_table_data_); - if (sin_table_data_) aclrtFree(sin_table_data_); - if (cos_data_) aclrtFree(cos_data_); - if (sin_data_) aclrtFree(sin_data_); - if (query_out_data_) aclrtFree(query_out_data_); - if (key_out_data_) aclrtFree(key_out_data_); + if (cos_table_data_) + ascend::CheckAcl(aclrtFree(cos_table_data_), "RoPE cosine table release"); + if (sin_table_data_) + ascend::CheckAcl(aclrtFree(sin_table_data_), "RoPE sine table release"); + if (cos_data_) + ascend::CheckAcl(aclrtFree(cos_data_), "RoPE gathered cosine release"); + if (sin_data_) + ascend::CheckAcl(aclrtFree(sin_data_), "RoPE gathered sine release"); + if (query_out_data_) + ascend::CheckAcl(aclrtFree(query_out_data_), "RoPE query output release"); + if (key_out_data_) + ascend::CheckAcl(aclrtFree(key_out_data_), "RoPE key output release"); } void operator()(const Tensor positions, Tensor query, @@ -149,32 +149,42 @@ class Operator auto t_sin_out = sin_out_cache_.get(sin_data_); if (!cos_index_executor_) { - aclnnIndexSelectGetWorkspaceSize(t_cos_table, 0, t_positions, t_cos_out, - &cos_index_ws_size_, - &cos_index_executor_); - aclSetAclOpExecutorRepeatable(cos_index_executor_); + ascend::CheckAcl(aclnnIndexSelectGetWorkspaceSize( + t_cos_table, 0, t_positions, t_cos_out, + &cos_index_ws_size_, &cos_index_executor_), + "RoPE cosine gather workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(cos_index_executor_), + "RoPE repeatable cosine gather executor"); } else { - aclSetInputTensorAddr(cos_index_executor_, 1, t_positions, - const_cast(positions.data())); + ascend::CheckAcl( + aclSetInputTensorAddr(cos_index_executor_, 1, t_positions, + const_cast(positions.data())), + "RoPE cosine positions address"); } if (!sin_index_executor_) { - aclnnIndexSelectGetWorkspaceSize(t_sin_table, 0, t_positions, t_sin_out, - &sin_index_ws_size_, - &sin_index_executor_); - aclSetAclOpExecutorRepeatable(sin_index_executor_); + ascend::CheckAcl(aclnnIndexSelectGetWorkspaceSize( + t_sin_table, 0, t_positions, t_sin_out, + &sin_index_ws_size_, &sin_index_executor_), + "RoPE sine gather workspace query"); + ascend::CheckAcl(aclSetAclOpExecutorRepeatable(sin_index_executor_), + "RoPE repeatable sine gather executor"); } else { - aclSetInputTensorAddr(sin_index_executor_, 1, t_positions, - const_cast(positions.data())); + ascend::CheckAcl( + aclSetInputTensorAddr(sin_index_executor_, 1, t_positions, + const_cast(positions.data())), + "RoPE sine positions address"); } auto index_ws_size = std::max(cos_index_ws_size_, sin_index_ws_size_); auto& index_arena = ascend::GetWorkspacePool().Ensure(stream, index_ws_size); - aclnnIndexSelect(index_arena.buf, cos_index_ws_size_, cos_index_executor_, - stream); - aclnnIndexSelect(index_arena.buf, sin_index_ws_size_, sin_index_executor_, - stream); + ascend::CheckAcl(aclnnIndexSelect(index_arena.buf, cos_index_ws_size_, + cos_index_executor_, stream), + "RoPE cosine gather execution"); + ascend::CheckAcl(aclnnIndexSelect(index_arena.buf, sin_index_ws_size_, + sin_index_executor_, stream), + "RoPE sine gather execution"); auto t_cos = cos_rotary_cache_.get(cos_data_); auto t_sin = sin_rotary_cache_.get(sin_data_); @@ -249,10 +259,14 @@ class Operator ascend::CheckAcl(aclSetAclOpExecutorRepeatable(executor), "aclSetAclOpExecutorRepeatable RoPE"); } else { - aclSetInputTensorAddr(executor, 0, input, input_data); - aclSetInputTensorAddr(executor, 1, cos, cos_data_); - aclSetInputTensorAddr(executor, 2, sin, sin_data_); - aclSetOutputTensorAddr(executor, 0, output, output_data); + ascend::CheckAcl(aclSetInputTensorAddr(executor, 0, input, input_data), + "RoPE input address"); + ascend::CheckAcl(aclSetInputTensorAddr(executor, 1, cos, cos_data_), + "RoPE cosine address"); + ascend::CheckAcl(aclSetInputTensorAddr(executor, 2, sin, sin_data_), + "RoPE sine address"); + ascend::CheckAcl(aclSetOutputTensorAddr(executor, 0, output, output_data), + "RoPE output address"); } auto& rotary_arena = @@ -287,8 +301,7 @@ class Operator auto ret = aclrtMemcpy(packed.data(), packed_bytes, cos_sin_cache.data(), packed_bytes, ACL_MEMCPY_DEVICE_TO_HOST); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to read cos/sin cache"); + ascend::CheckAcl(ret, "RoPE cos/sin cache download"); for (int64_t position = 0; position < max_seq_len_; ++position) { for (int64_t index = 0; index < head_size_; ++index) { const auto destination_index = position * head_size_ + index; @@ -312,12 +325,10 @@ class Operator ret = aclrtMemcpy(cos_table_data_, table_bytes, cosine.data(), table_bytes, ACL_MEMCPY_HOST_TO_DEVICE); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to upload cosine table"); + ascend::CheckAcl(ret, "RoPE cosine table upload"); ret = aclrtMemcpy(sin_table_data_, table_bytes, sine.data(), table_bytes, ACL_MEMCPY_HOST_TO_DEVICE); - assert(ret == ACL_SUCCESS && - "Ascend `RotaryEmbedding` failed to upload sine table"); + ascend::CheckAcl(ret, "RoPE sine table upload"); } void WriteOne(uint8_t* destination) const { From c7a4a4922e9255003772201187cf3cf186b45a43 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Wed, 21 Oct 2026 23:00:55 +0800 Subject: [PATCH 09/11] test(ascend): align paged prefill with the public page contract --- tests/test_flash_attn_varlen_func.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index d3f742a41..c77f7cb7a 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -245,12 +245,13 @@ def test_ascend_paged_prefill_follows_nontrivial_block_table( if device != "npu" or implementation_index != 0: pytest.skip("coverage requires the native Ascend provider") - q_lens = (128, 256) + q_lens = (256, 512) k_lens = q_lens num_heads = 8 num_kv_heads = 2 head_dim = 64 - page_size = 128 + # Preserve cross-page coverage within the public page size contract. + page_size = 256 block_rows = ((2, 0), (1, 3)) num_blocks = 4 From a005789095e66151b82b37a7f1d380b53a427e80 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Thu, 22 Oct 2026 11:04:43 +0800 Subject: [PATCH 10/11] fix(ascend): handle variable and paged attention on 310P --- .../ops/flash_attn_varlen_func/kernel.h | 184 +++++++++++++++++- tests/test_flash_attn_varlen_func.py | 73 ++++++- 2 files changed, 255 insertions(+), 2 deletions(-) diff --git a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h index 3ac4009d4..ef09039f1 100644 --- a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h +++ b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h @@ -87,7 +87,9 @@ class Operator ACL_FLOAT16, nullptr); } - if (causal || window_size != std::vector({-1, -1})) { + use_prompt_fallback_310p_ = ascend::IsAscend310P() && !use_prompt_310p_; + if (!use_prompt_fallback_310p_ && + (causal || window_size != std::vector({-1, -1}))) { InitializeAttentionMask(); } } @@ -148,6 +150,13 @@ class Operator auto actual_seq_lengths_kv = MakeSequenceLengths(cu_seqlens_k, k_lengths_i64_, stream, /*cumulative=*/!block_table.has_value()); + if (use_prompt_fallback_310p_) { + RunVariablePrompt310P(q, k, v, block_table, softmax_scale, causal, + window_size, out, stream); + aclDestroyIntArray(actual_seq_lengths); + aclDestroyIntArray(actual_seq_lengths_kv); + return; + } if (use_prompt_310p_) { assert(q_lengths_i64_[0] == static_cast(q.size(0)) && k_lengths_i64_[0] == static_cast(k.size(0)) && @@ -252,6 +261,178 @@ class Operator } private: + void RunVariablePrompt310P(const Tensor q, const Tensor k, const Tensor v, + const std::optional& block_table, + const std::optional scale, bool causal, + const std::vector& window_size, + Tensor out, aclrtStream stream) const { + assert(q.dtype() == DataType::kFloat16 && + "310P prompt attention requires FP16 tensors."); + const bool paged = block_table.has_value(); + const int64_t kv_heads = paged ? k.size(2) : k.size(1); + const size_t q_row_bytes = q.size(1) * q.size(2) * sizeof(uint16_t); + const size_t k_row_bytes = kv_heads * q.size(2) * sizeof(uint16_t); + const size_t v_row_bytes = + kv_heads * v.size(v.ndim() - 1) * sizeof(uint16_t); + std::vector blocks; + if (paged) { + blocks.resize(block_table->numel()); + ascend::CheckAcl( + aclrtMemcpyAsync(blocks.data(), blocks.size() * sizeof(int32_t), + block_table->data(), blocks.size() * sizeof(int32_t), + ACL_MEMCPY_DEVICE_TO_HOST, stream), + "read 310P attention block table"); + ascend::CheckAcl(aclrtSynchronizeStream(stream), + "synchronize 310P attention block table"); + } + int64_t q_begin = 0; + int64_t k_begin = 0; + for (size_t batch = 0; batch < sequence_count_; ++batch) { + const int64_t q_end = q_lengths_i64_[batch]; + const int64_t q_length = q_end - q_begin; + const int64_t k_length = + paged ? k_lengths_i64_[batch] : k_lengths_i64_[batch] - k_begin; + assert(q_length >= 0 && k_length >= 0 && + q_end <= static_cast(q.size(0)) && + "Invalid 310P attention sequence lengths."); + auto* output = static_cast(out.data()) + q_begin * q_row_bytes; + if (q_length == 0 || k_length == 0) { + if (q_length > 0) { + ascend::CheckAcl(aclrtMemsetAsync(output, q_length * q_row_bytes, 0, + q_length * q_row_bytes, stream), + "zero 310P attention with empty keys"); + } + q_begin = q_end; + if (!paged) k_begin += k_length; + continue; + } + // The 310P prompt kernel requires aligned sequence tiles. An explicit + // mask preserves right-down causal/window semantics after padding. + constexpr int64_t kAlignment = 256; + const int64_t padded_length = + ((std::max(q_length, k_length) + kAlignment - 1) / kAlignment) * + kAlignment; + const size_t q_bytes = padded_length * q_row_bytes; + const size_t k_bytes = padded_length * k_row_bytes; + const size_t v_bytes = padded_length * v_row_bytes; + const size_t out_bytes = padded_length * q_row_bytes; + const size_t total_bytes = q_bytes + k_bytes + v_bytes + out_bytes; + auto& staging = ascend::GetWorkspacePool().Ensure( + stream, total_bytes, "variable_prompt_310p_qkv_out"); + auto* q_data = static_cast(staging.buf); + auto* k_data = q_data + q_bytes; + auto* v_data = k_data + k_bytes; + auto* out_data = v_data + v_bytes; + ascend::CheckAcl( + aclrtMemsetAsync(staging.buf, total_bytes, 0, total_bytes, stream), + "zero 310P variable attention staging"); + const auto copy = [stream](void* dst, const void* src, size_t bytes) { + ascend::CheckAcl(aclrtMemcpyAsync(dst, bytes, src, bytes, + ACL_MEMCPY_DEVICE_TO_DEVICE, stream), + "stage 310P variable attention"); + }; + copy(q_data, static_cast(q.data()) + q_begin * q_row_bytes, + q_length * q_row_bytes); + if (paged) { + const int64_t page_size = k.size(1); + const int64_t page_count = (k_length + page_size - 1) / page_size; + assert(page_count <= static_cast(block_table->size(1)) && + "310P attention block table is too short."); + for (int64_t page = 0; page < page_count; ++page) { + const int64_t block = blocks[batch * block_table->size(1) + page]; + assert(block >= 0 && block < static_cast(k.size(0)) && + "Invalid active 310P attention cache block."); + const int64_t rows = std::min(page_size, k_length - page * page_size); + copy(k_data + page * page_size * k_row_bytes, + static_cast(k.data()) + + block * page_size * k_row_bytes, + rows * k_row_bytes); + copy(v_data + page * page_size * v_row_bytes, + static_cast(v.data()) + + block * page_size * v_row_bytes, + rows * v_row_bytes); + } + } else { + assert(k_begin + k_length <= static_cast(k.size(0)) && + "310P attention key lengths exceed storage."); + copy(k_data, static_cast(k.data()) + k_begin * k_row_bytes, + k_length * k_row_bytes); + copy(v_data, static_cast(v.data()) + k_begin * v_row_bytes, + k_length * v_row_bytes); + } + std::vector mask(padded_length * padded_length, 1); + std::vector empty_rows(q_length, true); + for (int64_t row = 0; row < q_length; ++row) { + const int64_t center = row + k_length - q_length; + for (int64_t column = 0; column < k_length; ++column) { + const bool masked = + (causal && column > center) || + (window_size[0] >= 0 && column < center - window_size[0]) || + (window_size[1] >= 0 && column > center + window_size[1]); + mask[row * padded_length + column] = masked; + if (!masked) empty_rows[row] = false; + } + } + // Avoid all-masked softmax rows inside CANN; discard these results below. + for (int64_t row = 0; row < padded_length; ++row) { + if (row >= q_length || empty_rows[row]) mask[row * padded_length] = 0; + } + auto& mask_buffer = ascend::GetWorkspacePool().Ensure( + stream, mask.size(), "variable_prompt_310p_mask"); + ascend::CheckAcl( + aclrtMemcpyAsync(mask_buffer.buf, mask.size(), mask.data(), + mask.size(), ACL_MEMCPY_HOST_TO_DEVICE, stream), + "upload 310P variable attention mask"); + ascend::AclTensorCache t_q( + {1, padded_length, static_cast(q.size(1) * q.size(2))}, + ACL_FLOAT16, q_data); + ascend::AclTensorCache t_k( + {1, padded_length, kv_heads * static_cast(q.size(2))}, + ACL_FLOAT16, k_data); + ascend::AclTensorCache t_v( + {1, padded_length, + kv_heads * static_cast(v.size(v.ndim() - 1))}, + ACL_FLOAT16, v_data); + ascend::AclTensorCache t_out( + {1, padded_length, static_cast(q.size(1) * q.size(2))}, + ACL_FLOAT16, out_data); + ascend::AclTensorCache t_mask({padded_length, padded_length}, ACL_BOOL, + mask_buffer.buf); + uint64_t workspace_bytes = 0; + aclOpExecutor* executor = nullptr; + char layout[] = "BSH"; + ascend::CheckAcl( + aclnnPromptFlashAttentionV3GetWorkspaceSize( + t_q.get(q_data), t_k.get(k_data), t_v.get(v_data), nullptr, + t_mask.get(mask_buffer.buf), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, q.size(1), + scale.value_or(1.0 / std::sqrt(q.size(2))), + std::numeric_limits::max(), + std::numeric_limits::max(), layout, kv_heads, 0, 0, + t_out.get(out_data), &workspace_bytes, &executor), + "query 310P variable PromptFlashAttentionV3 workspace"); + auto& workspace = ascend::GetWorkspacePool().Ensure( + stream, workspace_bytes, "variable_prompt_310p_workspace"); + ascend::CheckAcl(aclnnPromptFlashAttentionV3( + workspace.buf, workspace_bytes, executor, stream), + "run 310P variable PromptFlashAttentionV3"); + copy(output, out_data, q_length * q_row_bytes); + for (int64_t row = 0; row < q_length; ++row) { + if (empty_rows[row]) { + ascend::CheckAcl( + aclrtMemsetAsync(output + row * q_row_bytes, q_row_bytes, 0, + q_row_bytes, stream), + "zero fully masked 310P attention row"); + } + } + // Keep the host mask and ACL descriptors alive until this sequence ends. + ascend::CheckAcl(aclrtSynchronizeStream(stream), + "synchronize 310P variable prompt attention"); + q_begin = q_end; + if (!paged) k_begin += k_length; + } + } + static void ValidateSupportedOptions( const std::optional& alibi_slopes, bool return_attn_probs, const std::optional& softmax_lse, @@ -393,6 +574,7 @@ class Operator mutable std::vector k_lengths_i64_; void* attention_mask_data_{nullptr}; bool use_prompt_310p_{false}; + bool use_prompt_fallback_310p_{false}; int64_t padded_sequence_{0}; mutable ascend::AclTensorCache prompt_q_cache_; mutable ascend::AclTensorCache prompt_k_cache_; diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index c77f7cb7a..bd975c78b 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -783,6 +783,67 @@ def test_ascend_310p_dense_prefill_padding( torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) +@pytest.mark.parametrize("paged", (False, True)) +@pytest.mark.parametrize( + "causal, window_size", + ((False, (-1, -1)), (True, (-1, -1)), (True, (2, 1))), +) +@pytest.mark.parametrize("head_dim", (64, 128)) +def test_ascend_variable_prompt_reuses_updated_lengths( + device, implementation_index, paged, causal, window_size, head_dim +): + if device != "npu" or implementation_index != 0: + pytest.skip("coverage requires the native Ascend provider") + + generator = torch.Generator().manual_seed(42) + q_lens = (5, 2) + initial_k_lens = (3, 6) + updated_k_lens = (7, 2) + q_cpu = torch.randn((7, 4, head_dim), generator=generator).half() + kv_shape = (3, 256, 2, head_dim) if paged else (9, 2, head_dim) + k_cpu = torch.randn(kv_shape, generator=generator).half() + v_cpu = torch.randn(kv_shape, generator=generator).half() + blocks_cpu = ( + torch.tensor(((2, -1), (1, -1)), dtype=torch.int32) if paged else None + ) + q, k, v = (tensor.to(device) for tensor in (q_cpu, k_cpu, v_cpu)) + block_table = blocks_cpu.to(device) if paged else None + cu_q = _cumulative_lengths(q_lens, device) + cu_k = _cumulative_lengths(initial_k_lens, device) + out = torch.empty_like(q) + stream = torch.npu.Stream() + stream.wait_stream(torch.npu.current_stream()) + + for iteration, k_lens in enumerate( + (initial_k_lens, updated_k_lens, initial_k_lens) + ): + if paged and iteration == 1: + blocks_cpu[0, 0] = 0 + + expected = _reference_varlen_attention( + q_cpu, k_cpu, v_cpu, q_lens, k_lens, None, causal, + window_size, blocks_cpu, + ) + + with torch.npu.stream(stream): + cu_k.copy_(_cumulative_lengths(k_lens, device)) + if paged: + block_table.copy_(blocks_cpu) + + out.fill_(math.nan) + infini.ops.flash_attn_varlen_func( + q, k, v, cu_q, cu_k, None, block_table, max(q_lens), + 7, 0.0, None, causal, window_size, 0.0, False, False, + out, None, None, stream=get_stream(q.device), + implementation_index=0, + ) + actual = out.cpu() + + stream.synchronize() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + + def _pack_varlen_softmax_lse(softmax_lse, q_lens): if softmax_lse.ndim == 2: return softmax_lse @@ -807,6 +868,16 @@ def _reference_varlen_attention( block_table=None, alibi_slopes=None, ): + output_device = q.device + if ( + q.device.type == "npu" + and torch.npu.get_device_name(q.device.index).startswith("Ascend310P") + ): + # The 310P runtime does not provide `NanToNum` for the reference path. + q, k, v = (tensor.cpu() for tensor in (q, k, v)) + block_table = block_table.cpu() if block_table is not None else None + alibi_slopes = alibi_slopes.cpu() if alibi_slopes is not None else None + outputs = [] q_offset = 0 k_offset = 0 @@ -856,7 +927,7 @@ def _reference_varlen_attention( q_offset += q_len k_offset += k_len - return torch.cat(outputs) + return torch.cat(outputs).to(output_device) def _reference_varlen_softmax_lse( From e7d34f701235fa63d1d81d6d236d78750eee3eb5 Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Thu, 22 Oct 2026 11:14:38 +0800 Subject: [PATCH 11/11] test(ascend): format compatibility regression cases --- tests/test_flash_attn_varlen_func.py | 78 +++++++++++++++++++++------ tests/test_fused_add_rms_norm.py | 25 +++++++-- tests/test_reshape_and_cache_flash.py | 20 +++++-- tests/test_rms_norm.py | 8 +-- tests/test_rotary_embedding.py | 33 ++++++++---- tests/test_silu_and_mul.py | 12 +++-- 6 files changed, 130 insertions(+), 46 deletions(-) diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index bd975c78b..085f74073 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -768,15 +768,38 @@ def test_ascend_310p_dense_prefill_padding( lengths = torch.tensor([0, sequence_length], dtype=torch.int32, device=device) out = torch.empty_like(q) expected = _reference_varlen_attention( - q_cpu, k_cpu, v_cpu, (sequence_length,), (sequence_length,), - 128 ** -0.5, True, (-1, -1) + q_cpu, + k_cpu, + v_cpu, + (sequence_length,), + (sequence_length,), + 128**-0.5, + True, + (-1, -1), ) for _ in range(2): infini.ops.flash_attn_varlen_func( - q, k, v, lengths, lengths, None, None, - sequence_length, sequence_length, 0.0, 128 ** -0.5, - True, (-1, -1), 0.0, False, False, out, None, None, - stream=get_stream(q.device), implementation_index=0, + q, + k, + v, + lengths, + lengths, + None, + None, + sequence_length, + sequence_length, + 0.0, + 128**-0.5, + True, + (-1, -1), + 0.0, + False, + False, + out, + None, + None, + stream=get_stream(q.device), + implementation_index=0, ) actual = out.cpu() assert torch.isfinite(actual).all() @@ -803,9 +826,7 @@ def test_ascend_variable_prompt_reuses_updated_lengths( kv_shape = (3, 256, 2, head_dim) if paged else (9, 2, head_dim) k_cpu = torch.randn(kv_shape, generator=generator).half() v_cpu = torch.randn(kv_shape, generator=generator).half() - blocks_cpu = ( - torch.tensor(((2, -1), (1, -1)), dtype=torch.int32) if paged else None - ) + blocks_cpu = torch.tensor(((2, -1), (1, -1)), dtype=torch.int32) if paged else None q, k, v = (tensor.to(device) for tensor in (q_cpu, k_cpu, v_cpu)) block_table = blocks_cpu.to(device) if paged else None cu_q = _cumulative_lengths(q_lens, device) @@ -821,8 +842,15 @@ def test_ascend_variable_prompt_reuses_updated_lengths( blocks_cpu[0, 0] = 0 expected = _reference_varlen_attention( - q_cpu, k_cpu, v_cpu, q_lens, k_lens, None, causal, - window_size, blocks_cpu, + q_cpu, + k_cpu, + v_cpu, + q_lens, + k_lens, + None, + causal, + window_size, + blocks_cpu, ) with torch.npu.stream(stream): @@ -832,9 +860,26 @@ def test_ascend_variable_prompt_reuses_updated_lengths( out.fill_(math.nan) infini.ops.flash_attn_varlen_func( - q, k, v, cu_q, cu_k, None, block_table, max(q_lens), - 7, 0.0, None, causal, window_size, 0.0, False, False, - out, None, None, stream=get_stream(q.device), + q, + k, + v, + cu_q, + cu_k, + None, + block_table, + max(q_lens), + 7, + 0.0, + None, + causal, + window_size, + 0.0, + False, + False, + out, + None, + None, + stream=get_stream(q.device), implementation_index=0, ) actual = out.cpu() @@ -869,9 +914,8 @@ def _reference_varlen_attention( alibi_slopes=None, ): output_device = q.device - if ( - q.device.type == "npu" - and torch.npu.get_device_name(q.device.index).startswith("Ascend310P") + if q.device.type == "npu" and torch.npu.get_device_name(q.device.index).startswith( + "Ascend310P" ): # The 310P runtime does not provide `NanToNum` for the reference path. q, k, v = (tensor.cpu() for tensor in (q, k, v)) diff --git a/tests/test_fused_add_rms_norm.py b/tests/test_fused_add_rms_norm.py index 033900ee8..745902e33 100644 --- a/tests/test_fused_add_rms_norm.py +++ b/tests/test_fused_add_rms_norm.py @@ -73,7 +73,10 @@ def test_fused_add_rms_norm( @pytest.mark.parametrize("width", (4, 8, 16, 128)) @pytest.mark.parametrize("has_weight", (False, True)) def test_fused_add_rms_norm_repeated_shape( - device, implementation_index, width, has_weight, + device, + implementation_index, + width, + has_weight, ): if device != "npu" or implementation_index != 0: pytest.skip("Ascend native repeated-call regression") @@ -81,16 +84,28 @@ def test_fused_add_rms_norm_repeated_shape( for value in (0.25, 0.75): input = torch.full((2, 2, width), value, dtype=torch.float16, device=device) residual = torch.full_like(input, 0.5) - weight = torch.full((width,), 0.75, dtype=torch.float16, device=device) if has_weight else None + weight = ( + torch.full((width,), 0.75, dtype=torch.float16, device=device) + if has_weight + else None + ) expected_input = input.clone() expected_residual = residual.clone() _torch_fused_add_rms_norm(expected_input, expected_residual, weight, 1e-6) infini.ops.fused_add_rms_norm( - input, residual, weight, 1e-6, implementation_index=implementation_index, + input, + residual, + weight, + 1e-6, + implementation_index=implementation_index, stream=get_stream(input.device), ) - torch.testing.assert_close(input.cpu(), expected_input.cpu(), rtol=1e-2, atol=1e-2) - torch.testing.assert_close(residual.cpu(), expected_residual.cpu(), rtol=0, atol=0) + torch.testing.assert_close( + input.cpu(), expected_input.cpu(), rtol=1e-2, atol=1e-2 + ) + torch.testing.assert_close( + residual.cpu(), expected_residual.cpu(), rtol=0, atol=0 + ) def _torch_fused_add_rms_norm(input, residual, weight, epsilon): diff --git a/tests/test_reshape_and_cache_flash.py b/tests/test_reshape_and_cache_flash.py index 1001bf7a4..744a97138 100644 --- a/tests/test_reshape_and_cache_flash.py +++ b/tests/test_reshape_and_cache_flash.py @@ -155,12 +155,16 @@ def invoke(): @pytest.mark.parametrize("cache_layout", ("NHD", "HND")) -def test_reshape_and_cache_flash_repeated_slots(device, implementation_index, cache_layout): +def test_reshape_and_cache_flash_repeated_slots( + device, implementation_index, cache_layout +): if device != "npu" or implementation_index != 0: pytest.skip("Ascend native cache update regression") key_cache = _make_cache(3, 4, 2, 16, cache_layout, torch.float16, device).fill_(-8) - value_cache = _make_cache(3, 4, 2, 16, cache_layout, torch.float16, device).fill_(-9) + value_cache = _make_cache(3, 4, 2, 16, cache_layout, torch.float16, device).fill_( + -9 + ) expected_key = key_cache.cpu().clone() expected_value = value_cache.cpu().clone() scale = torch.ones((), dtype=torch.float32, device=device) @@ -175,8 +179,16 @@ def test_reshape_and_cache_flash_repeated_slots(device, implementation_index, ca expected_key[block, page] = key_cpu[token] expected_value[block, page] = value_cpu[token] infini.ops.reshape_and_cache_flash( - key, value, mapping, scale, scale, "auto", key_cache, value_cache, - implementation_index=implementation_index, stream=get_stream(key.device), + key, + value, + mapping, + scale, + scale, + "auto", + key_cache, + value_cache, + implementation_index=implementation_index, + stream=get_stream(key.device), ) torch.testing.assert_close(key_cache.cpu(), expected_key, rtol=0, atol=0) torch.testing.assert_close(value_cache.cpu(), expected_value, rtol=0, atol=0) diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 8e746e6d0..cb4bcb325 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -126,9 +126,7 @@ def test_rms_norm_qkv_slice( ) out = backing[:, 128 : 128 + width].view(seq_len, heads, 128) else: - out = torch.empty( - (seq_len, heads, 128), dtype=torch.float16, device=device - ) + out = torch.empty((seq_len, heads, 128), dtype=torch.float16, device=device) if transpose_rows: input = input.transpose(0, 1) @@ -166,9 +164,7 @@ def test_rms_norm_strided_npu_stream(device, implementation_index): source_cpu = torch.randn(36, 4096, dtype=torch.float16) weight_cpu = torch.randn(128, dtype=torch.float16) - expected = _torch_rms_norm( - source_cpu[:, 2048:3072].reshape(36, 8, 128), weight_cpu - ) + expected = _torch_rms_norm(source_cpu[:, 2048:3072].reshape(36, 8, 128), weight_cpu) stream = torch.npu.Stream() with torch.npu.stream(stream): diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index 1cfbf39ef..4d360a4b5 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -263,8 +263,7 @@ def test_rotary_embedding_ascend_strided_fused_qkv_view( @pytest.mark.parametrize("num_tokens", (1, 2, 24)) @pytest.mark.parametrize( "head_size, num_heads", - ((64, 8), (64, 9), (64, 36), (64, 40), - (128, 8), (128, 9), (128, 16), (128, 32)), + ((64, 8), (64, 9), (64, 36), (64, 40), (128, 8), (128, 9), (128, 16), (128, 32)), ) @pytest.mark.parametrize("has_key", (False, True)) @pytest.mark.parametrize("strided", (False, True)) @@ -276,9 +275,10 @@ def test_rotary_embedding_ascend_neox_decode( num_kv_heads = num_heads // 2 if num_heads % 2 == 0 else 1 half_dim = head_size // 2 - frequencies = torch.arange(256).float()[:, None] * ( - 1000000 ** (-torch.arange(half_dim).float() / half_dim) - )[None, :] + frequencies = ( + torch.arange(256).float()[:, None] + * (1000000 ** (-torch.arange(half_dim).float() / half_dim))[None, :] + ) cache_cpu = torch.cat((frequencies.cos(), frequencies.sin()), -1).half() cache = cache_cpu.to(device) # Reuse the cached operator with new positions and Q/K allocations. @@ -296,8 +296,10 @@ def to_device(value, heads): if not strided: return value.to(device) parent = torch.full( - (num_tokens, heads, 2 * head_size), -17, - dtype=torch.float16, device=device, + (num_tokens, heads, 2 * head_size), + -17, + dtype=torch.float16, + device=device, ) view = parent[..., :head_size] view.copy_(value.view(num_tokens, heads, head_size)) @@ -307,11 +309,22 @@ def to_device(value, heads): query = to_device(query_cpu, num_heads) key = to_device(key_cpu, num_kv_heads) if has_key else None _torch_rotary_embedding( - positions_cpu, query_cpu, key_cpu, cache_cpu, - head_size, True, 0, False, + positions_cpu, + query_cpu, + key_cpu, + cache_cpu, + head_size, + True, + 0, + False, ) infini.ops.rotary_embedding( - positions_cpu.to(device), query, key, cache, head_size, True, + positions_cpu.to(device), + query, + key, + cache, + head_size, + True, implementation_index=implementation_index, stream=get_stream(query.device), ) diff --git a/tests/test_silu_and_mul.py b/tests/test_silu_and_mul.py index cd1eb95f0..d99a88bd3 100644 --- a/tests/test_silu_and_mul.py +++ b/tests/test_silu_and_mul.py @@ -101,12 +101,16 @@ def test_silu_and_mul_repeated_small_shape(device, implementation_index, padded) for value in (0.25, 0.75): input = torch.empty_strided( - (13, 8), (12, 1) if padded else (8, 1), - dtype=torch.float16, device=device, + (13, 8), + (12, 1) if padded else (8, 1), + dtype=torch.float16, + device=device, ).fill_(value) out = torch.empty_strided( - (13, 4), (6, 1) if padded else (4, 1), - dtype=torch.float16, device=device, + (13, 4), + (6, 1) if padded else (4, 1), + dtype=torch.float16, + device=device, ).fill_(-9) expected = _torch_silu_and_mul(input, torch.empty_like(out)).cpu() actual = _silu_and_mul(input, out, implementation_index).cpu()