[Feature]:support kv_rmsnorm_rope_cache. · Pull Request !50 · MindSpore/ms_custom_ops

1. [Feature]:support kv_rmsnorm_rope_cache.
---
<!--  Thanks for sending a pull request!  Here are some tips for you:

1) If this is your first time, please read our contributor guidelines: https://gitee.com/mindspore/mindspore/blob/master/CONTRIBUTING.md

2) If you want to contribute your code but don't know who will review and merge, please add label `mindspore-assistant` to the pull request, we will find and do it as soon as possible.
-->

**What type of PR is this?**

/kind task

---

**What does this PR do / why do we need it**:
support deepseek v3.2
---

**Which issue(s) this PR fixes**:

Fixes #ID2B96

---

**Test Plan and Test result：What scenarios were tested, and what were the verification results（Function, performance, reliability, etc.）**：
 Function: aseert error.
---

**Self-checklist**:（**请自检，在[ ]内打上x，我们将检视你的完成情况，否则会导致pr无法合入**）

+ - [x] **设计**：PR对应的方案是否已经经过Maintainer评审，方案检视意见是否均已答复并完成方案修改
+ - [x] **测试**：PR中的代码是否已有UT/ST测试用例进行充分的覆盖，新增测试用例是否随本PR一并上库或已经上库
+ - [x] **验证**：PR描述信息中是否已包含对该PR对应的Feature、Refactor、Bugfix的预期目标达成情况的详细验证结果描述
+ - [ ] **接口**：是否涉及对外接口变更，相应变更已得到接口评审组织的通过，API对应的注释信息已经刷新正确
+ - [ ] **文档**：是否涉及官网文档修改，如果涉及请及时提交资料到Doc仓

**功能概述**
本次提交主要为MindSpore自定义算子库新增基于Ascend上的kv_rmsnorm_rope_cache算子支持。该算子实现了结合RMSNorm归一化、旋转位置编码（RoPE）和缓存机制的复合操作，主要用于优化Transformer类模型中的注意力机制，提供高效的计算和内存管理。

**主要功能特性**

- 1. 新增kv_rmsnorm_rope_cache算子

问题描述

原有自定义算子库缺少高效的结合归一化、位置编码和缓存功能的复合操作

用户需要实现复杂的注意力机制优化，但缺乏硬件加速支持

缺少针对昇腾硬件优化的RMSNorm和RoPE联合实现

解决方案

新增kv_rmsnorm_rope_cache算子的完整实现

实现完整的算子注册、接口定义和计算逻辑

- 2. 复合功能集成

问题描述

RMSNorm、RoPE和缓存操作通常需要分开执行，导致多次内存读写和计算开销

分开实现难以充分利用硬件并行性

解决方案

将RMSNorm归一化、旋转位置编码和缓存更新集成到单个算子中

优化内存访问模式，减少中间结果存储

实现计算流水线优化，提升整体性能

- 3. 缓存机制优化

问题描述

传统缓存实现需要手动管理索引和内存更新

不同缓存模式（如PA、PA_BLK等）的支持不够灵活

解决方案

提供多种缓存模式支持（PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ）

实现自动索引管理和缓存更新

支持量化和非量化两种缓存格式

**算子注册与接口**
# 算子注册信息
operator_name: "kv_rmsnorm_rope_cache"
| Name                | DType           | Shape                                  | Optional | Inplace | Format | Description                                            |
|---------------------|-----------------|----------------------------------------|----------|---------|--------|--------------------------------------------------------|
| kv               | Tensor          | 4维[B_kv, N, S_kv, D] | No       | No      | ND     | 用于切分出RMSNorm计算所需数据Dv和RoPE计算所需数据Dk的输入数据                              |
| gamma                 | Tensor          | 1维[D_v] | No       | No      | ND     | 用于RMSNorm计算的输入数据                              |  
| cos                 | Tensor          | 4维[Bkv, 1, Skv, Dk] 或 [Bkv, 1, 1, Dk]         | No       | No      | ND     | 用于计算Rope的余弦变换数据                                |
| sin                 | Tensor          | 4维[Bkv, 1, Skv, Dk] 或 [Bkv, 1, 1, Dk]          | No       | No      | ND     | 表用于计算Rope的正弦变换数据                                |
| index              | Tensor          | 多维，参见约束说明                                 | No      | No      | ND | 用于指定写入cache的具体索引位置，value为-1时跳过更新                                 |
| k_cache         | Tensor          |    4维，参见约束说明                                  | No      | Yes      | ND | 提前申请的Cache                              |
| c_kv_cache         | Tensor          | 4维，参见约束说明                                  | No      | Yes      | ND | 提前申请的Cache                              |
| k_rope_scale         | Tensor          | 多维，参见约束说明                               | Yes      | No      | ND | k_cache为int8时量化参数                              |
| c_kv_scale         | Tensor          | 多维，参见约束说明                                   | Yes      | No      | ND | 提c_kv_cache int8时量化参数                              |
| k_rope_offet         | Tensor          | 多维，参见约束说明                                 | Yes      | No      | ND | k_cache为int8时量化参数且非对称量化时的偏移参数                              |
| c_kv_offset_         | Tensor          | 多维，参见约束说明                                 | Yes      | No      | ND | c_kv_cache为int8时量化参数且非对称量化时的偏移参数                              |
| epsilon         | float          | No                                     | No      | No      | - | Rmsnorm进行计算时防除零的epsilon值                              |
| cache_mode         | int          | No                                     | No      | No      | - | 提cache格式选择枚举，枚举含义详见约束说明                              |
| is_output_kv         | bool          | No                                     | No      | No      | - | 控制是否输出k_rope_out和c_kv_out的标志位                              |

output_desc:
| Name   | DType      | Shape      | Description           |
|--------|------------|------------|-----------------------|
| k_rope_out| Tensor   | [B_kv, N, S_kv, D_k] | Rope计算的结果(is_output_kv为True时输出) |
| c_kv_out | Tensor    | [B_kv, N, S_kv, D_v] | Rmsnorm计算的结果(is_output_kv为True时输出) |

支持量化模式：int8量化存储，float32缩放因子

**测试验证**
1. 功能正确性测试
基础功能测试：验证算子在各种输入情况下的输出正确性

边界条件测试：测试索引越界、空张量等特殊情况的处理

数据类型测试：验证不同数据类型（float16/bfloat16/int8）的兼容性

缓存模式测试：测试所有支持的缓存模式（PA、PA_BNSD、PA_NZ、PA_BLK_BNSD、PA_BLK_NZ）

**兼容性说明**
向后兼容性
本次新增算子不会影响现有算子的功能

完全兼容现有MindSpore版本和模型

硬件兼容性
主要针对Ascend硬件优化

在其他硬件平台上提供基础功能支持

API兼容性
接口设计遵循MindSpore现有规范

确保与未来版本的兼容性

使用方式
```python
import mindspore as ms
import numpy as np
import ms_custom_ops

ms.set_device("Ascend")

@ms.jit
def kv_rmsnorm_rope_cache_func(kv,
        gamma,
        cos,
        sin,
        index,
        kCacheRef,
        ckvCacheRef,
        kRopeScale=None,
        ckvScale=None,
        kRopeOffset=None,
        cKvOffset=None,
        epsilon=1e-5,
        cache_mode=0,
        is_output_kv=False,
    ):
    return ms_custom_ops.kv_rmsnorm_rope_cache(
        kv,
        gamma,
        cos,
        sin,
        index,
        kCacheRef,
        ckvCacheRef,
        kRopeScale,
        ckvScale,
        kRopeOffset,
        cKvOffset,
        epsilon,
        cache_mode,
        is_output_kv,
    )

def to_tensor(arr, dtype=None):
    if arr is None:
        return None
    if dtype == np.int8:
        return Tensor(arr, dtype=ms.int8)
    return Tensor(arr, dtype=get_ms_dtype(dtype))

def generate_inputs(
    batch_size,
    seq_len,
    page_num,
    page_size,
    quant_mode,
    cache_mode,
    output_mode,
    input_dtype,
    epsilon,
):
    # 生成基础输入张量
    kv = np.random.randn(batch_size, 1, seq_len, 576).astype(input_dtype)
    gamma = np.random.randn(512).astype(input_dtype)
    cos = np.random.randn(batch_size, 1, seq_len, 64).astype(input_dtype)
    sin = np.random.randn(batch_size, 1, seq_len, 64).astype(input_dtype)

# 初始化缓存相关变量
    k_cache = None
    ckv_cache = None
    index = None
    k_rope_scale = None
    c_kv_scale = None

# 处理缓存模式
    if cache_mode != "Norm":
        # 创建初始缓存（全9张量）
        k_cache = np.ones((page_num, page_size, 1, 64), dtype=input_dtype) * 9
        ckv_cache = np.ones((page_num, page_size, 1, 512), dtype=input_dtype) * 9

# 创建索引数组
        if "BLK" in cache_mode:
            total_blocks = batch_size * ((seq_len + page_size - 1) // page_size)
            index = np.arange(0, total_blocks * page_size, page_size, dtype=np.int64)
        else:
            index = np.arange(0, batch_size * seq_len, 1, dtype=np.int64)

# 处理量化模式
    if quant_mode == 1:
        if k_cache is not None:
            k_cache = k_cache.astype(np.int8)
        if ckv_cache is not None:
            ckv_cache = ckv_cache.astype(np.int8)
        k_rope_scale = np.random.randn(64).astype(np.float32)
        c_kv_scale = np.random.randn(512).astype(np.float32)

# 应用与原始代码相同的变换
    kv = 8 * kv - 10  # (-2 + 10) = 8
    gamma = 990 * gamma - 1000  # (-10 + 1000) = 990
    sin = 0.02 * sin - 0.01  # (0.01 + 0.01) = 0.02
    return (
        kv,
        gamma,
        cos,
        sin,
        index,
        k_cache,
        ckv_cache,
        k_rope_scale,
        c_kv_scale,
        cache_mode,
        output_mode,
        input_dtype,
        epsilon,
    )

def get_kv_rmsnorm_rope_cache_mode_enum(cache_mode_str):
    if cache_mode_str == "PA":
        return 1
    elif cache_mode_str == "PA_BNSD":
        return 2
    elif cache_mode_str == "PA_NZ":
        return 3
    elif cache_mode_str == "PA_BLK_BNSD":
        return 4
    elif cache_mode_str == "PA_BLK_NZ":
        return 5
    return 0  # "Norm"

batch_size = 64
seq_len = 1
page_num = 576
page_size = 128
quant_mode = 0
cache_mode = "PA_BNSD"
output_mode = False
input_dtype = np.float16
epsilon = 1e-5
(kv, gamma, cos, sin, index, k_cache, ckv_cache, k_rope_scale, c_kv_scale, cache_mode, output_mode, _, _,) = generate_inputs(
    batch_size,
    seq_len,
    page_num,
    page_size,
    quant_mode,
    cache_mode,
    output_mode,
    input_dtype,
    epsilon)

kv_tensor = to_tensor(kv, input_dtype)
gamma_tensor = to_tensor(gamma, input_dtype)
cos_tensor = to_tensor(cos, input_dtype)
sin_tensor = to_tensor(sin, input_dtype)
index_tensor = to_tensor(index, np.int64)
k_rope_scale_tensor = None
c_kv_scale_tensor = None
if quant_mode == 1:
    k_cahce_tensor = to_tensor(k_cache, np.int8)
    ckv_cache_tensor = to_tensor(ckv_cache, np.int8)
    k_rope_scale_tensor = to_tensor(k_rope_scale, np.float32)
    c_kv_scale_tensor = to_tensor(c_kv_scale, np.float32)
else:
    k_cache_tensor = to_tensor(k_cache, input_dtype)
    ckv_cache_tensor = to_tensor(ckv_cache, input_dtype)

cache_mode_enum = get_kv_rmsnorm_rope_cache_mode_enum(cache_mode)
k_rope_offset = None
c_kv_offset = None
k_rope, c_kv = kv_rmsnorm_rope_cache_func(
        kv_tensor,
        gamma_tensor,
        cos_tensor,
        sin_tensor,
        index_tensor,
        k_cache_tensor,
        ckv_cache_tensor,
        k_rope_scale_tensor,
        c_kv_scale_tensor,
        k_rope_offset,
        c_kv_offset,
        epsilon,
        cache_mode_enum,
        output_mode,
    )
print(k_cache_tensor)
print(ckv_cache_tensor)
```

**影响范围**
**新增功能**
新增算子：kv_rmsnorm_rope_cache算子及其相关实现

API扩展：新增Python接口和C++后端实现

文档更新：新增算子使用说明和API文档

性能影响
性能提升：为复合操作提供硬件加速支持，显著减少计算时间

内存优化：减少中间结果存储，降低内存使用量

依赖关系
硬件依赖：最佳性能需要Ascend硬件支持

软件依赖：依赖特定版本的MindSpore和CANN库

类型	指派人员	状态
审查	fangwenyi	已完成 (0/0人)
测试	fangwenyi	已完成 (0/0人)

MindSpore/ms_custom_ops .gitee-modal { width: 500px !important; }

搜索帮助

MindSpore/ms_custom_ops