perf(ascend): accelerate packed Q/K and strided KV/SwiGLU kernels

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()

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-<CONFIG>/`, while

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<uint32_t>(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<uint8_t>();
-    AscendC::DataCopyExtParams copy_params{
-        1, static_cast<uint32_t>(head_bytes_), 0, 0, 0};
-    AscendC::DataCopyPadExtParams<uint8_t> 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<uint8_t>& source, int64_t source_offset,
-      AscendC::GlobalTensor<uint8_t>& destination, int64_t destination_offset,
-      AscendC::LocalTensor<uint8_t>& local,
-      const AscendC::DataCopyExtParams& copy_params,
-      const AscendC::DataCopyPadExtParams<uint8_t>& pad_params) {
-    AscendC::DataCopyPad(local, source[source_offset], copy_params, pad_params);
-    AscendC::PipeBarrier<PIPE_ALL>();
-    AscendC::DataCopyPad(destination[destination_offset], local, copy_params);
-    AscendC::PipeBarrier<PIPE_ALL>();
+  __aicore__ inline void Copy(AscendC::GlobalTensor<uint8_t>& source,
+                              int64_t source_offset,
+                              AscendC::GlobalTensor<uint8_t>& destination,
+                              int64_t destination_offset, int64_t bytes,
+                              AscendC::LocalTensor<uint8_t>& local) {
+    AscendC::DataCopyPadExtParams<uint8_t> pad{false, 0, 0, 0};
+    for (int64_t offset = 0; offset < bytes; offset += kTileBytes) {
+      auto count = static_cast<uint32_t>(
+          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<AscendC::HardEvent::MTE2_MTE3>(0);
+      AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(0);
+      AscendC::DataCopyPad(destination[destination_offset + offset], local,
+                           params);
+      AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(0);
+      AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(0);
+    }
   }
 
+  static constexpr int64_t kTileBytes = 2048;
+
+  bool combine_heads_{false};
+
   AscendC::TPipe pipe_;
   AscendC::TBuf<AscendC::TPosition::VECCALC> copy_buffer_;
   AscendC::GlobalTensor<uint8_t> key_;