add kv rsnorm rope cache archive task

archive_tasks/kv_rmsnorm_rope_cache/design/block_level/kv_rmsnorm_rope.py

@@ -0,0 +1,161 @@
+"""Block-level TileLang design for KvRmsnormRopeCache.
+
+Fused kernel that computes:
+  1. RMSNorm over the first rms_size elements of each position.
+  2. RoPE over the remaining rope_size elements of each position.
+
+Task partitioning:
+  - The total number of independent positions is B * S * N.
+  - Each block processes a contiguous chunk of positions.
+  - Within a block, each position is processed serially (row-wise) because
+    a single row (rms_size=512 or rope_size=64) easily fits in UB.
+"""
+
+import tilelang
+import tilelang.language as T
+
+pass_configs = {
+    tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
+    tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True,
+}
+
+
+@tilelang.jit(out_idx=[4, 5], pass_configs=pass_configs)
+def kv_rmsnorm_rope(total, rms_size, rope_size, eps=1e-5, dtype="float16"):
+    """Kernel generator for KV RMSNorm + RoPE.
+
+    Parameters
+    ----------
+    total : int
+        Total number of positions (B * S * N).
+    rms_size : int
+        Size of the RMSNorm dimension (e.g. 512).
+    rope_size : int
+        Size of the RoPE dimension (e.g. 64).
+    eps : float
+        RMSNorm epsilon.
+    dtype : str
+        Input/output dtype (float16 or bfloat16).
+    """
+    block_size = 64
+    num_physical_cores = 20
+    total_blocks = (total + block_size - 1) // block_size
+    used_core_num = min(num_physical_cores, total_blocks)
+    tasks_per_core = (total_blocks + used_core_num - 1) // used_core_num
+    vec_num = 2
+    sub_block_size = block_size // vec_num
+
+    need_cast = dtype != "float32"
+    out_cast_mode = "CAST_ROUND" if dtype == "bfloat16" else "CAST_NONE"
+    eps_const = T.float32(eps)
+    inv_rms_size = T.float32(1.0 / rms_size)
+
+    @T.prim_func
+    def main(
+        rms_in: T.Tensor((total, rms_size), dtype),
+        gamma: T.Tensor((rms_size,), dtype),
+        k_input: T.Tensor((total, rope_size), dtype),
+        cos: T.Tensor((total, rope_size), dtype),
+        sin: T.Tensor((total, rope_size), dtype),
+        v_out: T.Tensor((total, rms_size), dtype),
+        k_embed_out: T.Tensor((total, rope_size), dtype),
+    ):
+        with T.Kernel(used_core_num, is_npu=True) as (cid, vid):
+            core_idx = cid
+
+            with T.Scope("V"):
+                # Pre-load gamma into UB once per core
+                gamma_in_ub = T.alloc_ub((rms_size,), dtype)
+                gamma_ub = T.alloc_ub((rms_size,), "float32")
+                if need_cast:
+                    T.copy(gamma[0], gamma_in_ub)
+                    T.tile.cast(gamma_ub, gamma_in_ub, mode="CAST_NONE", count=rms_size)
+                else:
+                    T.copy(gamma[0], gamma_ub)
+
+                # EPS and inv_N constants broadcast buffers
+                eps_ub = T.alloc_ub((1,), "float32")
+                inv_n_ub = T.alloc_ub((1,), "float32")
+                T.tile.fill(eps_ub, eps_const)
+                T.tile.fill(inv_n_ub, inv_rms_size)
+
+                # Row buffers for RMSNorm
+                x_in_ub = T.alloc_ub((rms_size,), dtype)
+                x_ub = T.alloc_ub((rms_size,), "float32")
+                x_sq_ub = T.alloc_ub((rms_size,), "float32")
+                sum_sq_ub = T.alloc_ub((1,), "float32")
+                inv_rms_ub = T.alloc_ub((1,), "float32")
+                out_ub = T.alloc_ub((rms_size,), "float32")
+                out_cast_ub = T.alloc_ub((rms_size,), dtype)
+                reduce_tmp = T.alloc_ub((2 * rms_size,), "uint8")
+
+                # Row buffers for RoPE
+                k_in_ub = T.alloc_ub((rope_size,), dtype)
+                k_ub = T.alloc_ub((rope_size,), "float32")
+                cos_ub = T.alloc_ub((rope_size,), "float32")
+                sin_ub = T.alloc_ub((rope_size,), "float32")
+                rotate_half_ub = T.alloc_ub((rope_size,), "float32")
+                tmp1_ub = T.alloc_ub((rope_size,), "float32")
+                tmp2_ub = T.alloc_ub((rope_size,), "float32")
+                k_embed_ub = T.alloc_ub((rope_size,), "float32")
+                k_embed_cast_ub = T.alloc_ub((rope_size,), dtype)
+                rope_reduce_tmp = T.alloc_ub((2 * rope_size,), "uint8")
+
+                for local_idx in T.serial(tasks_per_core):
+                    bx = core_idx * tasks_per_core + local_idx
+                    if bx < total_blocks:
+                        for row in T.serial(sub_block_size):
+                            pos = bx * block_size + vid * sub_block_size + row
+                            if pos < total:
+                                # ---- RMSNorm ----
+                                if need_cast:
+                                    T.copy(rms_in[pos, :], x_in_ub)
+                                    T.tile.cast(x_ub, x_in_ub, mode="CAST_NONE", count=rms_size)
+                                else:
+                                    T.copy(rms_in[pos, :], x_ub)
+
+                                T.tile.mul(x_sq_ub, x_ub, x_ub)
+                                T.reduce_sum(x_sq_ub, sum_sq_ub, reduce_tmp, dim=-1)
+                                T.tile.mul(sum_sq_ub, sum_sq_ub, inv_n_ub[0])
+                                T.tile.add(sum_sq_ub, sum_sq_ub, eps_ub[0])
+                                T.tile.rsqrt(inv_rms_ub, sum_sq_ub)
+
+                                inv_rms = inv_rms_ub[0]
+                                T.tile.mul(out_ub, x_ub, inv_rms)
+                                T.tile.mul(out_ub, out_ub, gamma_ub)
+
+                                if need_cast:
+                                    T.tile.cast(out_cast_ub, out_ub, mode=out_cast_mode, count=rms_size)
+                                    T.copy(out_cast_ub, v_out[pos, :])
+                                else:
+                                    T.copy(out_ub, v_out[pos, :])
+
+                                # ---- RoPE ----
+                                if need_cast:
+                                    T.copy(k_input[pos, :], k_in_ub)
+                                    T.tile.cast(k_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                    T.copy(cos[pos, :], k_in_ub)
+                                    T.tile.cast(cos_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                    T.copy(sin[pos, :], k_in_ub)
+                                    T.tile.cast(sin_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                else:
+                                    T.copy(k_input[pos, :], k_ub)
+                                    T.copy(cos[pos, :], cos_ub)
+                                    T.copy(sin[pos, :], sin_ub)
+
+                                # rotate_half: [-k[rope_size//2:], k[:rope_size//2]]
+                                half = rope_size // 2
+                                T.tile.neg(rotate_half_ub[:half], k_ub[half:rope_size])
+                                T.copy(k_ub[:half], rotate_half_ub[half:rope_size])
+
+                                T.tile.mul(tmp1_ub, k_ub, cos_ub)
+                                T.tile.mul(tmp2_ub, rotate_half_ub, sin_ub)
+                                T.tile.add(k_embed_ub, tmp1_ub, tmp2_ub)
+
+                                if need_cast:
+                                    T.tile.cast(k_embed_cast_ub, k_embed_ub, mode=out_cast_mode, count=rope_size)
+                                    T.copy(k_embed_cast_ub, k_embed_out[pos, :])
+                                else:
+                                    T.copy(k_embed_ub, k_embed_out[pos, :])
+
+    return main

archive_tasks/kv_rmsnorm_rope_cache/design/tile_level/kv_rmsnorm_rope.py

@@ -0,0 +1,132 @@
+"""Tile-level TileLang design for KvRmsnormRopeCache.
+
+Completes the block-level skeleton with full tile-level compute details.
+"""
+
+import tilelang
+import tilelang.language as T
+
+pass_configs = {
+    tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
+    tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True,
+}
+
+
+@tilelang.jit(out_idx=[4, 5], pass_configs=pass_configs)
+def kv_rmsnorm_rope(total, rms_size, rope_size, eps=1e-5, dtype="float16"):
+    block_size = 64
+    num_physical_cores = 20
+    total_blocks = (total + block_size - 1) // block_size
+    used_core_num = min(num_physical_cores, total_blocks)
+    tasks_per_core = (total_blocks + used_core_num - 1) // used_core_num
+    vec_num = 2
+    sub_block_size = block_size // vec_num
+
+    need_cast = dtype != "float32"
+    out_cast_mode = "CAST_ROUND" if dtype == "bfloat16" else "CAST_NONE"
+    eps_const = T.float32(eps)
+    inv_rms_size = T.float32(1.0 / rms_size)
+
+    @T.prim_func
+    def main(
+        rms_in: T.Tensor((total, rms_size), dtype),
+        gamma: T.Tensor((rms_size,), dtype),
+        k_input: T.Tensor((total, rope_size), dtype),
+        cos: T.Tensor((total, rope_size), dtype),
+        sin: T.Tensor((total, rope_size), dtype),
+        v_out: T.Tensor((total, rms_size), dtype),
+        k_embed_out: T.Tensor((total, rope_size), dtype),
+    ):
+        with T.Kernel(used_core_num, is_npu=True) as (cid, vid):
+            core_idx = cid
+
+            with T.Scope("V"):
+                gamma_in_ub = T.alloc_ub((rms_size,), dtype)
+                gamma_ub = T.alloc_ub((rms_size,), "float32")
+                if need_cast:
+                    T.copy(gamma[0], gamma_in_ub)
+                    T.tile.cast(gamma_ub, gamma_in_ub, mode="CAST_NONE", count=rms_size)
+                else:
+                    T.copy(gamma[0], gamma_ub)
+
+                eps_ub = T.alloc_ub((1,), "float32")
+                inv_n_ub = T.alloc_ub((1,), "float32")
+                T.tile.fill(eps_ub, eps_const)
+                T.tile.fill(inv_n_ub, inv_rms_size)
+
+                x_in_ub = T.alloc_ub((rms_size,), dtype)
+                x_ub = T.alloc_ub((rms_size,), "float32")
+                x_sq_ub = T.alloc_ub((rms_size,), "float32")
+                sum_sq_ub = T.alloc_ub((1,), "float32")
+                inv_rms_ub = T.alloc_ub((1,), "float32")
+                out_ub = T.alloc_ub((rms_size,), "float32")
+                out_cast_ub = T.alloc_ub((rms_size,), dtype)
+                reduce_tmp = T.alloc_ub((2 * rms_size,), "uint8")
+
+                k_in_ub = T.alloc_ub((rope_size,), dtype)
+                k_ub = T.alloc_ub((rope_size,), "float32")
+                cos_ub = T.alloc_ub((rope_size,), "float32")
+                sin_ub = T.alloc_ub((rope_size,), "float32")
+                rotate_half_ub = T.alloc_ub((rope_size,), "float32")
+                tmp1_ub = T.alloc_ub((rope_size,), "float32")
+                tmp2_ub = T.alloc_ub((rope_size,), "float32")
+                k_embed_ub = T.alloc_ub((rope_size,), "float32")
+                k_embed_cast_ub = T.alloc_ub((rope_size,), dtype)
+
+                for local_idx in T.serial(tasks_per_core):
+                    bx = core_idx * tasks_per_core + local_idx
+                    if bx < total_blocks:
+                        for row in T.serial(sub_block_size):
+                            pos = bx * block_size + vid * sub_block_size + row
+                            if pos < total:
+                                # RMSNorm
+                                if need_cast:
+                                    T.copy(rms_in[pos, :], x_in_ub)
+                                    T.tile.cast(x_ub, x_in_ub, mode="CAST_NONE", count=rms_size)
+                                else:
+                                    T.copy(rms_in[pos, :], x_ub)
+
+                                T.tile.mul(x_sq_ub, x_ub, x_ub)
+                                T.reduce_sum(x_sq_ub, sum_sq_ub, reduce_tmp, dim=-1)
+                                T.tile.mul(sum_sq_ub, sum_sq_ub, inv_n_ub[0])
+                                T.tile.add(sum_sq_ub, sum_sq_ub, eps_ub[0])
+                                T.tile.rsqrt(inv_rms_ub, sum_sq_ub)
+
+                                inv_rms = inv_rms_ub[0]
+                                T.tile.mul(out_ub, x_ub, inv_rms)
+                                T.tile.mul(out_ub, out_ub, gamma_ub)
+
+                                if need_cast:
+                                    T.tile.cast(out_cast_ub, out_ub, mode=out_cast_mode, count=rms_size)
+                                    T.copy(out_cast_ub, v_out[pos, :])
+                                else:
+                                    T.copy(out_ub, v_out[pos, :])
+
+                                # RoPE
+                                if need_cast:
+                                    T.copy(k_input[pos, :], k_in_ub)
+                                    T.tile.cast(k_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                    T.copy(cos[pos, :], k_in_ub)
+                                    T.tile.cast(cos_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                    T.copy(sin[pos, :], k_in_ub)
+                                    T.tile.cast(sin_ub, k_in_ub, mode="CAST_NONE", count=rope_size)
+                                else:
+                                    T.copy(k_input[pos, :], k_ub)
+                                    T.copy(cos[pos, :], cos_ub)
+                                    T.copy(sin[pos, :], sin_ub)
+
+                                half = rope_size // 2
+                                T.tile.neg(rotate_half_ub[:half], k_ub[half:rope_size])
+                                T.copy(k_ub[:half], rotate_half_ub[half:rope_size])
+
+                                T.tile.mul(tmp1_ub, k_ub, cos_ub)
+                                T.tile.mul(tmp2_ub, rotate_half_ub, sin_ub)
+                                T.tile.add(k_embed_ub, tmp1_ub, tmp2_ub)
+
+                                if need_cast:
+                                    T.tile.cast(k_embed_cast_ub, k_embed_ub, mode=out_cast_mode, count=rope_size)
+                                    T.copy(k_embed_cast_ub, k_embed_out[pos, :])
+                                else:
+                                    T.copy(k_embed_ub, k_embed_out[pos, :])
+
+    return main

archive_tasks/kv_rmsnorm_rope_cache/kernel/kernel_common.h

@@ -0,0 +1,24 @@
+#ifndef KERNEL_COMMON_H
+#define KERNEL_COMMON_H
+
+#include <cstddef>
+#include <cstdint>
+
+#include "kernel_operator.h"
+
+__aicore__ inline uint32_t CeilDivU32(uint32_t a, uint32_t b)
+{
+    return (a + b - 1U) / b;
+}
+
+template <typename T>
+__aicore__ inline void CopyTiling(T *tiling, GM_ADDR tilingGM)
+{
+    int32_t *dst = reinterpret_cast<int32_t *>(tiling);
+    auto *src = reinterpret_cast<__gm__ int32_t *>(tilingGM);
+    for (size_t i = 0; i < sizeof(T) / sizeof(int32_t); ++i) {
+        dst[i] = src[i];
+    }
+}
+
+#endif

archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache.cpp

@@ -0,0 +1,59 @@
+#include "kernel_operator.h"
+#include "kv_rmsnorm_rope_cache_kernel.h"
+#include "kv_rmsnorm_rope_cache_tiling.h"
+
+extern "C" __global__ __aicore__ void kv_rmsnorm_rope_cache_custom_fp16(
+    GM_ADDR kv, GM_ADDR gamma, GM_ADDR cos, GM_ADDR sin,
+    GM_ADDR index, GM_ADDR k_cache, GM_ADDR ckv_cache,
+    GM_ADDR k_cache_out, GM_ADDR ckv_cache_out,
+    GM_ADDR k_embed_out, GM_ADDR v_out,
+    GM_ADDR tiling)
+{
+    AscendC::TPipe pipe;
+    KvRmsnormRopeCacheKernel<half> kernel;
+    kernel.Init(kv, gamma, cos, sin, index, k_cache, ckv_cache,
+                k_cache_out, ckv_cache_out, k_embed_out, v_out,
+                tiling, &pipe);
+    kernel.Process();
+}
+
+extern "C" void kv_rmsnorm_rope_cache_do_fp16(
+    uint32_t blockDim, void *stream,
+    uint8_t *kv, uint8_t *gamma, uint8_t *cos, uint8_t *sin,
+    uint8_t *index, uint8_t *k_cache, uint8_t *ckv_cache,
+    uint8_t *k_cache_out, uint8_t *ckv_cache_out,
+    uint8_t *k_embed_out, uint8_t *v_out,
+    uint8_t *tiling)
+{
+    kv_rmsnorm_rope_cache_custom_fp16<<<blockDim, nullptr, stream>>>(
+        kv, gamma, cos, sin, index, k_cache, ckv_cache,
+        k_cache_out, ckv_cache_out, k_embed_out, v_out, tiling);
+}
+
+extern "C" __global__ __aicore__ void kv_rmsnorm_rope_cache_custom_bf16(
+    GM_ADDR kv, GM_ADDR gamma, GM_ADDR cos, GM_ADDR sin,
+    GM_ADDR index, GM_ADDR k_cache, GM_ADDR ckv_cache,
+    GM_ADDR k_cache_out, GM_ADDR ckv_cache_out,
+    GM_ADDR k_embed_out, GM_ADDR v_out,
+    GM_ADDR tiling)
+{
+    AscendC::TPipe pipe;
+    KvRmsnormRopeCacheKernel<bfloat16_t> kernel;
+    kernel.Init(kv, gamma, cos, sin, index, k_cache, ckv_cache,
+                k_cache_out, ckv_cache_out, k_embed_out, v_out,
+                tiling, &pipe);
+    kernel.Process();
+}
+
+extern "C" void kv_rmsnorm_rope_cache_do_bf16(
+    uint32_t blockDim, void *stream,
+    uint8_t *kv, uint8_t *gamma, uint8_t *cos, uint8_t *sin,
+    uint8_t *index, uint8_t *k_cache, uint8_t *ckv_cache,
+    uint8_t *k_cache_out, uint8_t *ckv_cache_out,
+    uint8_t *k_embed_out, uint8_t *v_out,
+    uint8_t *tiling)
+{
+    kv_rmsnorm_rope_cache_custom_bf16<<<blockDim, nullptr, stream>>>(
+        kv, gamma, cos, sin, index, k_cache, ckv_cache,
+        k_cache_out, ckv_cache_out, k_embed_out, v_out, tiling);
+}