diff --git a/archive_tasks/kv_rmsnorm_rope_cache/design/block_level/kv_rmsnorm_rope.py b/archive_tasks/kv_rmsnorm_rope_cache/design/block_level/kv_rmsnorm_rope.py new file mode 100644 index 00000000..0d1e44cc --- /dev/null +++ b/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 diff --git a/archive_tasks/kv_rmsnorm_rope_cache/design/tile_level/kv_rmsnorm_rope.py b/archive_tasks/kv_rmsnorm_rope_cache/design/tile_level/kv_rmsnorm_rope.py new file mode 100644 index 00000000..0cf01b90 --- /dev/null +++ b/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 diff --git a/archive_tasks/kv_rmsnorm_rope_cache/kernel/kernel_common.h b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kernel_common.h new file mode 100644 index 00000000..56d556b7 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kernel_common.h @@ -0,0 +1,24 @@ +#ifndef KERNEL_COMMON_H +#define KERNEL_COMMON_H + +#include +#include + +#include "kernel_operator.h" + +__aicore__ inline uint32_t CeilDivU32(uint32_t a, uint32_t b) +{ + return (a + b - 1U) / b; +} + +template +__aicore__ inline void CopyTiling(T *tiling, GM_ADDR tilingGM) +{ + int32_t *dst = reinterpret_cast(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 diff --git a/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache.cpp b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache.cpp new file mode 100644 index 00000000..e57794e4 --- /dev/null +++ b/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 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<<>>( + 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 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<<>>( + kv, gamma, cos, sin, index, k_cache, ckv_cache, + k_cache_out, ckv_cache_out, k_embed_out, v_out, tiling); +} diff --git a/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_kernel.h b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_kernel.h new file mode 100644 index 00000000..4b65baef --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_kernel.h @@ -0,0 +1,612 @@ +#pragma once + +#include "kernel_operator.h" +#include "kernel_common.h" +#include "kv_rmsnorm_rope_cache_tiling.h" + +template +class KvRmsnormRopeCacheKernel { +public: + __aicore__ inline void Init( + 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 tilingGM, AscendC::TPipe *pipe) + { + CopyTiling(&tiling_, tilingGM); + + int32_t hidden_size = tiling_.hidden_size; + int32_t total = tiling_.total; + + kvGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(kv), total * hidden_size); + gammaGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(gamma), tiling_.rms_size); + cosGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(cos), total * tiling_.rope_size); + sinGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(sin), total * tiling_.rope_size); + indexGM_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(index), tiling_.index_numel); + kCacheGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(k_cache), + tiling_.k_cache_dim0 * tiling_.k_cache_dim1 * tiling_.k_cache_dim2 * tiling_.k_cache_dim3); + ckvCacheGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(ckv_cache), + tiling_.ckv_cache_dim0 * tiling_.ckv_cache_dim1 * tiling_.ckv_cache_dim2 * tiling_.ckv_cache_dim3); + kCacheOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(k_cache_out), + tiling_.k_cache_dim0 * tiling_.k_cache_dim1 * tiling_.k_cache_dim2 * tiling_.k_cache_dim3); + ckvCacheOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(ckv_cache_out), + tiling_.ckv_cache_dim0 * tiling_.ckv_cache_dim1 * tiling_.ckv_cache_dim2 * tiling_.ckv_cache_dim3); + + if (tiling_.is_output_kv) { + kEmbedOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(k_embed_out), total * tiling_.rope_size); + vOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ dataType *>(v_out), total * tiling_.rms_size); + } + + if ASCEND_IS_AIV { + pipe_ = pipe; + + pipe_->InitBuffer(gammaInQueue_, 1, tiling_.rms_size * sizeof(dataType)); + pipe_->InitBuffer(rmsInQueue_, 1, tiling_.rms_size * sizeof(dataType)); + pipe_->InitBuffer(ropeInQueue_, 1, tiling_.rope_size * sizeof(dataType)); + pipe_->InitBuffer(cosInQueue_, 1, tiling_.rope_size * sizeof(dataType)); + pipe_->InitBuffer(sinInQueue_, 1, tiling_.rope_size * sizeof(dataType)); + pipe_->InitBuffer(vOutQueue_, 1, tiling_.rms_size * sizeof(dataType)); + pipe_->InitBuffer(kEmbedOutQueue_, 1, tiling_.rope_size * sizeof(dataType)); + + pipe_->InitBuffer(reduceBuf_, tiling_.rms_size * sizeof(float)); + pipe_->InitBuffer(sumBuf_, 32 * sizeof(float)); + pipe_->InitBuffer(invRmsBuf_, 32 * sizeof(float)); + pipe_->InitBuffer(vFloatBuf_, tiling_.rms_size * sizeof(float)); + pipe_->InitBuffer(ropeInFloatBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(rotateHalfFloatBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(kEmbedFloatBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(scatterKBuf_, 32 * sizeof(dataType)); + pipe_->InitBuffer(scatterVBuf_, 32 * sizeof(dataType)); + + if constexpr (!std::is_same_v) { + pipe_->InitBuffer(rmsCastBuf_, tiling_.rms_size * sizeof(float)); + pipe_->InitBuffer(ropeCastBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(cosCastBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(sinCastBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(vCastBuf_, tiling_.rms_size * sizeof(float)); + pipe_->InitBuffer(kEmbedCastBuf_, tiling_.rope_size * sizeof(float)); + pipe_->InitBuffer(gammaCastBuf_, tiling_.rms_size * sizeof(float)); + } + + AscendC::LocalTensor gammaInTmp; + gammaInQueue_.AllocTensor(gammaInTmp); + LoadGmToUb(gammaInTmp, gammaGM_, static_cast(tiling_.rms_size)); + gammaInQueue_.EnQue(gammaInTmp); + gammaInQueue_.DeQue(gammaInTmp); + AscendC::PipeBarrier(); + + gammaFloatLocal_ = gammaCastBuf_.Get(); + AscendC::Cast(gammaFloatLocal_, gammaInTmp, AscendC::RoundMode::CAST_NONE, tiling_.rms_size); + AscendC::PipeBarrier(); + gammaInQueue_.FreeTensor(gammaInTmp); + } + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIV { + const int blockIdx = AscendC::GetBlockIdx(); + + for (int localIdx = 0; localIdx < tiling_.tasksPerCore; ++localIdx) { + const int bx = blockIdx * tiling_.tasksPerCore + localIdx; + if (bx >= BlockCount()) { + continue; + } + for (int row = 0; row < tiling_.block_size; ++row) { + const int pos = bx * tiling_.block_size + row; + if (pos < tiling_.total) { + ProcessPosition(pos); + } + } + } + } + } + +private: + __aicore__ inline int32_t BlockCount() const + { + return (tiling_.total + tiling_.block_size - 1) / tiling_.block_size; + } + + __aicore__ inline AscendC::RoundMode OutputRoundMode() const + { + if constexpr (std::is_same_v) { + return AscendC::RoundMode::CAST_ROUND; + } + return AscendC::RoundMode::CAST_NONE; + } + + __aicore__ inline void LoadGmToUb(AscendC::LocalTensor dst, + AscendC::GlobalTensor src, uint32_t count) + { + AscendC::DataCopy(dst, src, count); + } + + __aicore__ inline void StoreUbToGm(AscendC::GlobalTensor dst, + AscendC::LocalTensor src, uint32_t count) + { + AscendC::DataCopy(dst, src, count); + } + + __aicore__ inline void PrepareInputTensor( + AscendC::LocalTensor &dst, + AscendC::LocalTensor &src, + AscendC::TBuf &castBuf, + int32_t count) + { + if constexpr (std::is_same_v) { + dst = src.template ReinterpretCast(); + } else { + dst = castBuf.Get(); + AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void PrepareOutputTensor( + AscendC::LocalTensor &dst, + AscendC::LocalTensor &out, + AscendC::TBuf &castBuf, + int32_t count) + { + if constexpr (std::is_same_v) { + dst = out.template ReinterpretCast(); + } else { + (void)out; + dst = castBuf.Get(); + } + } + + __aicore__ inline void FinalizeOutputTensor( + AscendC::LocalTensor &out, + AscendC::LocalTensor &src, + int32_t count) + { + if constexpr (!std::is_same_v) { + AscendC::Cast(out, src, OutputRoundMode(), count); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void CopyInInputs(int32_t pos) + { + rmsInQueue_.AllocTensor(rmsInLocal_); + LoadGmToUb(rmsInLocal_, kvGM_[pos * tiling_.hidden_size], static_cast(tiling_.rms_size)); + rmsInQueue_.EnQue(rmsInLocal_); + + ropeInQueue_.AllocTensor(ropeInLocal_); + LoadGmToUb(ropeInLocal_, kvGM_[pos * tiling_.hidden_size + tiling_.rms_size], static_cast(tiling_.rope_size)); + ropeInQueue_.EnQue(ropeInLocal_); + + cosInQueue_.AllocTensor(cosInLocal_); + LoadGmToUb(cosInLocal_, cosGM_[pos * tiling_.rope_size], static_cast(tiling_.rope_size)); + cosInQueue_.EnQue(cosInLocal_); + + sinInQueue_.AllocTensor(sinInLocal_); + LoadGmToUb(sinInLocal_, sinGM_[pos * tiling_.rope_size], static_cast(tiling_.rope_size)); + sinInQueue_.EnQue(sinInLocal_); + } + + __aicore__ inline void CopyOutOutputs(int32_t pos, int32_t b, int32_t s, int32_t n) + { + vOutQueue_.DeQue(vOutLocal_); + kEmbedOutQueue_.DeQue(kEmbedOutLocal_); + + UpdateCache(b, s, n, vOutLocal_, kEmbedOutLocal_); + + if (tiling_.is_output_kv) { + StoreUbToGm(vOutGM_[pos * tiling_.rms_size], vOutLocal_, static_cast(tiling_.rms_size)); + StoreUbToGm(kEmbedOutGM_[pos * tiling_.rope_size], kEmbedOutLocal_, static_cast(tiling_.rope_size)); + } + + vOutQueue_.FreeTensor(vOutLocal_); + kEmbedOutQueue_.FreeTensor(kEmbedOutLocal_); + } + + __aicore__ inline void ProcessPosition(int32_t pos) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + + CopyInInputs(pos); + + // DeQue inputs + rmsInQueue_.DeQue(rmsInLocal_); + ropeInQueue_.DeQue(ropeInLocal_); + cosInQueue_.DeQue(cosInLocal_); + sinInQueue_.DeQue(sinInLocal_); + + // Cast inputs to float + AscendC::LocalTensor rmsInFloat; + PrepareInputTensor(rmsInFloat, rmsInLocal_, rmsCastBuf_, rms_size); + AscendC::LocalTensor ropeInFloat; + PrepareInputTensor(ropeInFloat, ropeInLocal_, ropeCastBuf_, rope_size); + AscendC::LocalTensor cosFloat; + PrepareInputTensor(cosFloat, cosInLocal_, cosCastBuf_, rope_size); + AscendC::LocalTensor sinFloat; + PrepareInputTensor(sinFloat, sinInLocal_, sinCastBuf_, rope_size); + + // Allocate outputs + vOutQueue_.AllocTensor(vOutLocal_); + kEmbedOutQueue_.AllocTensor(kEmbedOutLocal_); + + AscendC::LocalTensor vFloat; + AscendC::LocalTensor kEmbedFloat; + PrepareOutputTensor(vFloat, vOutLocal_, vCastBuf_, rms_size); + PrepareOutputTensor(kEmbedFloat, kEmbedOutLocal_, kEmbedCastBuf_, rope_size); + + // Compute RMSNorm + ComputeRmsNorm(rmsInFloat, vFloat, rms_size); + + // Compute RoPE + ComputeRoPE(ropeInFloat, cosFloat, sinFloat, kEmbedFloat, rope_size); + + // Finalize outputs (cast back from float) + FinalizeOutputTensor(vOutLocal_, vFloat, rms_size); + FinalizeOutputTensor(kEmbedOutLocal_, kEmbedFloat, rope_size); + + // Free input queues + rmsInQueue_.FreeTensor(rmsInLocal_); + ropeInQueue_.FreeTensor(ropeInLocal_); + cosInQueue_.FreeTensor(cosInLocal_); + sinInQueue_.FreeTensor(sinInLocal_); + + // EnQue outputs + vOutQueue_.EnQue(vOutLocal_); + kEmbedOutQueue_.EnQue(kEmbedOutLocal_); + + // Compute b, s, n from pos + int32_t B = tiling_.B; + int32_t N = tiling_.N; + int32_t S = tiling_.S; + int32_t b = pos / (S * N); + int32_t rem = pos % (S * N); + int32_t s = rem / N; + int32_t n = rem % N; + + CopyOutOutputs(pos, b, s, n); + } + + __aicore__ inline void ComputeRmsNorm(AscendC::LocalTensor x, + AscendC::LocalTensor y, int32_t count) + { + AscendC::LocalTensor xSq = reduceBuf_.Get(); + AscendC::Mul(xSq, x, x, count); + AscendC::PipeBarrier(); + + float sumVal = 0.0f; + for (int i = 0; i < count; ++i) { + sumVal += xSq.GetValue(i); + } + + float meanSq = sumVal * tiling_.invRmsSize + tiling_.eps; + float invRmsVal = 1.0f / sqrt(meanSq); + + AscendC::LocalTensor invRms = invRmsBuf_.Get(); + AscendC::Duplicate(invRms, invRmsVal, 1); + AscendC::PipeBarrier(); + + AscendC::Muls(y, x, invRmsVal, count); + AscendC::PipeBarrier(); + + AscendC::Mul(y, y, gammaFloatLocal_, count); + AscendC::PipeBarrier(); + } + + __aicore__ inline void ComputeRoPE(AscendC::LocalTensor ropeIn, + AscendC::LocalTensor cosLocal, + AscendC::LocalTensor sinLocal, + AscendC::LocalTensor kEmbedOut, + int32_t rope_size) + { + int32_t half = rope_size / 2; + AscendC::LocalTensor rotateHalf = rotateHalfFloatBuf_.Get(); + + // rotate_half: [-k2, k1] using element-wise operations + for (int i = 0; i < half; ++i) { + rotateHalf.SetValue(i, -ropeIn.GetValue(half + i)); + } + for (int i = 0; i < half; ++i) { + rotateHalf.SetValue(half + i, ropeIn.GetValue(i)); + } + AscendC::PipeBarrier(); + + // k_embed = k * cos + rotate_half * sin + // First half: k1 * cos1 + (-k2) * sin1 + for (int i = 0; i < half; ++i) { + float k1 = ropeIn.GetValue(i); + float c1 = cosLocal.GetValue(i); + float s1 = sinLocal.GetValue(i); + float rh = rotateHalf.GetValue(i); + kEmbedOut.SetValue(i, k1 * c1 + rh * s1); + } + // Second half: k2 * cos2 + k1 * sin2 + for (int i = 0; i < half; ++i) { + float k2 = ropeIn.GetValue(half + i); + float c2 = cosLocal.GetValue(half + i); + float s2 = sinLocal.GetValue(half + i); + float rh = rotateHalf.GetValue(half + i); + kEmbedOut.SetValue(half + i, k2 * c2 + rh * s2); + } + AscendC::PipeBarrier(); + } + + __aicore__ inline void UpdateCache(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t mode = tiling_.cache_mode; + + if (mode == 0) { + UpdateCacheNorm(b, s, n, vLocal, kEmbedLocal); + } else if (mode == 1 || mode == 2) { + UpdateCachePA(b, s, n, vLocal, kEmbedLocal); + } else if (mode == 3) { + UpdateCachePANZ(b, s, n, vLocal, kEmbedLocal); + } else if (mode == 4) { + UpdateCachePABlkBNSD(b, s, n, vLocal, kEmbedLocal); + } else if (mode == 5) { + UpdateCachePABlkNZ(b, s, n, vLocal, kEmbedLocal); + } + } + + __aicore__ inline void UpdateCacheNorm(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + int32_t N = tiling_.N; + int32_t S = tiling_.S; + int32_t max_seq = tiling_.k_cache_dim2; + int32_t B = tiling_.B; + + int64_t idx = 0; + if (tiling_.index_numel == B * S) { + idx = indexGM_.GetValue(b * S + s); + } else { + int pos = b * S + s; + if (pos < tiling_.index_numel) { + idx = indexGM_.GetValue(pos); + } + } + if (idx < 0 || idx >= max_seq) { + return; + } + + int64_t kOffset = ((int64_t)b * N * max_seq + n * max_seq + idx) * rope_size; + StoreUbToGm(kCacheOutGM_[kOffset], kEmbedLocal, rope_size); + + int64_t vOffset = ((int64_t)b * N * max_seq + n * max_seq + idx) * rms_size; + StoreUbToGm(ckvCacheOutGM_[vOffset], vLocal, rms_size); + } + + __aicore__ inline void UpdateCachePA(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + int32_t N = tiling_.N; + int32_t S = tiling_.S; + int pos = b * S + s; + if (pos >= tiling_.index_numel) { + return; + } + int64_t idx = indexGM_.GetValue(pos); + if (idx < 0) { + return; + } + + int64_t kCacheFlatSize = (int64_t)tiling_.k_cache_dim0 * tiling_.k_cache_dim1 * + tiling_.k_cache_dim2 * tiling_.k_cache_dim3; + int64_t cacheKN = kCacheFlatSize / rope_size; + if (idx >= cacheKN) { + return; + } + + int64_t kOffset = idx * rope_size; + StoreUbToGm(kCacheOutGM_[kOffset], kEmbedLocal, rope_size); + + int64_t ckvCacheFlatSize = (int64_t)tiling_.ckv_cache_dim0 * tiling_.ckv_cache_dim1 * + tiling_.ckv_cache_dim2 * tiling_.ckv_cache_dim3; + int64_t cacheVN = ckvCacheFlatSize / rms_size; + if (idx >= cacheVN) { + return; + } + int64_t vOffset = idx * rms_size; + StoreUbToGm(ckvCacheOutGM_[vOffset], vLocal, rms_size); + } + + __aicore__ inline void UpdateCachePANZ(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + int32_t N = tiling_.N; + int32_t S = tiling_.S; + int pos = b * S + s; + if (pos >= tiling_.index_numel) { + return; + } + int64_t idx = indexGM_.GetValue(pos); + if (idx < 0) { + return; + } + + int32_t block_size = tiling_.k_cache_dim1; + int32_t dk = tiling_.k_cache_dim3; + int32_t dv = tiling_.ckv_cache_dim3; + int32_t dk0 = 16; + int32_t dv0 = 16; + int32_t dk1 = dk / dk0; + int32_t dv1 = dv / dv0; + int32_t bn = tiling_.k_cache_dim0; + + int64_t bn_id = idx / block_size; + int64_t block_offset = idx % block_size; + if (bn_id >= bn) { + return; + } + + AscendC::LocalTensor scatterK = scatterKBuf_.Get(); + for (int d = 0; d < dk1; ++d) { + int64_t offset = ((bn_id * N * dk1 + n * dk1 + d) * block_size + block_offset) * dk0; + for (int i = 0; i < dk0; ++i) { + scatterK.SetValue(i, kEmbedLocal.GetValue(d * dk0 + i)); + } + StoreUbToGm(kCacheOutGM_[offset], scatterK, dk0); + AscendC::PipeBarrier(); + } + AscendC::LocalTensor scatterV = scatterVBuf_.Get(); + for (int d = 0; d < dv1; ++d) { + int64_t offset = ((bn_id * N * dv1 + n * dv1 + d) * block_size + block_offset) * dv0; + for (int i = 0; i < dv0; ++i) { + scatterV.SetValue(i, vLocal.GetValue(d * dv0 + i)); + } + StoreUbToGm(ckvCacheOutGM_[offset], scatterV, dv0); + AscendC::PipeBarrier(); + } + } + + __aicore__ inline void UpdateCachePABlkBNSD(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + int32_t S = tiling_.S; + int32_t block_size = tiling_.k_cache_dim1; + int32_t ceil_div_s = (S + block_size - 1) / block_size; + int32_t seq_id = s / block_size; + int32_t seq_start = seq_id * block_size; + int32_t idx_pos = b * ceil_div_s + seq_id; + if (idx_pos >= tiling_.index_numel) { + return; + } + int64_t idx_val = indexGM_.GetValue(idx_pos); + if (idx_val < 0) { + return; + } + int64_t cache_b = idx_val / block_size; + if (cache_b >= tiling_.k_cache_dim0) { + return; + } + + int32_t offset_in_block = s - seq_start; + int64_t kOffset = ((cache_b * block_size + offset_in_block) * tiling_.N + n) * rope_size; + int64_t vOffset = ((cache_b * block_size + offset_in_block) * tiling_.N + n) * rms_size; + StoreUbToGm(kCacheOutGM_[kOffset], kEmbedLocal, rope_size); + StoreUbToGm(ckvCacheOutGM_[vOffset], vLocal, rms_size); + } + + __aicore__ inline void UpdateCachePABlkNZ(int32_t b, int32_t s, int32_t n, + AscendC::LocalTensor vLocal, + AscendC::LocalTensor kEmbedLocal) + { + int32_t rms_size = tiling_.rms_size; + int32_t rope_size = tiling_.rope_size; + int32_t S = tiling_.S; + int32_t block_size = tiling_.k_cache_dim1; + int32_t ceil_div_s = (S + block_size - 1) / block_size; + int32_t seq_id = s / block_size; + int32_t seq_start = seq_id * block_size; + int32_t idx_pos = b * ceil_div_s + seq_id; + if (idx_pos >= tiling_.index_numel) { + return; + } + int64_t idx_val = indexGM_.GetValue(idx_pos); + if (idx_val < 0) { + return; + } + int64_t cache_b = idx_val / block_size; + int32_t bn = tiling_.k_cache_dim0; + if (cache_b >= bn) { + return; + } + + int32_t dk = tiling_.k_cache_dim3; + int32_t dv = tiling_.ckv_cache_dim3; + int32_t dk0 = 16; + int32_t dv0 = 16; + int32_t dk1 = dk / dk0; + int32_t dv1 = dv / dv0; + int32_t offset_in_block = s - seq_start; + + AscendC::LocalTensor scatterK = scatterKBuf_.Get(); + for (int d = 0; d < dk1; ++d) { + int64_t offset = ((cache_b * tiling_.N * dk1 + n * dk1 + d) * block_size + offset_in_block) * dk0; + for (int i = 0; i < dk0; ++i) { + scatterK.SetValue(i, kEmbedLocal.GetValue(d * dk0 + i)); + } + StoreUbToGm(kCacheOutGM_[offset], scatterK, dk0); + AscendC::PipeBarrier(); + } + AscendC::LocalTensor scatterV = scatterVBuf_.Get(); + for (int d = 0; d < dv1; ++d) { + int64_t offset = ((cache_b * tiling_.N * dv1 + n * dv1 + d) * block_size + offset_in_block) * dv0; + for (int i = 0; i < dv0; ++i) { + scatterV.SetValue(i, vLocal.GetValue(d * dv0 + i)); + } + StoreUbToGm(ckvCacheOutGM_[offset], scatterV, dv0); + AscendC::PipeBarrier(); + } + } + + KvRmsnormRopeCacheTiling tiling_; + AscendC::TPipe *pipe_; + + AscendC::GlobalTensor kvGM_; + AscendC::GlobalTensor gammaGM_; + AscendC::GlobalTensor cosGM_; + AscendC::GlobalTensor sinGM_; + AscendC::GlobalTensor indexGM_; + AscendC::GlobalTensor kCacheGM_; + AscendC::GlobalTensor ckvCacheGM_; + AscendC::GlobalTensor kCacheOutGM_; + AscendC::GlobalTensor ckvCacheOutGM_; + AscendC::GlobalTensor kEmbedOutGM_; + AscendC::GlobalTensor vOutGM_; + + // Input queues + AscendC::TQue gammaInQueue_; + AscendC::TQue rmsInQueue_; + AscendC::TQue ropeInQueue_; + AscendC::TQue cosInQueue_; + AscendC::TQue sinInQueue_; + + // Output queues + AscendC::TQue vOutQueue_; + AscendC::TQue kEmbedOutQueue_; + + // Compute buffers (TBuf) + AscendC::TBuf reduceBuf_; + AscendC::TBuf sumBuf_; + AscendC::TBuf invRmsBuf_; + AscendC::TBuf vFloatBuf_; + AscendC::TBuf ropeInFloatBuf_; + AscendC::TBuf rotateHalfFloatBuf_; + AscendC::TBuf kEmbedFloatBuf_; + AscendC::TBuf scatterKBuf_; + AscendC::TBuf scatterVBuf_; + + // Cast buffers (TBuf) + AscendC::TBuf rmsCastBuf_; + AscendC::TBuf ropeCastBuf_; + AscendC::TBuf cosCastBuf_; + AscendC::TBuf sinCastBuf_; + AscendC::TBuf vCastBuf_; + AscendC::TBuf kEmbedCastBuf_; + AscendC::TBuf gammaCastBuf_; + + // LocalTensor members + AscendC::LocalTensor rmsInLocal_; + AscendC::LocalTensor ropeInLocal_; + AscendC::LocalTensor cosInLocal_; + AscendC::LocalTensor sinInLocal_; + AscendC::LocalTensor vOutLocal_; + AscendC::LocalTensor kEmbedOutLocal_; + AscendC::LocalTensor gammaFloatLocal_; +}; diff --git a/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_tiling.h b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_tiling.h new file mode 100644 index 00000000..d7023a82 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/kernel/kv_rmsnorm_rope_cache_tiling.h @@ -0,0 +1,36 @@ +#ifndef KV_RMSNORM_ROPE_CACHE_TILING_H +#define KV_RMSNORM_ROPE_CACHE_TILING_H + +#include + +constexpr uint32_t DEFAULT_BLOCK_SIZE = 64; +constexpr uint32_t DEFAULT_NUM_PHYSICAL_CORES = 20; + +struct KvRmsnormRopeCacheTiling { + int32_t B; + int32_t N; + int32_t S; + int32_t rms_size; + int32_t rope_size; + int32_t hidden_size; + int32_t total; + int32_t block_size; + int32_t usedCoreNum; + int32_t tasksPerCore; + float eps; + float invRmsSize; + int32_t cache_mode; // 0=Norm, 1=PA, 2=PA_BNSD, 3=PA_NZ, 4=PA_BLK_BNSD, 5=PA_BLK_NZ + int32_t is_output_kv; + int32_t k_cache_dim0; + int32_t k_cache_dim1; + int32_t k_cache_dim2; + int32_t k_cache_dim3; + int32_t ckv_cache_dim0; + int32_t ckv_cache_dim1; + int32_t ckv_cache_dim2; + int32_t ckv_cache_dim3; + int32_t index_numel; + int32_t quant_enabled; // 0=no quant, 1=quant +}; + +#endif diff --git a/archive_tasks/kv_rmsnorm_rope_cache/kernel/pybind11.cpp b/archive_tasks/kv_rmsnorm_rope_cache/kernel/pybind11.cpp new file mode 100644 index 00000000..fd41ff33 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/kernel/pybind11.cpp @@ -0,0 +1,170 @@ +#include +#include + +#include +#include + +#include "acl/acl.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" + +#include "kv_rmsnorm_rope_cache_tiling.h" + +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); + +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); + +namespace kv_rmsnorm_rope_cache_ext { + +using LaunchFn = void (*)(uint32_t, void *, + uint8_t *, uint8_t *, uint8_t *, uint8_t *, + uint8_t *, uint8_t *, uint8_t *, + uint8_t *, uint8_t *, uint8_t *, uint8_t *, + uint8_t *); + +int GetCacheModeId(const std::string &mode) +{ + if (mode == "Norm") return 0; + if (mode == "PA") return 1; + if (mode == "PA_BNSD") return 2; + if (mode == "PA_NZ") return 3; + if (mode == "PA_BLK_BNSD") return 4; + if (mode == "PA_BLK_NZ") return 5; + return 0; +} + +pybind11::tuple run_kv_rmsnorm_rope_cache( + const at::Tensor &kv, const at::Tensor &gamma, + const at::Tensor &cos, const at::Tensor &sin, + const at::Tensor &index, const at::Tensor &k_cache, const at::Tensor &ckv_cache, + double epsilon, const std::string &cache_mode, bool is_output_kv) +{ + TORCH_CHECK(kv.dim() == 4, "kv must be 4D [B, N, S, hidden_size]"); + TORCH_CHECK(gamma.dim() == 1, "gamma must be 1D"); + TORCH_CHECK(cos.dim() == 4, "cos must be 4D"); + TORCH_CHECK(sin.dim() == 4, "sin must be 4D"); + TORCH_CHECK(index.dtype() == at::kLong, "index must be int64"); + TORCH_CHECK(k_cache.dim() == 4, "k_cache must be 4D"); + TORCH_CHECK(ckv_cache.dim() == 4, "ckv_cache must be 4D"); + TORCH_CHECK(kv.is_contiguous(), "kv must be contiguous"); + TORCH_CHECK(gamma.is_contiguous(), "gamma must be contiguous"); + TORCH_CHECK(cos.is_contiguous(), "cos must be contiguous"); + TORCH_CHECK(sin.is_contiguous(), "sin must be contiguous"); + TORCH_CHECK(index.is_contiguous(), "index must be contiguous"); + TORCH_CHECK(k_cache.is_contiguous(), "k_cache must be contiguous"); + TORCH_CHECK(ckv_cache.is_contiguous(), "ckv_cache must be contiguous"); + + const int32_t B = static_cast(kv.sizes()[0]); + const int32_t N = static_cast(kv.sizes()[1]); + const int32_t S = static_cast(kv.sizes()[2]); + const int32_t hidden_size = static_cast(kv.sizes()[3]); + const int32_t rms_size = static_cast(gamma.sizes()[0]); + const int32_t rope_size = hidden_size - rms_size; + const int32_t total = B * S * N; + + // Rearrange kv, cos, sin from BNSD to BSND on host + at::Tensor kv_bsnd = kv.permute({0, 2, 1, 3}).contiguous(); + at::Tensor cos_bsnd = cos.permute({0, 2, 1, 3}).contiguous(); + at::Tensor sin_bsnd = sin.permute({0, 2, 1, 3}).contiguous(); + + // Extract rms_in and rope_in, then interleave rope_in + at::Tensor rms_in = kv_bsnd.narrow(3, 0, rms_size).contiguous(); + at::Tensor rope_in = kv_bsnd.narrow(3, rms_size, rope_size).contiguous(); + at::Tensor k_input = rope_in.reshape({B, S, N, rope_size / 2, 2}) + .permute({0, 1, 2, 4, 3}) + .reshape({B, S, N, rope_size}) + .contiguous(); + + at::Tensor kv_input = at::cat({rms_in.reshape({total, rms_size}), + k_input.reshape({total, rope_size})}, 1).contiguous(); + + at::Tensor k_cache_out = k_cache.clone(); + at::Tensor ckv_cache_out = ckv_cache.clone(); + at::Tensor k_embed_out = at::empty({B, N, S, rope_size}, kv.options()); + at::Tensor v_out = at::empty({B, N, S, rms_size}, kv.options()); + + const int32_t mNum = (total + DEFAULT_BLOCK_SIZE - 1) / DEFAULT_BLOCK_SIZE; + const int32_t usedCoreNum = std::min(DEFAULT_NUM_PHYSICAL_CORES, mNum); + const int32_t tasksPerCore = (mNum + usedCoreNum - 1) / usedCoreNum; + + at::Tensor tilingCpu = at::empty( + {static_cast(sizeof(KvRmsnormRopeCacheTiling))}, + at::device(at::kCPU).dtype(at::kByte)); + auto *tiling = reinterpret_cast(tilingCpu.data_ptr()); + tiling->B = B; + tiling->N = N; + tiling->S = S; + tiling->rms_size = rms_size; + tiling->rope_size = rope_size; + tiling->hidden_size = hidden_size; + tiling->total = total; + tiling->block_size = DEFAULT_BLOCK_SIZE; + tiling->usedCoreNum = usedCoreNum; + tiling->tasksPerCore = tasksPerCore; + tiling->eps = static_cast(epsilon); + tiling->invRmsSize = 1.0f / static_cast(rms_size); + tiling->cache_mode = GetCacheModeId(cache_mode); + tiling->is_output_kv = is_output_kv ? 1 : 0; + tiling->k_cache_dim0 = static_cast(k_cache.sizes()[0]); + tiling->k_cache_dim1 = static_cast(k_cache.sizes()[1]); + tiling->k_cache_dim2 = static_cast(k_cache.sizes()[2]); + tiling->k_cache_dim3 = static_cast(k_cache.sizes()[3]); + tiling->ckv_cache_dim0 = static_cast(ckv_cache.sizes()[0]); + tiling->ckv_cache_dim1 = static_cast(ckv_cache.sizes()[1]); + tiling->ckv_cache_dim2 = static_cast(ckv_cache.sizes()[2]); + tiling->ckv_cache_dim3 = static_cast(ckv_cache.sizes()[3]); + tiling->index_numel = static_cast(index.numel()); + tiling->quant_enabled = 0; + + auto tilingNpu = tilingCpu.to(at::kPrivateUse1); + + auto aclStream = c10_npu::getCurrentNPUStream().stream(false); + LaunchFn launch = nullptr; + if (kv.scalar_type() == at::kHalf) { + launch = kv_rmsnorm_rope_cache_do_fp16; + } else if (kv.scalar_type() == at::kBFloat16) { + launch = kv_rmsnorm_rope_cache_do_bf16; + } else { + TORCH_CHECK(false, "unsupported dtype, only float16 and bfloat16 are supported"); + } + + launch( + usedCoreNum, + aclStream, + static_cast(kv_input.data_ptr()), + static_cast(gamma.data_ptr()), + static_cast(cos_bsnd.data_ptr()), + static_cast(sin_bsnd.data_ptr()), + static_cast(index.data_ptr()), + static_cast(k_cache.data_ptr()), + static_cast(ckv_cache.data_ptr()), + static_cast(k_cache_out.data_ptr()), + static_cast(ckv_cache_out.data_ptr()), + static_cast(k_embed_out.data_ptr()), + static_cast(v_out.data_ptr()), + static_cast(tilingNpu.data_ptr())); + + if (is_output_kv) { + return pybind11::make_tuple(k_cache_out, ckv_cache_out, k_embed_out, v_out); + } + return pybind11::make_tuple(k_cache_out, ckv_cache_out, pybind11::none(), pybind11::none()); +} + +} // namespace kv_rmsnorm_rope_cache_ext + +PYBIND11_MODULE(_kv_rmsnorm_rope_cache_ext, m) +{ + m.doc() = "kv_rmsnorm_rope_cache extension"; + m.def("run_kv_rmsnorm_rope_cache", &kv_rmsnorm_rope_cache_ext::run_kv_rmsnorm_rope_cache, ""); +} diff --git a/archive_tasks/kv_rmsnorm_rope_cache/model.py b/archive_tasks/kv_rmsnorm_rope_cache/model.py new file mode 100644 index 00000000..24558d27 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/model.py @@ -0,0 +1,362 @@ +import json +import os +import torch +import torch.nn as nn +import torch_npu + +class Model(nn.Module): + """ + Simple model that performs KV RMSNorm and RoPE with cache operations. + torch_npu.npu_kv_rmsnorm_rope_cache(kv, gamma, cos, sin, index, k_cache, ckv_cache, *, k_rope_scale=None, c_kv_scale=None, k_rope_offset=None, c_kv_offset=None, epsilon=1e-5, cache_mode='Norm', is_output_kv=False) -> (Tensor, Tensor, Tensor, Tensor) + PyTorch native implementation of forward function + def forward(self, kv: torch.Tensor, gamma: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, + index: torch.Tensor, k_cache: torch.Tensor, ckv_cache: torch.Tensor, + k_rope_scale: torch.Tensor = None, c_kv_scale: torch.Tensor = None, + k_rope_offset: torch.Tensor = None, c_kv_offset: torch.Tensor = None, + epsilon: float = 1e-5, cache_mode: str = 'Norm', is_output_kv: bool = False) -> tuple: + + B, N, S, hidden_size = kv.shape + rms_size = gamma.shape[0] + rope_size = hidden_size - rms_size + orig_dtype = kv.dtype + + kv = kv.float() + gamma = gamma.float() + cos = cos.float() + sin = sin.float() + + kv = rearrange(kv, 'b n s d -> b s n d') + cos = rearrange(cos, 'b n s d -> b s n d') + sin = rearrange(sin, 'b n s d -> b s n d') + + rms_in = kv[..., :rms_size] + rope_in = kv[..., rms_size:] + + rms_mean = torch.mean(rms_in ** 2, dim=-1, keepdim=True) + rms_normalized = rms_in / torch.sqrt(rms_mean + epsilon) + v = gamma * rms_normalized + + k = rope_in.reshape(B, S, N, rope_size // 2, 2).transpose(-1, -2).reshape(B, S, N, rope_size) + k1 = k[..., :rope_size // 2] + k2 = k[..., rope_size // 2:] + rotate_half_k = torch.cat((-k2, k1), dim=-1) + + if cos.shape[1] == 1 and S > 1: + cos = cos.expand(B, S, N, rope_size) + if sin.shape[1] == 1 and S > 1: + sin = sin.expand(B, S, N, rope_size) + + k_embed = k * cos + rotate_half_k * sin + + v_out = rearrange(v, 'b s n d -> b n s d').to(orig_dtype) + k_embed_out = rearrange(k_embed, 'b s n d -> b n s d').to(orig_dtype) + + v_for_cache = v.clone() + k_for_cache = k_embed.clone() + + if c_kv_scale is not None: + v_for_cache = v_for_cache * c_kv_scale.float() + if c_kv_offset is not None: + v_for_cache = v_for_cache + c_kv_offset.float() + if c_kv_scale is not None: + v_for_cache = torch.round(v_for_cache).clamp(-128, 127) + + if k_rope_scale is not None: + k_for_cache = k_for_cache * k_rope_scale.float() + if k_rope_offset is not None: + k_for_cache = k_for_cache + k_rope_offset.float() + if k_rope_scale is not None: + k_for_cache = torch.round(k_for_cache).clamp(-128, 127) + + # Cache operations + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + + if cache_mode == 'Norm': + if index.dim() == 2: + for b in range(min(B, index.shape[0])): + for s in range(min(S, index.shape[1])): + idx = index[b, s].item() + if idx < 0: + continue + if idx < k_cache_out.shape[2]: + k_cache_out[b, :, idx, :] = k_for_cache[b, s, :, :].to(k_cache_out.dtype) + ckv_cache_out[b, :, idx, :] = v_for_cache[b, s, :, :].to(ckv_cache_out.dtype) + elif index.dim() == 1: + for i, idx_t in enumerate(index): + idx = idx_t.item() + if idx < 0: + continue + b_idx = i // S + s_idx = i % S + if b_idx < B and s_idx < S and idx < k_cache_out.shape[2]: + k_cache_out[b_idx, :, idx, :] = k_for_cache[b_idx, s_idx, :, :].to(k_cache_out.dtype) + ckv_cache_out[b_idx, :, idx, :] = v_for_cache[b_idx, s_idx, :, :].to(ckv_cache_out.dtype) + + elif cache_mode in ['PA', 'PA_BNSD']: + if index.dim() == 1: + block_size = k_cache_out.shape[1] + k_flat = k_for_cache.reshape(B * S, N, -1) + v_flat = v_for_cache.reshape(B * S, N, -1) + cache_shape_k = k_cache_out.shape + cache_shape_v = ckv_cache_out.shape + k_cache_flat = k_cache_out.reshape(-1, N, cache_shape_k[-1]) + v_cache_flat = ckv_cache_out.reshape(-1, N, cache_shape_v[-1]) + + for i in range(min(len(index), B * S)): + idx = index[i].item() + if idx < 0: + continue + if idx < k_cache_flat.shape[0]: + k_cache_flat[idx, :, :] = k_flat[i, :, :].to(k_cache_flat.dtype) + v_cache_flat[idx, :, :] = v_flat[i, :, :].to(v_cache_flat.dtype) + + k_cache_out = k_cache_flat.reshape(cache_shape_k) + ckv_cache_out = v_cache_flat.reshape(cache_shape_v) + + elif cache_mode == 'PA_NZ': + if index.dim() == 1: + block_size = k_cache_out.shape[1] + dk = k_cache_out.shape[-1] + dv = ckv_cache_out.shape[-1] + dk0 = 32 if k_cache_out.dtype == torch.int8 else 16 + dv0 = 32 if ckv_cache_out.dtype == torch.int8 else 16 + dk1 = dk // dk0 + dv1 = dv // dv0 + bn = k_cache_out.shape[0] + num_head = k_cache_out.shape[2] + + k_cache_nz = k_cache_out.reshape(bn, num_head, dk1, block_size, dk0) + v_cache_nz = ckv_cache_out.reshape(bn, num_head, dv1, block_size, dv0) + + k_flat = k_for_cache.reshape(B * S, N, -1) + v_flat = v_for_cache.reshape(B * S, N, -1) + + for i in range(min(len(index), B * S)): + idx = index[i].item() + if idx < 0: + continue + bn_id = idx // block_size + block_offset = idx % block_size + if bn_id < bn: + for d in range(dk1): + k_cache_nz[bn_id, :, d, block_offset, :] = k_flat[i, :, d*dk0:(d+1)*dk0].to(k_cache_nz.dtype) + for d in range(dv1): + v_cache_nz[bn_id, :, d, block_offset, :] = v_flat[i, :, d*dv0:(d+1)*dv0].to(v_cache_nz.dtype) + + k_cache_out = k_cache_nz.reshape(k_cache_out.shape) + ckv_cache_out = v_cache_nz.reshape(ckv_cache_out.shape) + + elif cache_mode == 'PA_BLK_BNSD': + if index.dim() == 1: + block_size = k_cache_out.shape[1] + ceil_div_s = (S + block_size - 1) // block_size + + for batch in range(B): + for seq_id in range(ceil_div_s): + seq_start = seq_id * block_size + seq_end = S if seq_id == (ceil_div_s - 1) else (seq_id + 1) * block_size + copy_len = seq_end - seq_start + idx_pos = batch * ceil_div_s + seq_id + if idx_pos >= len(index): + continue + idx_val = index[idx_pos].item() + if idx_val < 0: + continue + cache_b = idx_val // block_size + if cache_b < k_cache_out.shape[0]: + k_cache_out[cache_b, :copy_len, :, :] = k_for_cache[batch, seq_start:seq_end, :, :].to(k_cache_out.dtype) + ckv_cache_out[cache_b, :copy_len, :, :] = v_for_cache[batch, seq_start:seq_end, :, :].to(ckv_cache_out.dtype) + + elif cache_mode == 'PA_BLK_NZ': + if index.dim() == 1: + block_size = k_cache_out.shape[1] + dk = k_cache_out.shape[-1] + dv = ckv_cache_out.shape[-1] + dk0 = 32 if k_cache_out.dtype == torch.int8 else 16 + dv0 = 32 if ckv_cache_out.dtype == torch.int8 else 16 + dk1 = dk // dk0 + dv1 = dv // dv0 + bn = k_cache_out.shape[0] + num_head = k_cache_out.shape[2] + ceil_div_s = (S + block_size - 1) // block_size + + k_cache_nz = k_cache_out.reshape(bn, num_head, dk1, block_size, dk0) + v_cache_nz = ckv_cache_out.reshape(bn, num_head, dv1, block_size, dv0) + + for batch in range(B): + for seq_id in range(ceil_div_s): + seq_start = seq_id * block_size + seq_end = S if seq_id == (ceil_div_s - 1) else (seq_id + 1) * block_size + copy_len = seq_end - seq_start + idx_pos = batch * ceil_div_s + seq_id + if idx_pos >= len(index): + continue + idx_val = index[idx_pos].item() + if idx_val < 0: + continue + cache_b = idx_val // block_size + if cache_b < bn: + for n_idx in range(num_head): + for d in range(dk1): + k_cache_nz[cache_b, n_idx, d, :copy_len, :] = k_for_cache[batch, seq_start:seq_end, n_idx, d*dk0:(d+1)*dk0].to(k_cache_nz.dtype) + for d in range(dv1): + v_cache_nz[cache_b, n_idx, d, :copy_len, :] = v_for_cache[batch, seq_start:seq_end, n_idx, d*dv0:(d+1)*dv0].to(v_cache_nz.dtype) + + k_cache_out = k_cache_nz.reshape(k_cache_out.shape) + ckv_cache_out = v_cache_nz.reshape(ckv_cache_out.shape) + + if is_output_kv: + k_embed_ret = k_embed_out + y_ret = v_out + else: + k_embed_ret = None + y_ret = None + + return k_cache_out, ckv_cache_out, k_embed_ret, y_ret + """ + def __init__(self): + super(Model, self).__init__() + + def postprocess_output(self, output, inputs): + """ + KV RMSNorm RoPE Cache 专用输出裁剪 + 规则：Norm模式 或 is_output_kv=False → 只校验前两个输出 + """ + # 输入顺序与 forward 完全一致 + # kv, gamma, cos, sin, index, k_cache, ckv_cache, + # k_rope_scale, c_kv_scale, k_rope_offset, c_kv_offset, + # epsilon, cache_mode, is_output_kv + if len(inputs) >= 14: + cache_mode = inputs[12] + is_output_kv = inputs[13] + + if cache_mode == 'Norm' or not is_output_kv: + return output[:2] + + return output + + def forward(self, kv: torch.Tensor, gamma: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, + index: torch.Tensor, k_cache: torch.Tensor, ckv_cache: torch.Tensor, + k_rope_scale: torch.Tensor = None, c_kv_scale: torch.Tensor = None, + k_rope_offset: torch.Tensor = None, c_kv_offset: torch.Tensor = None, + epsilon: float = 1e-5, cache_mode: str = 'Norm', is_output_kv: bool = False) -> tuple: + """ + Performs KV RMSNorm and RoPE with cache operations. + + Args: + kv (torch.Tensor): Input feature tensor. Must be 4D [batch_size, 1, seq_len, hidden_size]. + hidden_size = rms_size + rope_size. + dtype: bfloat16, float16, format: BNSD. + gamma (torch.Tensor): RMS normalization scale parameter. Must be 1D [rms_size]. + dtype: bfloat16, float16, format: ND. + cos (torch.Tensor): RoPE cosine component. Must be 4D [batch_size, 1, seq_len, rope_size]. + dtype: bfloat16, float16, format: ND. + sin (torch.Tensor): RoPE sine component. Must be 4D [batch_size, 1, seq_len, rope_size]. + dtype: bfloat16, float16, format: ND. + index (torch.Tensor): Cache index tensor for locating write positions in caches. + dtype: int64, format: ND. Shape depends on cache_mode. + k_cache (torch.Tensor): Storage for quantized/non-quantized key vectors. + dtype: bfloat16, float16, int8, format: ND. Shape depends on cache_mode. + ckv_cache (torch.Tensor): Storage for quantized/non-quantized compressed KV vectors. + dtype: bfloat16, float16, int8, format: ND. Shape depends on cache_mode. + k_rope_scale (torch.Tensor, optional): K RoPE quantization scale. Must be 1D [rope_size]. + dtype: float32, format: ND. Required in quantization mode. + c_kv_scale (torch.Tensor, optional): Compressed KV quantization scale. Must be 1D [rms_size]. + dtype: float32, format: ND. Required in quantization mode. + k_rope_offset (torch.Tensor, optional): K RoPE quantization offset. Must be 1D [rope_size]. + dtype: float32, format: ND. Required in quantization mode. + c_kv_offset (torch.Tensor, optional): Compressed KV quantization offset. Must be 1D [rms_size]. + dtype: float32, format: ND. Required in quantization mode. + epsilon (float, optional): Small value for RMS normalization to prevent division by zero. + Default: 1e-5. + cache_mode (str, optional): Cache mode. Options: 'Norm', 'PA', 'PA_BNSD', 'PA_NZ', + 'PA_BLK_BNSD', 'PA_BLK_NZ'. Default: 'Norm'. + is_output_kv (bool, optional): Whether to output processed k_embed_out and y_out. + Default: False. Only effective in PA modes. + + Returns: + tuple: (k_cache, ckv_cache, k_embed_out, y_out) tensors. Last two are None if is_output_kv=False. + """ + return torch_npu.npu_kv_rmsnorm_rope_cache(kv, gamma, cos, sin, index, k_cache, ckv_cache, + k_rope_scale=k_rope_scale, c_kv_scale=c_kv_scale, + k_rope_offset=k_rope_offset, c_kv_offset=c_kv_offset, + epsilon=epsilon, cache_mode=cache_mode, + is_output_kv=is_output_kv) + + +def get_input_groups(): + """Generate input groups from JSON test cases.""" + json_path = os.path.join(os.path.dirname(__file__), os.path.splitext(os.path.basename(__file__))[0] + '.json') + input_groups = [] + with open(json_path, 'r') as f: + for line in f: + line = line.strip() + if not line: + continue + case = json.loads(line) + inputs = case['inputs'] + tensors = {} + attrs = {} + for inp in inputs: + if inp['type'] == 'tensor': + name = inp['name'] + dtype_str = inp.get('dtype', 'float32') + shape = inp.get('shape') + if shape is None: + tensors[name] = None + elif dtype_str == 'bool': + tensors[name] = (torch.rand(shape) > 0.5).to(torch.bool) + elif dtype_str in ('int32', 'int64', 'int8'): + max_val = {'int32': 1000, 'int64': 10000, 'int8': 127}.get(dtype_str, 100) + dtype = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, 'int32': torch.int32, 'int64': torch.int64, 'int8': torch.int8, 'bool': torch.bool}[dtype_str] + tensors[name] = torch.randint(0, max_val, shape, dtype=dtype) + else: + dtype = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, 'int32': torch.int32, 'int64': torch.int64, 'int8': torch.int8, 'bool': torch.bool}.get(dtype_str, torch.float32) + tensors[name] = torch.randn(shape, dtype=dtype) + elif inp['type'] == 'attr': + attrs[inp['name']] = inp['value'] + + # Generate valid index values based on cache_mode and cache shapes + index = tensors.get('index') + cache_mode = attrs.get('cache_mode', 'Norm') + k_cache = tensors.get('k_cache') + + if index is not None and k_cache is not None: + device = index.device + if cache_mode == 'Norm': + # Norm：全局唯一索引 mod 最大序列长度，避免重复 + max_seq = k_cache.shape[2] + if index.dim() == 2: + B, S = index.shape + total = B * S + index = (torch.arange(total, dtype=torch.int64, device=device) % max_seq).reshape(B, S) + else: + total = index.numel() + index = (torch.arange(total, dtype=torch.int64, device=device) % max_seq).reshape(index.shape) + + elif cache_mode in ('PA', 'PA_BNSD', 'PA_NZ'): + # PA 系列：直接生成连续唯一索引 + index = torch.arange(index.numel(), dtype=torch.int64, device=device) + + elif cache_mode in ('PA_BLK_BNSD', 'PA_BLK_NZ'): + # 分块 PA：索引 = 连续序号 × block_size + block_size = k_cache.shape[1] + length = index.numel() + index = torch.arange(length, dtype=torch.int64, device=device) * block_size + + tensors['index'] = index + group = [ + tensors['kv'], tensors['gamma'], tensors['cos'], tensors['sin'], + tensors['index'], tensors['k_cache'], tensors['ckv_cache'], + tensors.get('k_rope_scale'), tensors.get('c_kv_scale'), + tensors.get('k_rope_offset'), tensors.get('c_kv_offset'), + attrs.get('epsilon', 1e-5), attrs.get('cache_mode', 'Norm'), + attrs.get('is_output_kv', False) + ] + input_groups.append(group) + return input_groups + + +def get_init_inputs(): + return [] diff --git a/archive_tasks/kv_rmsnorm_rope_cache/model_new_ascendc.py b/archive_tasks/kv_rmsnorm_rope_cache/model_new_ascendc.py new file mode 100644 index 00000000..53f8bf16 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/model_new_ascendc.py @@ -0,0 +1,24 @@ +import torch +import torch.nn as nn +import _kv_rmsnorm_rope_cache_ext + + +class ModelNew(nn.Module): + def __init__(self): + super(ModelNew, self).__init__() + + def postprocess_output(self, output, inputs): + if len(inputs) >= 14: + cache_mode = inputs[12] + is_output_kv = inputs[13] + if cache_mode == 'Norm' or not is_output_kv: + return output[:2] + return output + + def forward(self, kv, gamma, cos, sin, index, k_cache, ckv_cache, + k_rope_scale=None, c_kv_scale=None, + k_rope_offset=None, c_kv_offset=None, + epsilon=1e-5, cache_mode='Norm', is_output_kv=False): + return _kv_rmsnorm_rope_cache_ext.run_kv_rmsnorm_rope_cache( + kv, gamma, cos, sin, index, k_cache, ckv_cache, + epsilon, cache_mode, is_output_kv) diff --git a/archive_tasks/kv_rmsnorm_rope_cache/model_new_tilelang.py b/archive_tasks/kv_rmsnorm_rope_cache/model_new_tilelang.py new file mode 100644 index 00000000..23775dc8 --- /dev/null +++ b/archive_tasks/kv_rmsnorm_rope_cache/model_new_tilelang.py @@ -0,0 +1,216 @@ +import torch +from design.tile_level.kv_rmsnorm_rope import kv_rmsnorm_rope + + +def _update_cache_norm(k_cache, ckv_cache, k_embed, v, index): + """Norm mode cache update.""" + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + B, N, S = k_embed.shape[0], k_embed.shape[1], k_embed.shape[2] + if index.dim() == 2: + for b in range(B): + for s in range(S): + idx = index[b, s].item() + if idx < 0: + continue + if idx < k_cache_out.shape[2]: + k_cache_out[b, :, idx, :] = k_embed[b, :, s, :] + ckv_cache_out[b, :, idx, :] = v[b, :, s, :] + else: + for i in range(index.numel()): + idx = index[i].item() + if idx < 0: + continue + b_idx = i // S + s_idx = i % S + if b_idx < B and s_idx < S and idx < k_cache_out.shape[2]: + k_cache_out[b_idx, :, idx, :] = k_embed[b_idx, :, s_idx, :] + ckv_cache_out[b_idx, :, idx, :] = v[b_idx, :, s_idx, :] + return k_cache_out, ckv_cache_out + + +def _update_cache_pa(k_cache, ckv_cache, k_embed, v, index): + """PA / PA_BNSD mode cache update.""" + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + B, N, S = k_embed.shape[0], k_embed.shape[1], k_embed.shape[2] + k_flat = k_embed.reshape(B * S, N, -1) + v_flat = v.reshape(B * S, N, -1) + cache_shape_k = k_cache_out.shape + cache_shape_v = ckv_cache_out.shape + k_cache_flat = k_cache_out.reshape(-1, N, cache_shape_k[-1]) + v_cache_flat = ckv_cache_out.reshape(-1, N, cache_shape_v[-1]) + for i in range(min(len(index), B * S)): + idx = index[i].item() + if idx < 0: + continue + if idx < k_cache_flat.shape[0]: + k_cache_flat[idx, :, :] = k_flat[i, :, :] + v_cache_flat[idx, :, :] = v_flat[i, :, :] + return k_cache_flat.reshape(cache_shape_k), v_cache_flat.reshape(cache_shape_v) + + +def _update_cache_pa_nz(k_cache, ckv_cache, k_embed, v, index): + """PA_NZ mode cache update.""" + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + B, N, S = k_embed.shape[0], k_embed.shape[1], k_embed.shape[2] + block_size = k_cache_out.shape[1] + dk = k_cache_out.shape[-1] + dv = ckv_cache_out.shape[-1] + dk0 = 32 if k_cache_out.dtype == torch.int8 else 16 + dv0 = 32 if ckv_cache_out.dtype == torch.int8 else 16 + dk1 = dk // dk0 + dv1 = dv // dv0 + bn = k_cache_out.shape[0] + num_head = k_cache_out.shape[2] + k_cache_nz = k_cache_out.reshape(bn, num_head, dk1, block_size, dk0) + v_cache_nz = ckv_cache_out.reshape(bn, num_head, dv1, block_size, dv0) + k_flat = k_embed.reshape(B * S, N, -1) + v_flat = v.reshape(B * S, N, -1) + for i in range(min(len(index), B * S)): + idx = index[i].item() + if idx < 0: + continue + bn_id = idx // block_size + block_offset = idx % block_size + if bn_id < bn: + for d in range(dk1): + k_cache_nz[bn_id, :, d, block_offset, :] = k_flat[i, :, d * dk0:(d + 1) * dk0] + for d in range(dv1): + v_cache_nz[bn_id, :, d, block_offset, :] = v_flat[i, :, d * dv0:(d + 1) * dv0] + return k_cache_nz.reshape(k_cache_out.shape), v_cache_nz.reshape(ckv_cache_out.shape) + + +def _update_cache_pa_blk_bnsd(k_cache, ckv_cache, k_embed, v, index): + """PA_BLK_BNSD mode cache update.""" + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + B, N, S = k_embed.shape[0], k_embed.shape[1], k_embed.shape[2] + block_size = k_cache_out.shape[1] + ceil_div_s = (S + block_size - 1) // block_size + for batch in range(B): + for seq_id in range(ceil_div_s): + seq_start = seq_id * block_size + seq_end = S if seq_id == (ceil_div_s - 1) else (seq_id + 1) * block_size + copy_len = seq_end - seq_start + idx_pos = batch * ceil_div_s + seq_id + if idx_pos >= len(index): + continue + idx_val = index[idx_pos].item() + if idx_val < 0: + continue + cache_b = idx_val // block_size + if cache_b < k_cache_out.shape[0]: + k_cache_out[cache_b, :copy_len, :, :] = k_embed[batch, seq_start:seq_end, :, :] + ckv_cache_out[cache_b, :copy_len, :, :] = v[batch, seq_start:seq_end, :, :] + return k_cache_out, ckv_cache_out + + +def _update_cache_pa_blk_nz(k_cache, ckv_cache, k_embed, v, index): + """PA_BLK_NZ mode cache update.""" + k_cache_out = k_cache.clone() + ckv_cache_out = ckv_cache.clone() + B, N, S = k_embed.shape[0], k_embed.shape[1], k_embed.shape[2] + block_size = k_cache_out.shape[1] + dk = k_cache_out.shape[-1] + dv = ckv_cache_out.shape[-1] + dk0 = 32 if k_cache_out.dtype == torch.int8 else 16 + dv0 = 32 if ckv_cache_out.dtype == torch.int8 else 16 + dk1 = dk // dk0 + dv1 = dv // dv0 + bn = k_cache_out.shape[0] + num_head = k_cache_out.shape[2] + ceil_div_s = (S + block_size - 1) // block_size + k_cache_nz = k_cache_out.reshape(bn, num_head, dk1, block_size, dk0) + v_cache_nz = ckv_cache_out.reshape(bn, num_head, dv1, block_size, dv0) + for batch in range(B): + for seq_id in range(ceil_div_s): + seq_start = seq_id * block_size + seq_end = S if seq_id == (ceil_div_s - 1) else (seq_id + 1) * block_size + copy_len = seq_end - seq_start + idx_pos = batch * ceil_div_s + seq_id + if idx_pos >= len(index): + continue + idx_val = index[idx_pos].item() + if idx_val < 0: + continue + cache_b = idx_val // block_size + if cache_b < bn: + for n_idx in range(num_head): + for d in range(dk1): + k_cache_nz[cache_b, n_idx, d, :copy_len, :] = k_embed[batch, seq_start:seq_end, n_idx, d * dk0:(d + 1) * dk0] + for d in range(dv1): + v_cache_nz[cache_b, n_idx, d, :copy_len, :] = v[batch, seq_start:seq_end, n_idx, d * dv0:(d + 1) * dv0] + return k_cache_nz.reshape(k_cache_out.shape), v_cache_nz.reshape(ckv_cache_out.shape) + + +class ModelNew(torch.nn.Module): + def __init__(self): + super(ModelNew, self).__init__() + + def postprocess_output(self, output, inputs): + if len(inputs) >= 14: + cache_mode = inputs[12] + is_output_kv = inputs[13] + if cache_mode == 'Norm' or not is_output_kv: + return output[:2] + return output + + def forward(self, kv, gamma, cos, sin, index, k_cache, ckv_cache, + k_rope_scale=None, c_kv_scale=None, + k_rope_offset=None, c_kv_offset=None, + epsilon=1e-5, cache_mode='Norm', is_output_kv=False): + B, N, S, hidden_size = kv.shape + rms_size = gamma.shape[0] + rope_size = hidden_size - rms_size + + # Rearrange BNSD -> BSND + kv_bsnd = kv.permute(0, 2, 1, 3) + cos_bsnd = cos.permute(0, 2, 1, 3) + sin_bsnd = sin.permute(0, 2, 1, 3) + + rms_in = kv_bsnd[..., :rms_size] + rope_in = kv_bsnd[..., rms_size:] + + # Preprocess rope_in: interleave pairs (same as reshape+transpose+reshape in reference) + k_input = rope_in.reshape(B, S, N, rope_size // 2, 2).permute(0, 1, 2, 4, 3).reshape(B, S, N, rope_size) + + # Flatten position dimension + total = B * S * N + rms_in_flat = rms_in.reshape(total, rms_size) + k_input_flat = k_input.reshape(total, rope_size) + cos_flat = cos_bsnd.reshape(total, rope_size) + sin_flat = sin_bsnd.reshape(total, rope_size) + + # Build and call TileLang kernel + dtype_str = str(kv.dtype).split('.')[-1] + kernel = kv_rmsnorm_rope(total, rms_size, rope_size, eps=epsilon, dtype=dtype_str) + v_flat, k_embed_flat = kernel(rms_in_flat, gamma, k_input_flat, cos_flat, sin_flat) + + # Reshape outputs back to BNSD + v = v_flat.reshape(B, S, N, rms_size).permute(0, 2, 1, 3) + k_embed = k_embed_flat.reshape(B, S, N, rope_size).permute(0, 2, 1, 3) + + # Cache update (dispatch by cache_mode) + if cache_mode == 'Norm': + k_cache_out, ckv_cache_out = _update_cache_norm(k_cache, ckv_cache, k_embed, v, index) + elif cache_mode in ('PA', 'PA_BNSD'): + k_cache_out, ckv_cache_out = _update_cache_pa(k_cache, ckv_cache, k_embed, v, index) + elif cache_mode == 'PA_NZ': + k_cache_out, ckv_cache_out = _update_cache_pa_nz(k_cache, ckv_cache, k_embed, v, index) + elif cache_mode == 'PA_BLK_BNSD': + k_cache_out, ckv_cache_out = _update_cache_pa_blk_bnsd(k_cache, ckv_cache, k_embed, v, index) + elif cache_mode == 'PA_BLK_NZ': + k_cache_out, ckv_cache_out = _update_cache_pa_blk_nz(k_cache, ckv_cache, k_embed, v, index) + else: + raise ValueError(f"Unsupported cache_mode: {cache_mode}") + + if is_output_kv: + k_embed_ret = k_embed + y_ret = v + else: + k_embed_ret = None + y_ret = None + + return k_cache_out, ckv_cache_out, k_embed_ret, y_ret