From 33e10fb2f291e316d8218e31af39d343905033dd Mon Sep 17 00:00:00 2001 From: zhangxin <77204052+JoeZhang-0x000@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:37:27 +0800 Subject: [PATCH] perf(ascend): parallelize Q/K, SwiGLU and strided KV kernels --- scripts/bench_ascend_qk.py | 135 +++++++++++++++ src/native/ascend/custom/CMakeLists.txt | 3 + .../op_kernel/reshape_and_cache_flash.cpp | 93 +++++----- .../rms_norm_qk/op_kernel/rms_norm_qk.cpp | 152 +++++++++++++++++ .../op_kernel/rotary_embedding_qk.cpp | 160 ++++++++++++++++++ .../silu_and_mul/op_kernel/silu_and_mul.cpp | 87 ++++++++++ .../ops/reshape_and_cache_flash/kernel.h | 9 +- src/native/ascend/ops/rms_norm/kernel.h | 37 ++++ .../ascend/ops/rotary_embedding/kernel.h | 54 ++++++ src/native/ascend/ops/silu_and_mul/kernel.h | 41 +++++ tests/test_reshape_and_cache_flash.py | 53 ++++++ tests/test_rms_norm.py | 54 ++++++ tests/test_rotary_embedding.py | 92 ++++++++++ tests/test_silu_and_mul.py | 34 ++++ 14 files changed, 959 insertions(+), 45 deletions(-) create mode 100644 scripts/bench_ascend_qk.py create mode 100644 src/native/ascend/custom/rms_norm_qk/op_kernel/rms_norm_qk.cpp create mode 100644 src/native/ascend/custom/rotary_embedding_qk/op_kernel/rotary_embedding_qk.cpp create mode 100644 src/native/ascend/custom/silu_and_mul/op_kernel/silu_and_mul.cpp diff --git a/scripts/bench_ascend_qk.py b/scripts/bench_ascend_qk.py new file mode 100644 index 000000000..9db915efc --- /dev/null +++ b/scripts/bench_ascend_qk.py @@ -0,0 +1,135 @@ +"""Compare warm Ascend operator latency using the selected InfiniOps package.""" + +import argparse +import json +import statistics +import time +from pathlib import Path + +import infini.ops +import torch +import torch_npu # noqa: F401 + + +def measure(invoke, iterations, repeats): + for _ in range(5): + invoke() + + torch.npu.synchronize() + samples = [] + + for _ in range(repeats): + start = time.perf_counter() + + for _ in range(iterations): + invoke() + + torch.npu.synchronize() + samples.append((time.perf_counter() - start) * 1e6 / iterations) + + return {"median_us": statistics.median(samples), "samples_us": samples} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--case-prefix", default="") + args = parser.parse_args() + torch.npu.set_device(0) + torch.manual_seed(41) + torch.set_num_threads(2) + stream = torch.npu.current_stream().npu_stream + options = {"implementation_index": 0, "stream": stream} + result = { + "package": infini.ops.__file__, + "iterations": args.iterations, + "repeats": args.repeats, + "measurement": "warm host plus device latency; median of synchronized batches", + "cases": {}, + } + + def record(name, invoke): + if not name.startswith(args.case_prefix): + return + + result["cases"][name] = measure(invoke, args.iterations, args.repeats) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(name, result["cases"][name]["median_us"], flush=True) + + input = torch.randn((64, 35, 128), dtype=torch.bfloat16, device="npu")[:, :32] + weight = torch.ones(128, dtype=input.dtype, device=input.device) + out = torch.empty((64, 32, 128), dtype=input.dtype, device=input.device) + record( + "rms_norm_bf16_64x32x128_packed", + lambda: infini.ops.rms_norm(input, weight, 1e-6, out, **options), + ) + + qkv = torch.randn((64, 48, 128), dtype=torch.bfloat16, device="npu") + query, key = qkv[:, :32], qkv[:, 32:40] + positions = torch.arange(64, dtype=torch.int64, device="npu") + cache = torch.cat( + (torch.ones((128, 64), device="npu"), torch.zeros((128, 64), device="npu")), + dim=1, + ).bfloat16() + record( + "rotary_embedding_bf16_64x32_8x128_packed", + lambda: infini.ops.rotary_embedding( + positions, query, key, cache, 128, True, **options + ), + ) + + for hidden, padded in ((3072, False), (6144, False), (3072, True), (6144, True)): + input = torch.randn((64, hidden * 2), dtype=torch.bfloat16, device="npu") + out = torch.empty_strided( + (64, hidden), + (hidden + 128 if padded else hidden, 1), + dtype=input.dtype, + device=input.device, + ) + record( + f"silu_and_mul_bf16_64x{hidden}{'_padded' if padded else ''}", + lambda: infini.ops.silu_and_mul(input, out, **options), + ) + + for layout in ("NHD", "HND"): + tokens, heads, dim, block_size = 8192, 8, 128, 128 + packed = torch.randn( + (tokens, heads * 3, dim), dtype=torch.bfloat16, device="npu" + ) + key, value = packed[:, :heads], packed[:, heads : heads * 2] + shape = ( + (tokens // block_size, block_size, heads, dim) + if layout == "NHD" + else (tokens // block_size, heads, block_size, dim) + ) + key_cache = torch.zeros(shape, dtype=key.dtype, device=key.device) + value_cache = torch.zeros_like(key_cache) + + if layout == "HND": + key_cache, value_cache = ( + key_cache.transpose(1, 2), + value_cache.transpose(1, 2), + ) + + slots = torch.arange(tokens, dtype=torch.int64, device="npu") + scale = torch.ones((), dtype=torch.float32, device="npu") + record( + f"reshape_and_cache_bf16_8192x8x128_{layout}", + lambda: infini.ops.reshape_and_cache_flash( + key, + value, + slots, + scale, + scale, + "auto", + key_cache, + value_cache, + **options, + ), + ) + + +if __name__ == "__main__": + main() diff --git a/src/native/ascend/custom/CMakeLists.txt b/src/native/ascend/custom/CMakeLists.txt index 94f19850a..cb21426d6 100644 --- a/src/native/ascend/custom/CMakeLists.txt +++ b/src/native/ascend/custom/CMakeLists.txt @@ -63,6 +63,9 @@ ascendc_library(no_workspace_kernel STATIC rms_norm/op_kernel/rms_norm.cpp add_rms_norm/op_kernel/add_rms_norm.cpp reshape_and_cache_flash/op_kernel/reshape_and_cache_flash.cpp + rms_norm_qk/op_kernel/rms_norm_qk.cpp + rotary_embedding_qk/op_kernel/rotary_embedding_qk.cpp + silu_and_mul/op_kernel/silu_and_mul.cpp ) # `CANN 8.5.1` installs host objects under `objects-/`, while 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..21d5a4d88 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 @@ -38,62 +38,71 @@ class KernelReshapeAndCacheFlashStrided { value_cache_page_stride_ = value_cache_page_stride; value_cache_head_stride_ = value_cache_head_stride; - int64_t aligned_bytes = (head_bytes_ + 31) / 32 * 32; - pipe_.InitBuffer(copy_buffer_, static_cast(aligned_bytes)); + combine_heads_ = key_head_stride == head_size && + value_head_stride == head_size && + key_cache_head_stride == head_size && + value_cache_head_stride == head_size; + pipe_.InitBuffer(copy_buffer_, kTileBytes); } __aicore__ inline void Process() { auto local = copy_buffer_.Get(); - AscendC::DataCopyExtParams copy_params{ - 1, static_cast(head_bytes_), 0, 0, 0}; - AscendC::DataCopyPadExtParams pad_params{false, 0, 0, 0}; - - for (int64_t token = 0; token < num_tokens_; ++token) { + int64_t heads_per_task = combine_heads_ ? num_heads_ : 1; + int64_t tasks_per_token = combine_heads_ ? 1 : num_heads_; + for (int64_t task = AscendC::GetBlockIdx(); + task < num_tokens_ * tasks_per_token; task += AscendC::GetBlockNum()) { + int64_t token = task / tasks_per_token; + int64_t head = task % tasks_per_token; int64_t slot = slot_mapping_.GetValue(token); - if (slot < 0 || slot >= num_blocks_ * block_size_) { - continue; - } + if (slot < 0 || slot >= num_blocks_ * block_size_) continue; int64_t block = slot / block_size_; int64_t page = slot % block_size_; - - for (int64_t head = 0; head < num_heads_; ++head) { - int64_t key_offset = - (token * key_token_stride_ + head * key_head_stride_) * - element_size_; - int64_t key_cache_offset = - (block * key_cache_block_stride_ + page * key_cache_page_stride_ + - head * key_cache_head_stride_) * - element_size_; - Copy(key_, key_offset, key_cache_, key_cache_offset, local, copy_params, - pad_params); - - int64_t value_offset = - (token * value_token_stride_ + head * value_head_stride_) * - element_size_; - int64_t value_cache_offset = (block * value_cache_block_stride_ + - page * value_cache_page_stride_ + - head * value_cache_head_stride_) * - element_size_; - Copy(value_, value_offset, value_cache_, value_cache_offset, local, - copy_params, pad_params); - } + int64_t key_offset = + (token * key_token_stride_ + head * key_head_stride_) * element_size_; + int64_t key_cache_offset = + (block * key_cache_block_stride_ + page * key_cache_page_stride_ + + head * key_cache_head_stride_) * + element_size_; + Copy(key_, key_offset, key_cache_, key_cache_offset, + heads_per_task * head_bytes_, local); + int64_t value_offset = + (token * value_token_stride_ + head * value_head_stride_) * + element_size_; + int64_t value_cache_offset = + (block * value_cache_block_stride_ + page * value_cache_page_stride_ + + head * value_cache_head_stride_) * + element_size_; + Copy(value_, value_offset, value_cache_, value_cache_offset, + heads_per_task * head_bytes_, local); } } private: - __aicore__ inline void Copy( - AscendC::GlobalTensor& source, int64_t source_offset, - AscendC::GlobalTensor& destination, int64_t destination_offset, - AscendC::LocalTensor& local, - const AscendC::DataCopyExtParams& copy_params, - const AscendC::DataCopyPadExtParams& pad_params) { - AscendC::DataCopyPad(local, source[source_offset], copy_params, pad_params); - AscendC::PipeBarrier(); - AscendC::DataCopyPad(destination[destination_offset], local, copy_params); - AscendC::PipeBarrier(); + __aicore__ inline void Copy(AscendC::GlobalTensor& source, + int64_t source_offset, + AscendC::GlobalTensor& destination, + int64_t destination_offset, int64_t bytes, + AscendC::LocalTensor& local) { + AscendC::DataCopyPadExtParams pad{false, 0, 0, 0}; + for (int64_t offset = 0; offset < bytes; offset += kTileBytes) { + auto count = static_cast( + bytes - offset < kTileBytes ? bytes - offset : kTileBytes); + AscendC::DataCopyExtParams params{1, count, 0, 0, 0}; + AscendC::DataCopyPad(local, source[source_offset + offset], params, pad); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + AscendC::DataCopyPad(destination[destination_offset + offset], local, + params); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + } } + static constexpr int64_t kTileBytes = 2048; + + bool combine_heads_{false}; + AscendC::TPipe pipe_; AscendC::TBuf copy_buffer_; AscendC::GlobalTensor key_; diff --git a/src/native/ascend/custom/rms_norm_qk/op_kernel/rms_norm_qk.cpp b/src/native/ascend/custom/rms_norm_qk/op_kernel/rms_norm_qk.cpp new file mode 100644 index 000000000..afef03e8e --- /dev/null +++ b/src/native/ascend/custom/rms_norm_qk/op_kernel/rms_norm_qk.cpp @@ -0,0 +1,152 @@ +#include + +#include "kernel_operator.h" + +// One token's contiguous 128-element heads are normalized together. +// Reduction and broadcast stay in UB; packed QKV token strides are preserved. +class KernelRmsNormQk { + public: + __aicore__ inline void Init(GM_ADDR y, GM_ADDR x, GM_ADDR weight, + size_t heads, ptrdiff_t y_token, + ptrdiff_t x_token, float epsilon) { + heads_ = heads; + y_token_ = y_token; + x_token_ = x_token; + epsilon_ = epsilon; + x_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(x)); + y_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(y)); + weight_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(weight)); + const size_t elements = heads * 128; + pipe_.InitBuffer(input_, 2, elements * sizeof(bfloat16_t)); + pipe_.InitBuffer(output_, 2, elements * sizeof(bfloat16_t)); + pipe_.InitBuffer(weight_input_, 1, 128 * sizeof(bfloat16_t)); + pipe_.InitBuffer(x_f_, elements * sizeof(float)); + pipe_.InitBuffer(square_f_, elements * sizeof(float)); + pipe_.InitBuffer(out_f_, elements * sizeof(float)); + pipe_.InitBuffer(weight_f_, 128 * sizeof(float)); + pipe_.InitBuffer(rstd_, 32 * sizeof(float)); + pipe_.InitBuffer(variance_, 32 * sizeof(float)); + pipe_.InitBuffer(ones_, 32 * sizeof(float)); + pipe_.InitBuffer(rstd_blocks_, 32 * 8 * sizeof(float)); + } + + __aicore__ inline void Process(size_t tokens) { + AscendC::LocalTensor weight = + weight_input_.AllocTensor(); + AscendC::DataCopy(weight, weight_, 128); + weight_input_.EnQue(weight); + weight = weight_input_.DeQue(); + AscendC::LocalTensor wf = weight_f_.Get(); + AscendC::Cast(wf, weight, AscendC::RoundMode::CAST_NONE, 128); + weight_input_.FreeTensor(weight); + AscendC::PipeBarrier(); + AscendC::LocalTensor ones = ones_.Get(); + AscendC::Duplicate(ones, 1.0f, 32); + AscendC::PipeBarrier(); + for (size_t token = AscendC::GetBlockIdx(); token < tokens; + token += AscendC::GetBlockNum()) { + AscendC::LocalTensor input = input_.AllocTensor(); + AscendC::DataCopy(input, x_[token * x_token_], heads_ * 128); + input_.EnQue(input); + input = input_.DeQue(); + AscendC::LocalTensor output = + output_.AllocTensor(); + AscendC::LocalTensor x = x_f_.Get(); + AscendC::LocalTensor squares = square_f_.Get(); + AscendC::LocalTensor out = out_f_.Get(); + AscendC::LocalTensor rstd = rstd_.Get(); + AscendC::LocalTensor variance = variance_.Get(); + AscendC::LocalTensor factors = rstd_blocks_.Get(); + AscendC::Cast(x, input, AscendC::RoundMode::CAST_NONE, heads_ * 128); + AscendC::PipeBarrier(); + AscendC::Mul(squares, x, x, heads_ * 128); + AscendC::PipeBarrier(); + // Sum corresponding elements of the two halves, then reduce 64. + AscendC::BinaryRepeatParams halves{1, 1, 1, 16, 16, 16}; + AscendC::Add(out, squares, squares[64], uint64_t(64), + static_cast(heads_), halves); + AscendC::PipeBarrier(); + AscendC::WholeReduceSum(rstd, out, 64, static_cast(heads_), 1, 1, + 16); + AscendC::PipeBarrier(); + AscendC::Muls(variance, rstd, 1.0f / 128.0f, heads_); + AscendC::PipeBarrier(); + AscendC::Adds(variance, variance, epsilon_, heads_); + AscendC::PipeBarrier(); + // Match the native sqrt/divide sequence before BF16 rounding. + AscendC::Sqrt(rstd, variance, heads_); + AscendC::PipeBarrier(); + AscendC::Div(rstd, ones, rstd, heads_); + AscendC::PipeBarrier(); + AscendC::Brcb(factors, rstd, static_cast((heads_ + 7) / 8), + {1, 8}); + AscendC::PipeBarrier(); + AscendC::BinaryRepeatParams normalize{1, 1, 0, 16, 16, 1}; + AscendC::Mul(out, x, factors, uint64_t(64), static_cast(heads_), + normalize); + AscendC::Mul(out[64], x[64], factors, uint64_t(64), + static_cast(heads_), normalize); + AscendC::PipeBarrier(); + AscendC::BinaryRepeatParams broadcast_weight{1, 1, 1, 16, 16, 0}; + AscendC::Mul(out, out, wf, uint64_t(64), static_cast(heads_), + broadcast_weight); + AscendC::Mul(out[64], out[64], wf[64], uint64_t(64), + static_cast(heads_), broadcast_weight); + AscendC::PipeBarrier(); + AscendC::Cast(output, out, AscendC::RoundMode::CAST_RINT, heads_ * 128); + output_.EnQue(output); + input_.FreeTensor(input); + output = output_.DeQue(); + AscendC::DataCopy(y_[token * y_token_], output, heads_ * 128); + output_.FreeTensor(output); + } + } + + private: + AscendC::TPipe pipe_; + + AscendC::TQue input_; + + AscendC::TQue weight_input_; + + AscendC::TQue output_; + + AscendC::TBuf x_f_; + + AscendC::TBuf square_f_; + + AscendC::TBuf out_f_; + + AscendC::TBuf weight_f_; + + AscendC::TBuf rstd_; + + AscendC::TBuf variance_; + + AscendC::TBuf ones_; + + AscendC::TBuf rstd_blocks_; + + AscendC::GlobalTensor x_; + + AscendC::GlobalTensor y_; + + AscendC::GlobalTensor weight_; + + size_t heads_; + + ptrdiff_t x_token_; + + ptrdiff_t y_token_; + + float epsilon_; +}; + +extern "C" __global__ __aicore__ void rms_norm_qk( + GM_ADDR input, GM_ADDR weight, GM_ADDR out, int64_t tokens, int64_t heads, + int64_t input_token_stride, int64_t out_token_stride, float epsilon) { + KernelRmsNormQk op; + op.Init(out, input, weight, heads, out_token_stride, input_token_stride, + epsilon); + op.Process(tokens); +} diff --git a/src/native/ascend/custom/rotary_embedding_qk/op_kernel/rotary_embedding_qk.cpp b/src/native/ascend/custom/rotary_embedding_qk/op_kernel/rotary_embedding_qk.cpp new file mode 100644 index 000000000..2fd377219 --- /dev/null +++ b/src/native/ascend/custom/rotary_embedding_qk/op_kernel/rotary_embedding_qk.cpp @@ -0,0 +1,160 @@ +#include + +#include "kernel_operator.h" + +// A token's contiguous heads share one DMA transfer and one sin/cos load. +// FP32 vector repeats broadcast the half-head tables across all heads. +class KernelRotaryEmbeddingQk { + public: + __aicore__ inline void Init(GM_ADDR positions, GM_ADDR data, GM_ADDR cache, + size_t heads, ptrdiff_t token_stride, + ptrdiff_t cache_stride) { + heads_ = heads; + token_stride_ = token_stride; + cache_stride_ = cache_stride; + x_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(data)); + y_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(data)); + sin_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(cache) + 64); + cos_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(cache)); + positions_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(positions)); + const size_t elements = heads_ * 128; + pipe_.InitBuffer(input_, 2, elements * sizeof(bfloat16_t)); + pipe_.InitBuffer(output_, 2, elements * sizeof(bfloat16_t)); + pipe_.InitBuffer(sin_input_, 2, 64 * sizeof(bfloat16_t)); + pipe_.InitBuffer(cos_input_, 2, 64 * sizeof(bfloat16_t)); + pipe_.InitBuffer(input_f_, elements * sizeof(float)); + pipe_.InitBuffer(output_f_, elements * sizeof(float)); + pipe_.InitBuffer(sin_f_, 64 * sizeof(float)); + pipe_.InitBuffer(cos_f_, 64 * sizeof(float)); + pipe_.InitBuffer(first_, heads_ * 64 * sizeof(float)); + pipe_.InitBuffer(second_, heads_ * 64 * sizeof(float)); + pipe_.InitBuffer(product1_, heads_ * 64 * sizeof(float)); + pipe_.InitBuffer(product2_, heads_ * 64 * sizeof(float)); + pipe_.InitBuffer(result1_, heads_ * 64 * sizeof(float)); + pipe_.InitBuffer(result2_, heads_ * 64 * sizeof(float)); + } + + __aicore__ inline void Process(size_t tokens) { + for (size_t work = AscendC::GetBlockIdx(); work < tokens; + work += AscendC::GetBlockNum()) { + const size_t token = work; + const size_t position = static_cast(positions_.GetValue(token)); + AscendC::LocalTensor input = input_.AllocTensor(); + AscendC::LocalTensor sin = + sin_input_.AllocTensor(); + AscendC::LocalTensor cos = + cos_input_.AllocTensor(); + AscendC::DataCopy(input, x_[token * token_stride_], heads_ * 128); + AscendC::DataCopy(sin, sin_[position * cache_stride_], 64); + AscendC::DataCopy(cos, cos_[position * cache_stride_], 64); + input_.EnQue(input); + sin_input_.EnQue(sin); + cos_input_.EnQue(cos); + input = input_.DeQue(); + sin = sin_input_.DeQue(); + cos = cos_input_.DeQue(); + AscendC::LocalTensor output = + output_.AllocTensor(); + AscendC::LocalTensor x = input_f_.Get(); + AscendC::LocalTensor out = output_f_.Get(); + AscendC::LocalTensor sf = sin_f_.Get(); + AscendC::LocalTensor cf = cos_f_.Get(); + AscendC::LocalTensor first = first_.Get(); + AscendC::LocalTensor second = second_.Get(); + AscendC::LocalTensor p1 = product1_.Get(); + AscendC::LocalTensor p2 = product2_.Get(); + AscendC::LocalTensor r1 = result1_.Get(); + AscendC::LocalTensor r2 = result2_.Get(); + AscendC::Cast(x, input, AscendC::RoundMode::CAST_NONE, heads_ * 128); + AscendC::Cast(sf, sin, AscendC::RoundMode::CAST_NONE, 64); + AscendC::Cast(cf, cos, AscendC::RoundMode::CAST_NONE, 64); + AscendC::PipeBarrier(); + AscendC::DataCopyParams split{static_cast(heads_), 8, 8, 0}; + AscendC::DataCopy(first, x, split); + AscendC::DataCopy(second, x[64], split); + AscendC::PipeBarrier(); + AscendC::BinaryRepeatParams broadcast{1, 1, 1, 8, 8, 0}; + AscendC::Mul(p1, first, cf, uint64_t(64), static_cast(heads_), + broadcast); + AscendC::Mul(p2, second, sf, uint64_t(64), static_cast(heads_), + broadcast); + AscendC::PipeBarrier(); + AscendC::Sub(r1, p1, p2, heads_ * 64); + AscendC::PipeBarrier(); + AscendC::Mul(p1, first, sf, uint64_t(64), static_cast(heads_), + broadcast); + AscendC::Mul(p2, second, cf, uint64_t(64), static_cast(heads_), + broadcast); + AscendC::PipeBarrier(); + AscendC::Add(r2, p1, p2, heads_ * 64); + AscendC::PipeBarrier(); + AscendC::DataCopyParams merge{static_cast(heads_), 8, 0, 8}; + AscendC::DataCopy(out, r1, merge); + AscendC::DataCopy(out[64], r2, merge); + AscendC::PipeBarrier(); + AscendC::Cast(output, out, AscendC::RoundMode::CAST_RINT, heads_ * 128); + output_.EnQue(output); + input_.FreeTensor(input); + sin_input_.FreeTensor(sin); + cos_input_.FreeTensor(cos); + output = output_.DeQue(); + AscendC::DataCopy(y_[token * token_stride_], output, heads_ * 128); + output_.FreeTensor(output); + } + } + + private: + AscendC::TPipe pipe_; + + AscendC::TQue input_; + + AscendC::TQue sin_input_; + + AscendC::TQue cos_input_; + + AscendC::TQue output_; + + AscendC::TBuf input_f_; + + AscendC::TBuf output_f_; + + AscendC::TBuf sin_f_; + + AscendC::TBuf cos_f_; + + AscendC::TBuf first_; + + AscendC::TBuf second_; + + AscendC::TBuf product1_; + + AscendC::TBuf product2_; + + AscendC::TBuf result1_; + + AscendC::TBuf result2_; + + AscendC::GlobalTensor x_; + + AscendC::GlobalTensor y_; + + AscendC::GlobalTensor sin_; + + AscendC::GlobalTensor cos_; + + AscendC::GlobalTensor positions_; + + size_t heads_{0}; + + ptrdiff_t token_stride_{0}; + + ptrdiff_t cache_stride_{0}; +}; + +extern "C" __global__ __aicore__ void rotary_embedding_qk( + GM_ADDR positions, GM_ADDR data, GM_ADDR cos_sin_cache, int64_t tokens, + int64_t heads, int64_t token_stride, int64_t cache_stride) { + KernelRotaryEmbeddingQk op; + op.Init(positions, data, cos_sin_cache, heads, token_stride, cache_stride); + op.Process(tokens); +} diff --git a/src/native/ascend/custom/silu_and_mul/op_kernel/silu_and_mul.cpp b/src/native/ascend/custom/silu_and_mul/op_kernel/silu_and_mul.cpp new file mode 100644 index 000000000..bb217b9c8 --- /dev/null +++ b/src/native/ascend/custom/silu_and_mul/op_kernel/silu_and_mul.cpp @@ -0,0 +1,87 @@ +#include "kernel_operator.h" + +class KernelSiluAndMul { + public: + __aicore__ inline void Init(GM_ADDR input, GM_ADDR out, int64_t hidden, + int64_t input_stride, int64_t out_stride) { + hidden_ = hidden; + input_stride_ = input_stride; + out_stride_ = out_stride; + tile_ = hidden / 8; + input_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(input)); + out_.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t*>(out)); + pipe_.InitBuffer(gate_queue_, 2, tile_ * sizeof(bfloat16_t)); + pipe_.InitBuffer(up_queue_, 2, tile_ * sizeof(bfloat16_t)); + pipe_.InitBuffer(out_queue_, 2, tile_ * sizeof(bfloat16_t)); + pipe_.InitBuffer(gate_float_, tile_ * sizeof(float)); + pipe_.InitBuffer(up_float_, tile_ * sizeof(float)); + pipe_.InitBuffer(out_float_, tile_ * sizeof(float)); + } + + __aicore__ inline void Process(int64_t rows) { + int64_t tile_offset = (AscendC::GetBlockIdx() % 8) * tile_; + for (int64_t row = AscendC::GetBlockIdx() / 8; row < rows; + row += AscendC::GetBlockNum() / 8) { + auto gate = gate_queue_.AllocTensor(); + auto up = up_queue_.AllocTensor(); + AscendC::DataCopy(gate, input_[row * input_stride_ + tile_offset], tile_); + AscendC::DataCopy(up, input_[row * input_stride_ + hidden_ + tile_offset], + tile_); + gate_queue_.EnQue(gate); + up_queue_.EnQue(up); + gate = gate_queue_.DeQue(); + up = up_queue_.DeQue(); + auto out = out_queue_.AllocTensor(); + auto gate_float = gate_float_.Get(); + auto up_float = up_float_.Get(); + auto out_float = out_float_.Get(); + AscendC::Cast(gate_float, gate, AscendC::RoundMode::CAST_NONE, tile_); + AscendC::Cast(up_float, up, AscendC::RoundMode::CAST_NONE, tile_); + // AscendC takes the multiplier before the SiLU gate. + AscendC::SwiGLU(out_float, up_float, gate_float, 1.0f, + tile_); + AscendC::Cast(out, out_float, AscendC::RoundMode::CAST_RINT, tile_); + out_queue_.EnQue(out); + gate_queue_.FreeTensor(gate); + up_queue_.FreeTensor(up); + out = out_queue_.DeQue(); + AscendC::DataCopy(out_[row * out_stride_ + tile_offset], out, tile_); + out_queue_.FreeTensor(out); + } + } + + private: + AscendC::TPipe pipe_; + + AscendC::TQue gate_queue_; + + AscendC::TQue up_queue_; + + AscendC::TQue out_queue_; + + AscendC::TBuf gate_float_; + + AscendC::TBuf up_float_; + + AscendC::TBuf out_float_; + + AscendC::GlobalTensor input_; + + AscendC::GlobalTensor out_; + + int64_t hidden_{0}; + + int64_t tile_{0}; + + int64_t input_stride_{0}; + + int64_t out_stride_{0}; +}; + +extern "C" __global__ __aicore__ void silu_and_mul_tiled( + GM_ADDR input, GM_ADDR out, int64_t rows, int64_t hidden, + int64_t input_stride, int64_t out_stride) { + KernelSiluAndMul op; + op.Init(input, out, hidden, input_stride, out_stride); + op.Process(rows); +} 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..a4f8717bb 100644 --- a/src/native/ascend/ops/reshape_and_cache_flash/kernel.h +++ b/src/native/ascend/ops/reshape_and_cache_flash/kernel.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_ASCEND_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ #define INFINI_OPS_ASCEND_RESHAPE_AND_CACHE_FLASH_KERNEL_H_ +#include #include #include #include @@ -121,9 +122,11 @@ class Operator const Tensor slot_mapping, Tensor key_cache, Tensor value_cache, aclrtStream stream) const { auto ret = aclrtlaunch_reshape_and_cache_flash_strided( - 1, stream, const_cast(key.data()), - const_cast(value.data()), const_cast(slot_mapping.data()), - key_cache.data(), value_cache.data(), static_cast(num_tokens_), + static_cast( + std::min(num_tokens_ * num_heads_, 64)), + stream, const_cast(key.data()), const_cast(value.data()), + const_cast(slot_mapping.data()), key_cache.data(), + value_cache.data(), static_cast(num_tokens_), static_cast(num_heads_), static_cast(head_size_), static_cast(block_size_), static_cast(num_blocks_), static_cast(key_token_stride_), diff --git a/src/native/ascend/ops/rms_norm/kernel.h b/src/native/ascend/ops/rms_norm/kernel.h index 734d8ed42..9a01da8c9 100644 --- a/src/native/ascend/ops/rms_norm/kernel.h +++ b/src/native/ascend/ops/rms_norm/kernel.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_ASCEND_RMS_NORM_KERNEL_H_ #define INFINI_OPS_ASCEND_RMS_NORM_KERNEL_H_ +#include #include #include @@ -13,6 +14,10 @@ #include "native/ascend/workspace_pool_.h" #include "operator.h" +#ifdef INFINI_HAS_CUSTOM_KERNELS +#include "aclrtlaunch_rms_norm_qk.h" +#endif + namespace infini::ops { template <> @@ -24,6 +29,19 @@ class Operator : public RmsNorm { assert(input.strides().back() == 1 && out.strides().back() == 1 && "`RmsNorm` Ascend path requires a contiguous last dimension"); +#ifdef INFINI_HAS_CUSTOM_KERNELS + use_qk_kernel_ = + input.dtype() == DataType::kBFloat16 && + weight.dtype() == input.dtype() && dim_ == 128 && + (ndim_ == 2 || ndim_ == 3) && nhead_ > 0 && nhead_ <= 32 && + input.shape() == out.shape() && weight.ndim() == 1 && + weight.numel() == 128 && weight.stride(0) == 1 && + (ndim_ == 2 || (input.stride(1) == 128 && out.stride(1) == 128)) && + input.stride(0) >= static_cast(nhead_ * 128) && + out.stride(0) >= static_cast(nhead_ * 128); + if (use_qk_kernel_) return; +#endif + needs_input_staging_ = !input.IsContiguous(); needs_out_staging_ = !out.IsContiguous(); logical_rows_ = input.numel() / dim_; @@ -109,6 +127,23 @@ class Operator : public RmsNorm { void operator()(const Tensor input, const Tensor weight, float eps, Tensor out) const override { auto stream = static_cast(stream_); +#ifdef INFINI_HAS_CUSTOM_KERNELS + if (use_qk_kernel_) { + if (input.size(0) == 0) return; + + auto blocks = + static_cast(std::min(input.size(0), 64)); + auto ret = aclrtlaunch_rms_norm_qk( + blocks, stream, const_cast(input.data()), + const_cast(weight.data()), out.data(), + static_cast(input.size(0)), static_cast(nhead_), + static_cast(input.stride(0)), + static_cast(out.stride(0)), eps); + assert(ret == ACL_SUCCESS && "Ascend Q/K `RmsNorm` execution failed"); + return; + } +#endif + void* input_data = const_cast(input.data()); void* out_data = out.data(); @@ -202,6 +237,8 @@ class Operator : public RmsNorm { mutable ascend::AclTensorCache weight_cache_; mutable ascend::AclTensorCache norm_out_cache_; mutable ascend::AclTensorCache weight_fp32_cache_; + bool use_qk_kernel_{false}; + bool needs_input_staging_{false}; bool needs_out_staging_{false}; bool needs_weight_cast_{false}; diff --git a/src/native/ascend/ops/rotary_embedding/kernel.h b/src/native/ascend/ops/rotary_embedding/kernel.h index b8e922e5b..a5000bc5d 100644 --- a/src/native/ascend/ops/rotary_embedding/kernel.h +++ b/src/native/ascend/ops/rotary_embedding/kernel.h @@ -18,6 +18,10 @@ #include "native/ascend/workspace_pool_.h" #include "operator.h" +#ifdef INFINI_HAS_CUSTOM_KERNELS +#include "aclrtlaunch_rotary_embedding_qk.h" +#endif + namespace infini::ops { // NeoX/GPT-J RoPE. Positions select rows from the packed [cos, sin] cache, @@ -42,6 +46,25 @@ class Operator assert(cos_sin_cache_type_ == query_type_ && "Ascend `RotaryEmbedding` requires cache and query dtypes to " "match"); +#ifdef INFINI_HAS_CUSTOM_KERNELS + use_qk_kernel_ = + query_type_ == DataType::kBFloat16 && is_neox_ && head_size_ == 128 && + rot_dim_ == 128 && num_heads_ > 0 && num_heads_ <= 32 && + query_head_stride_ == 128 && + query_token_stride_ >= static_cast(num_heads_ * 128) && + cos_sin_cache.stride(0) >= 128 && + (positions_ndim_ == 1 || + query.stride(0) == + static_cast(positions.size(1)) * query_token_stride_) && + (!key.has_value() || + (num_kv_heads_ <= 32 && key_head_stride_ == 128 && + key_token_stride_ >= static_cast(num_kv_heads_ * 128) && + (positions_ndim_ == 1 || + key->stride(0) == + static_cast(positions.size(1)) * key_token_stride_))); + if (use_qk_kernel_) return; +#endif + assert(cos_sin_cache.IsContiguous() && "Ascend `RotaryEmbedding` requires a contiguous cache tensor"); @@ -148,6 +171,20 @@ class Operator "Ascend `RotaryEmbedding` key presence changed after planning"); auto stream = static_cast(stream_); +#ifdef INFINI_HAS_CUSTOM_KERNELS + if (use_qk_kernel_) { + if (num_tokens_ == 0) return; + + RunQkKernel(positions, query, cos_sin_cache, num_heads_, + query_token_stride_, stream); + if (key.has_value()) { + RunQkKernel(positions, *key, cos_sin_cache, num_kv_heads_, + key_token_stride_, stream); + } + return; + } +#endif + if (cos_sin_cache.data() != cos_sin_cache_data_) { UploadCosSinCache(cos_sin_cache); cos_sin_cache_data_ = cos_sin_cache.data(); @@ -203,6 +240,21 @@ class Operator } private: +#ifdef INFINI_HAS_CUSTOM_KERNELS + void RunQkKernel(const Tensor positions, Tensor data, const Tensor cache, + int64_t heads, int64_t token_stride, + aclrtStream stream) const { + auto blocks = + static_cast(std::min(num_tokens_, 64)); + auto ret = aclrtlaunch_rotary_embedding_qk( + blocks, stream, const_cast(positions.data()), data.data(), + const_cast(cache.data()), static_cast(num_tokens_), + heads, token_stride, static_cast(cache.stride(0))); + assert(ret == ACL_SUCCESS && + "Ascend Q/K `RotaryEmbedding` execution failed"); + } +#endif + void RunOne(aclTensor* input, aclTensor* cos, aclTensor* sin, void* input_data, ascend::AclTensorCache& output_cache, void* output_data, aclOpExecutor*& executor, @@ -301,6 +353,8 @@ class Operator } } + bool use_qk_kernel_{false}; + int64_t max_seq_len_{0}; size_t element_size_{0}; bool has_key_{false}; diff --git a/src/native/ascend/ops/silu_and_mul/kernel.h b/src/native/ascend/ops/silu_and_mul/kernel.h index 84cb39471..07c8d9d3a 100644 --- a/src/native/ascend/ops/silu_and_mul/kernel.h +++ b/src/native/ascend/ops/silu_and_mul/kernel.h @@ -1,6 +1,7 @@ #ifndef INFINI_OPS_ASCEND_SILU_AND_MUL_KERNEL_H_ #define INFINI_OPS_ASCEND_SILU_AND_MUL_KERNEL_H_ +#include #include #include @@ -13,6 +14,10 @@ #include "native/ascend/workspace_pool_.h" #include "operator.h" +#ifdef INFINI_HAS_CUSTOM_KERNELS +#include "aclrtlaunch_silu_and_mul_tiled.h" +#endif + namespace infini::ops { // `aclnnSwiGlu` splits the last input dimension into `[gate, up]` and @@ -23,6 +28,22 @@ class Operator : public SiluAndMul { public: Operator(const Tensor input, Tensor out) : SiluAndMul(input, out), input_cache_(input), out_cache_(out) { +#ifdef INFINI_HAS_CUSTOM_KERNELS + // Padded outputs avoid the ACLNN staging and copy; contiguous outputs + // retain the vendor kernel's lower launch overhead. + use_tiled_kernel_ = + !is_out_contiguous_ && (ndim_ == 2 || ndim_ == 3) && hidden_size_ > 0 && + hidden_size_ <= 8192 && hidden_size_ % 128 == 0 && + input.stride(-1) == 1 && out.stride(-1) == 1 && + input.stride(-2) >= static_cast(hidden_size_ * 2) && + out.stride(-2) >= static_cast(hidden_size_) && + (ndim_ == 2 || (input.stride(0) == static_cast(input.size(1)) * + input.stride(1) && + out.stride(0) == static_cast(out.size(1)) * + out.stride(1))) && + input_type_ == DataType::kBFloat16; +#endif + needs_copy_ = !is_out_contiguous_; if (needs_copy_) { out_staging_size_ = out.numel() * kDataTypeToSize.at(out.dtype()); @@ -38,6 +59,24 @@ class Operator : public SiluAndMul { } void operator()(const Tensor input, Tensor out) const override { +#ifdef INFINI_HAS_CUSTOM_KERNELS + if (use_tiled_kernel_) { + auto rows = output_size_ / hidden_size_; + if (rows == 0) return; + + auto blocks = static_cast(std::min(rows, 8) * 8); + auto ret = aclrtlaunch_silu_and_mul_tiled( + blocks, static_cast(stream_), + const_cast(input.data()), out.data(), + static_cast(rows), static_cast(hidden_size_), + static_cast(input.stride(-2)), + static_cast(out.stride(-2))); + assert(ret == ACL_SUCCESS && + "Ascend tiled `SiluAndMul` execution failed"); + return; + } +#endif + auto t_input = input_cache_.get(const_cast(input.data())); auto t_out = out_cache_.get(out.data()); auto stream = static_cast(stream_); @@ -94,6 +133,8 @@ class Operator : public SiluAndMul { mutable std::optional out_staging_cache_; + bool use_tiled_kernel_{false}; + bool needs_copy_{false}; uint64_t out_staging_size_{0}; diff --git a/tests/test_reshape_and_cache_flash.py b/tests/test_reshape_and_cache_flash.py index 3191a5d0d..c577be10b 100644 --- a/tests/test_reshape_and_cache_flash.py +++ b/tests/test_reshape_and_cache_flash.py @@ -161,3 +161,56 @@ def _make_cache(num_blocks, block_size, num_heads, head_size, layout, dtype, dev physical_shape = (num_blocks, num_heads, block_size, head_size) return torch.randn(physical_shape, dtype=dtype, device=device).permute(0, 2, 1, 3) + + +@pytest.mark.parametrize( + "tokens, heads, head_size", + ((1, 8, 128), (8191, 8, 128), (8192, 8, 128), (3, 2, 8193)), +) +@pytest.mark.parametrize("layout", ("NHD", "HND")) +@pytest.mark.parametrize("padded_heads", (False, True)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_reshape_and_cache_flash_ascend_grid_stride( + tokens, heads, head_size, layout, padded_heads, dtype, device, implementation_index +): + if device != "npu" or implementation_index != 0: + pytest.skip("grid stride regression targets native Ascend") + + head_stride = head_size + 3 if padded_heads else head_size + # Token padding selects the custom path even for NHD caches. + strides = (heads * head_stride + 128, head_stride, 1) + key = randn_strided((tokens, heads, head_size), strides, dtype=dtype, device=device) + value = randn_strided( + (tokens, heads, head_size), strides, dtype=dtype, device=device + ) + block_size = 128 + blocks = (tokens + block_size - 1) // block_size + key_cache = _make_cache(blocks, block_size, heads, head_size, layout, dtype, device) + value_cache = _make_cache( + blocks, block_size, heads, head_size, layout, dtype, device + ) + slots = torch.arange(tokens, dtype=torch.int64, device=device) + slots[0] = -1 + expected_key = key_cache.clone(memory_format=torch.preserve_format) + expected_value = value_cache.clone(memory_format=torch.preserve_format) + token_indices = torch.arange(1, tokens, dtype=torch.int64, device=device) + block_indices, offsets = token_indices // block_size, token_indices % block_size + expected_key[block_indices, offsets] = key[1:] + expected_value[block_indices, offsets] = value[1:] + scale = torch.ones((), dtype=torch.float32, device=device) + + infini.ops.reshape_and_cache_flash( + key, + value, + slots, + scale, + scale, + "auto", + key_cache, + value_cache, + implementation_index=implementation_index, + stream=get_stream(key.device), + ) + + torch.testing.assert_close(key_cache, expected_key, rtol=0, atol=0) + torch.testing.assert_close(value_cache, expected_value, rtol=0, atol=0) diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 707f9a85f..b38db20f7 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -126,3 +126,57 @@ def _fallback(input, _normalized_shape, weight, *, eps=1e-6): out = result return out + + +@pytest.mark.parametrize("tokens, heads", ((1, 1), (31, 8), (8192, 32))) +@pytest.mark.parametrize("packed", (False, True)) +@pytest.mark.parametrize("eps", (1e-6, 1e-5)) +def test_rms_norm_ascend_packed_qk( + tokens, heads, packed, eps, device, implementation_index +): + if device != "npu" or implementation_index != 0: + pytest.skip("packed Q/K regression targets native Ascend") + + import torch_npu + + shape = (tokens, heads, 128) + strides = ((heads + 3) * 128, 128, 1) if packed else None + input = randn_strided(shape, strides, dtype=torch.bfloat16, device=device) + weight = torch.randn(128, dtype=torch.bfloat16, device=device) + out = empty_strided(shape, strides, dtype=torch.bfloat16, device=device) + expected = torch_npu.npu_rms_norm(input.contiguous(), weight, eps)[0] + + _rms_norm( + input, weight, eps=eps, out=out, implementation_index=implementation_index + ) + + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +def test_rms_norm_ascend_graph_updates_input_and_weight(device, implementation_index): + if device != "npu" or implementation_index != 0: + pytest.skip("Q/K graph regression targets native Ascend") + + import torch_npu + + input = torch.randn((64, 35, 128), dtype=torch.bfloat16, device=device)[:, :32] + weight = torch.randn(128, dtype=input.dtype, device=device) + out = torch.empty_like(input) + + def invoke(): + _rms_norm(input, weight, out=out, implementation_index=implementation_index) + + invoke() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + + with torch.npu.graph(graph): + invoke() + + input.add_(0.5) + weight.mul_(0.25) + graph.replay() + torch.npu.synchronize() + expected = torch_npu.npu_rms_norm(input.contiguous(), weight, 1e-6)[0] + + torch.testing.assert_close(out, expected, rtol=0, atol=0) diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index c306fa5e7..60f403e9c 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -324,3 +324,95 @@ def _apply_rotary( y = rotary[..., 1::2].float().clone() rotary[..., 0::2].copy_(x * cos - y * sin) rotary[..., 1::2].copy_(y * cos + x * sin) + + +@pytest.mark.parametrize("tokens", (1, 32, 8192)) +@pytest.mark.parametrize("has_key", (False, True)) +@pytest.mark.parametrize("padded_cache", (False, True)) +def test_rotary_embedding_ascend_packed_qk( + tokens, has_key, padded_cache, device, implementation_index +): + if device != "npu" or implementation_index != 0: + pytest.skip("packed Q/K regression targets native Ascend") + + query_width, key_width = 32 * 128, 8 * 128 + fused = torch.randn( + (tokens, query_width + 2 * key_width), dtype=torch.bfloat16, device=device + ) + query = fused[:, :query_width].view(tokens, 32, 128) + key = ( + fused[:, query_width : query_width + key_width].view(tokens, 8, 128) + if has_key + else None + ) + positions = torch.arange(tokens, dtype=torch.int64, device=device) % 128 + cache = torch.randn( + (128, 160 if padded_cache else 128), dtype=query.dtype, device=device + )[:, :128] + expected_query = query.clone() + expected_key = key.clone() if key is not None else None + _torch_rotary_embedding( + positions, expected_query, expected_key, cache, 128, True, 0, False + ) + + infini.ops.rotary_embedding( + positions, + query, + key, + cache, + 128, + True, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + + torch.testing.assert_close(query, expected_query, rtol=0, atol=0) + + if key is not None: + torch.testing.assert_close(key, expected_key, rtol=0, atol=0) + + +def test_rotary_embedding_ascend_graph_updates_cache_and_positions( + device, implementation_index +): + if device != "npu" or implementation_index != 0: + pytest.skip("Q/K graph regression targets native Ascend") + + source = torch.randn((32, 40, 128), dtype=torch.bfloat16, device=device) + query = source[:, :32].clone() + key = source[:, 32:].clone() + positions = torch.arange(32, dtype=torch.int64, device=device) + cache = torch.randn((64, 128), dtype=query.dtype, device=device) + + def invoke(): + infini.ops.rotary_embedding( + positions, + query, + key, + cache, + 128, + True, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + + invoke() + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + + with torch.npu.graph(graph): + invoke() + + query.copy_(source[:, :32]) + key.copy_(source[:, 32:]) + cache.mul_(0.25) + positions.add_(16) + expected_query, expected_key = query.clone(), key.clone() + _torch_rotary_embedding( + positions, expected_query, expected_key, cache, 128, True, 0, False + ) + graph.replay() + torch.npu.synchronize() + + torch.testing.assert_close(query, expected_query, rtol=0, atol=0) + torch.testing.assert_close(key, expected_key, rtol=0, atol=0) diff --git a/tests/test_silu_and_mul.py b/tests/test_silu_and_mul.py index c5d23e784..99a2f5f6e 100644 --- a/tests/test_silu_and_mul.py +++ b/tests/test_silu_and_mul.py @@ -111,3 +111,37 @@ def _torch_silu_and_mul(input, out): up = input[..., hidden_size:] return torch.mul(torch.nn.functional.silu(gate), up, out=out) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("tokens", (1, 31, 8192)) +@pytest.mark.parametrize("hidden", (3072, 6144, 8192)) +@pytest.mark.parametrize("padded", (False, True)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_silu_and_mul_ascend_token_tiles( + tokens, hidden, padded, dtype, device, implementation_index +): + if device != "npu" or implementation_index != 0: + pytest.skip("token tile regression targets native Ascend") + + input = rand_strided( + (tokens, hidden * 2), + (hidden * 2 + 128, 1) if padded else None, + dtype=dtype, + device=device, + ) + out = empty_strided( + (tokens, hidden), + (hidden + 128, 1) if padded else None, + dtype=dtype, + device=device, + ) + + return Payload( + lambda *args: _silu_and_mul(*args, implementation_index=implementation_index), + _torch_silu_and_mul, + (input, out), + {}, + rtol=1e-2 if dtype == torch.bfloat16 else 1e-3, + atol=5e-3 if dtype == torch.bfloat16 else 1e-5, + )