[feature]接入算子aclnnRingAttentionUpdateV2-op-plugin-AtomGit

Mcodegen/templates/_op_plugin_docs.py+56-1

	@@ -2822,6 +2822,61 @@ out, lse_out = torch_npu.npu_attention_update(lse, local_out, update_type)
2822	)	2822	)
2823		2823
2824		2824
		2825	+_add_torch_npu_docstr(
		2826	+ "npu_ring_attention_update",
		2827	+ """
		2828	+接口原型：
		2829	+npu_ring_attention_update(Tensor prev_attn_out, Tensor prev_softmax_max, Tensor prev_softmax_sum, Tensor cur_attn_out, Tensor cur_softmax_max, Tensor cur_softmax_sum, *, Tensor? actual_seq_qlen=None, str input_layout="SBH", str input_softmax_layout="") -> (Tensor, Tensor, Tensor)
		2830	+
		2831	+功能描述
		2832	+将两次FlashAttention的输出按照softmax的max和sum进行更新，输出新的attention结果、softmax_max和softmax_sum。
		2833	+
		2834	+计算公式：
		2835	+softmax_max = max(prev_softmax_max, cur_softmax_max)
		2836	+softmax_sum = prev_softmax_sum * exp(prev_softmax_max - softmax_max) + cur_softmax_sum * exp(cur_softmax_max - softmax_max)
		2837	+attn_out = prev_attn_out * exp(prev_softmax_max - softmax_max) * prev_softmax_sum / softmax_sum
		2838	+ + cur_attn_out * exp(cur_softmax_max - softmax_max) * cur_softmax_sum / softmax_sum
		2839	+
		2840	+参数说明
		2841	+prev_attn_out: Tensor类型, 第一次FlashAttention的输出。数据类型支持FLOAT16、FLOAT、BFLOAT16, 数据格式支持ND。
		2842	+prev_softmax_max: Tensor类型, 第一次FlashAttention的softmax max结果。数据类型支持FLOAT, 数据格式支持ND。
		2843	+prev_softmax_sum: Tensor类型, 第一次FlashAttention的softmax sum结果。数据类型支持FLOAT, 数据格式支持ND。
		2844	+cur_attn_out: Tensor类型, 第二次FlashAttention的输出。数据类型和shape需与prev_attn_out一致。
		2845	+cur_softmax_max: Tensor类型, 第二次FlashAttention的softmax max结果。数据类型和shape需与prev_softmax_max一致。
		2846	+cur_softmax_sum: Tensor类型, 第二次FlashAttention的softmax sum结果。数据类型和shape需与prev_softmax_sum一致。
		2847	+actual_seq_qlen: Tensor类型, 可选参数, TND场景下必选, 表示从0开始累计的query序列长度前缀和。数据类型支持int64, 数据格式支持ND。
		2848	+input_layout: string类型, 可选参数, attention输入输出排布。支持"SBH"和"TND", 默认值为"SBH"。
		2849	+input_softmax_layout: string类型, 可选参数, softmax相关输入排布。支持""、"SBH"、"TND", 默认值为""。仅在input_layout为"TND"时生效。
		2850	+
		2851	+输出说明
		2852	+attn_out: Tensor类型, 更新后的attention输出。shape和数据类型与prev_attn_out一致。
		2853	+softmax_max: Tensor类型, 更新后的softmax max。shape与prev_softmax_max一致, 数据类型为FLOAT。
		2854	+softmax_sum: Tensor类型, 更新后的softmax sum。shape与prev_softmax_sum一致, 数据类型为FLOAT。
		2855	+
		2856	+支持的型号
		2857	+----------------
		2858	+Ascend 950PR/Ascend 950DT
		2859	+Atlas A3训练系列产品/Atlas A3推理系列产品
		2860	+Atlas A2训练系列产品/Atlas A2推理系列产品
		2861	+
		2862	+调用示例
		2863	+----------------
		2864	+import torch
		2865	+import torch_npu
		2866	+
		2867	+prev_attn_out = torch.randn((4, 2, 32), dtype=torch.float16, device="npu")
		2868	+cur_attn_out = torch.randn((4, 2, 32), dtype=torch.float16, device="npu")
		2869	+prev_softmax_max = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
		2870	+prev_softmax_sum = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
		2871	+cur_softmax_max = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
		2872	+cur_softmax_sum = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
		2873	+attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
		2874	+ prev_attn_out, prev_softmax_max, prev_softmax_sum,
		2875	+ cur_attn_out, cur_softmax_max, cur_softmax_sum)
		2876	+"""
		2877	+)
		2878	+
		2879	+
2825	_add_torch_npu_docstr(	2880	_add_torch_npu_docstr(
2826	"npu_linear",	2881	"npu_linear",
2827	"""	2882	"""
	@@ -15678,4 +15733,4 @@ def test_npu_apply_rotary_pos_emb():
15678	if __name__ == "__main__":	15733	if __name__ == "__main__":
15679	test_npu_apply_rotary_pos_emb()	15734	test_npu_apply_rotary_pos_emb()
15680	"""	15735	"""
15681	-)	15736	+)

Mdocs/zh/custom_APIs/menu_Pytorch_API.md+1-0

	@@ -79,6 +79,7 @@
79	- [torch_npu.npu_anti_quant](./torch_npu/torch_npu-npu_anti_quant.md)	79	- [torch_npu.npu_anti_quant](./torch_npu/torch_npu-npu_anti_quant.md)
80	- [torch_npu.npu_attention_to_ffn](./torch_npu/torch_npu-npu_attention_to_ffn.md)	80	- [torch_npu.npu_attention_to_ffn](./torch_npu/torch_npu-npu_attention_to_ffn.md)
81	- [torch_npu.npu_attention_update](./torch_npu/torch_npu-npu_attention_update.md)	81	- [torch_npu.npu_attention_update](./torch_npu/torch_npu-npu_attention_update.md)
		82	+ - [torch_npu.npu_ring_attention_update](./torch_npu/torch_npu-npu_ring_attention_update.md)
82	- [torch_npu.npu_block_sparse_attention](./torch_npu/torch_npu-npu_block_sparse_attention.md)	83	- [torch_npu.npu_block_sparse_attention](./torch_npu/torch_npu-npu_block_sparse_attention.md)
83	- [torch_npu-npu_chunk_gated_delta_rule](./torch_npu/torch_npu-npu_chunk_gated_delta_rule.md)	84	- [torch_npu-npu_chunk_gated_delta_rule](./torch_npu/torch_npu-npu_chunk_gated_delta_rule.md)
84	- [torch_npu.npu_convert_weight_to_int4pack](./torch_npu/torch_npu-npu_convert_weight_to_int4pack.md)	85	- [torch_npu.npu_convert_weight_to_int4pack](./torch_npu/torch_npu-npu_convert_weight_to_int4pack.md)

Adocs/zh/custom_APIs/torch_npu/torch_npu-npu_ring_attention_update.md+129-0

@@ -0,0 +1,129 @@
	1	+# torch_npu.npu_ring_attention_update
	2	+
	3	+## 产品支持情况
	4	+
	5	+\| 产品 \| 是否支持 \|
	6	+\| ------------------------------------------------------------ \| :------: \|
	7	+\|<term>Ascend 950PR/Ascend 950DT</term> \| √ \|
	M M molly1233215月21日 Atlas 350 加速卡
	8	+\|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> \| √ \|
	9	+\|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> \| √ \|
	10	+
	11	+## 功能说明
	12	+
	13	+- API功能：将两次 FlashAttention 的输出结果按照对应的 softmax max 和 softmax sum 做增量更新，得到新的 attention 输出、softmax max 和 softmax sum。
	M M molly1233215月21日中英文间不要空格，全文排查 API的目的是啥，得到新的输出要干啥之类的
	14	+- 计算公式：
	15	+
	16	+ $$
	17	+ softmax\_max = max(prev\_softmax\_max, cur\_softmax\_max)
	18	+ $$
	19	+
	20	+ $$
	21	+ softmax\_sum = prev\_softmax\_sum \times exp(prev\_softmax\_max - softmax\_max) + cur\_softmax\_sum \times exp(cur\_softmax\_max - softmax\_max)
	22	+ $$
	23	+
	24	+ $$
	25	+ attn\_out = prev\_attn\_out \times exp(prev\_softmax\_max - softmax\_max) \times prev\_softmax\_sum / softmax\_sum + cur\_attn\_out \times exp(cur\_softmax\_max - softmax\_max) \times cur\_softmax\_sum / softmax\_sum
	26	+ $$
	27	+
	28	+> [!NOTE]
	29	+>
	30	+> - 该接口底层调用 CANN 的 `aclnnRingAttentionUpdateV2`，支持在 `input_layout="TND"` 场景下通过 `input_softmax_layout` 控制 softmax 相关输入是否采用 `TND` 排布。
	31	+> - `input_layout="TND"` 时，`actual_seq_qlen` 为必选参数。
	M M molly1233215月21日为啥放这里，和约束和参数说明都有重复
	32	+> - `input_softmax_layout` 仅支持 `""`、`"SBH"`、`"TND"` 三种取值。
	33	+
	34	+## 函数原型
	35	+
	36	+```python
	37	+torch_npu.npu_ring_attention_update(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out, cur_softmax_max, cur_softmax_sum, *, actual_seq_qlen=None, input_layout="SBH", input_softmax_layout="") -> (Tensor, Tensor, Tensor)
	38	+```
	39	+
	40	+## 参数说明
	41	+
	42	+- prev_attn_out (`Tensor`)：必选参数，第一次 FlashAttention 的输出。数据类型支持 `float16`、`float32`、`bfloat16`，数据格式支持 $ND$，支持非连续的 Tensor。`input_layout="SBH"` 时 shape 为 `[S, B, H]`，`input_layout="TND"` 时 shape 为 `[T, N, D]`。
	43	+- prev_softmax_max (`Tensor`)：必选参数，第一次 FlashAttention 的 softmax max 结果。数据类型支持 `float32`，数据格式支持 $ND$，支持非连续的 Tensor。`input_layout="SBH"` 时 shape 为 `[B, N, S, 8]`，`input_layout="TND"` 且 `input_softmax_layout="TND"` 时 shape 为 `[T, N, 8]`。最后一维 8 个元素应保持相同且为正数。
	44	+- prev_softmax_sum (`Tensor`)：必选参数，第一次 FlashAttention 的 softmax sum 结果。数据类型支持 `float32`，数据格式支持 $ND$，支持非连续的 Tensor。shape 需要与 `prev_softmax_max` 一致，最后一维 8 个元素应保持相同且为正数。
	45	+- cur_attn_out (`Tensor`)：必选参数，第二次 FlashAttention 的输出。数据类型、数据格式和 shape 需要与 `prev_attn_out` 保持一致。
	46	+- cur_softmax_max (`Tensor`)：必选参数，第二次 FlashAttention 的 softmax max 结果。数据类型支持 `float32`，数据格式支持 $ND$，shape 需要与 `prev_softmax_max` 保持一致。
	47	+- cur_softmax_sum (`Tensor`)：必选参数，第二次 FlashAttention 的 softmax sum 结果。数据类型支持 `float32`，数据格式支持 $ND$，shape 需要与 `prev_softmax_max` 保持一致。
	M M molly1233215月21日是不是和哪个API配合使用呀，第一次XX，第二次XX都从哪里来呀
	48	+- <strong>*</strong>：位置参数与关键字参数的分隔符。其之前的参数为位置参数，需按顺序传入；其之后的参数为关键字参数，未显式赋值时使用默认值。
	49	+- actual_seq_qlen (`Tensor`)：可选参数，表示从 0 开始累计的 query 序列长度前缀和。数据类型支持 `int64`，数据格式支持 $ND$。当 `input_layout="TND"` 时必须传入，且张量中的值需要单调递增至总 token 数。
	50	+- input_layout (`str`)：可选参数，表示 `prev_attn_out` 和 `cur_attn_out` 的数据排布。支持的取值：
	51	+ - `"SBH"`：attention 输入输出按 `[S, B, H]` 排布。
	52	+ - `"TND"`：attention 输入输出按 `[T, N, D]` 排布。
	53	+ 默认值为 `"SBH"`。
	54	+- input_softmax_layout (`str`)：可选参数，表示 softmax 相关输入的排布方式。支持的取值：
	55	+ - `""`：使用默认排布。
	56	+ - `"SBH"`：softmax 输入按 SBH 语义排布。
	57	+ - `"TND"`：softmax 输入按 TND 语义排布。
	58	+ 默认值为 `""`。仅在 `input_layout="TND"` 时生效。
	M M molly1233215月21日都是默认值了，又仅在这个时候生效，逻辑有点问题把
	59	+
	60	+## 返回值说明
	61	+
	62	+- attn_out (`Tensor`)：更新后的 attention 输出。shape 与数据类型与输入 `prev_attn_out` 保持一致。
	63	+- softmax_max (`Tensor`)：更新后的 softmax max。shape 与 `prev_softmax_max` 保持一致，数据类型为 `float32`。
	64	+- softmax_sum (`Tensor`)：更新后的 softmax sum。shape 与 `prev_softmax_sum` 保持一致，数据类型为 `float32`。
	65	+
	66	+## 约束说明
	67	+
	68	+- 该接口支持推理、训练场景下使用。
	69	+- 该接口支持图模式。
	70	+- `prev_attn_out` 与 `cur_attn_out` 的 shape 和数据类型必须一致。
	71	+- `prev_softmax_max`、`prev_softmax_sum`、`cur_softmax_max`、`cur_softmax_sum` 的 shape 必须一致，且数据类型均需为 `float32`。
	72	+- `input_layout="TND"` 时，`actual_seq_qlen` 为必选参数。
	73	+- `input_layout="TND"` 时，`input_softmax_layout` 才生效，且只支持 `""`、`"SBH"`、`"TND"`。
	74	+- <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> 在 `input_layout="TND"` 场景下额外要求：
	75	+
	76	+\| 约束项 \| 限制 \|
	77	+\| ------ \| ---- \|
	78	+\| `N` \| `N <= 256` \|
	79	+\| `D` \| `D <= 768` 且 `D` 为 `64` 的倍数 \|
	80	+
	81	+## 调用示例
	82	+
	83	+- 单算子模式调用
	84	+
	85	+ ```python
	86	+ >>> import torch
	87	+ >>> import torch_npu
	88	+ >>>
	89	+ >>> prev_attn_out = torch.randn((4, 2, 32), dtype=torch.float16, device="npu")
	90	+ >>> cur_attn_out = torch.randn((4, 2, 32), dtype=torch.float16, device="npu")
	91	+ >>> prev_softmax_max = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
	92	+ >>> prev_softmax_sum = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
	93	+ >>> cur_softmax_max = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
	94	+ >>> cur_softmax_sum = torch.rand((2, 2, 4, 1), dtype=torch.float32, device="npu").repeat(1, 1, 1, 8)
	95	+ >>> attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
	96	+ ... prev_attn_out, prev_softmax_max, prev_softmax_sum,
	97	+ ... cur_attn_out, cur_softmax_max, cur_softmax_sum)
	98	+ >>> attn_out.shape
	99	+ torch.Size([4, 2, 32])
	100	+ >>> softmax_max.shape
	101	+ torch.Size([2, 2, 4, 8])
	102	+ >>> softmax_sum.shape
	103	+ torch.Size([2, 2, 4, 8])
	104	+ ```
	105	+
	106	+- 图模式调用
	107	+
	108	+ ```python
	109	+ import torch
	110	+ import torch_npu
	111	+ import torchair as tng
	112	+ from torchair.configs.compiler_config import CompilerConfig
	113	+
	114	+ config = CompilerConfig()
	115	+ npu_backend = tng.get_npu_backend(compiler_config=config)
	116	+
	117	+ class Model(torch.nn.Module):
	118	+ def forward(self, prev_attn_out, prev_softmax_max, prev_softmax_sum,
	119	+ cur_attn_out, cur_softmax_max, cur_softmax_sum):
	120	+ return torch_npu.npu_ring_attention_update(
	121	+ prev_attn_out, prev_softmax_max, prev_softmax_sum,
	122	+ cur_attn_out, cur_softmax_max, cur_softmax_sum)
	123	+
	124	+ model = Model().npu()
	125	+ model = torch.compile(model, backend=npu_backend, dynamic=False, fullgraph=True)
	126	+ outputs = model(prev_attn_out, prev_softmax_max, prev_softmax_sum,
	127	+ cur_attn_out, cur_softmax_max, cur_softmax_sum)
	128	+ print(outputs[0].shape)
	129	+ ```

Mdocs/zh/custom_APIs/torch_npu/torch_npu_list.md+5-0

	@@ -837,5 +837,10 @@
837	<td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p_npu_apply_rotary_pos_emb_desc"><a name="p_npu_apply_rotary_pos_emb_desc"></a><a name="p_npu_apply_rotary_pos_emb_desc"></a><span>为提升推理网络性能，将query和key两路算子融合为单路，在旋转位置编码计算中直接对结果执行原地更新。</span></p>	837	<td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p_npu_apply_rotary_pos_emb_desc"><a name="p_npu_apply_rotary_pos_emb_desc"></a><a name="p_npu_apply_rotary_pos_emb_desc"></a><span>为提升推理网络性能，将query和key两路算子融合为单路，在旋转位置编码计算中直接对结果执行原地更新。</span></p>
838	</td>	838	</td>
839	</tr>	839	</tr>
		840	+<tr id="row_ring_attention_update"><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p id="p_ring_attention_update"><a name="p_ring_attention_update"></a><a name="p_ring_attention_update"></a><a href="torch_npu-npu_ring_attention_update.md">torch_npu.npu_ring_attention_update</a></p>
		841	+</td>
		842	+<td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p_ring_attention_update_desc"><a name="p_ring_attention_update_desc"></a><a name="p_ring_attention_update_desc"></a><span>将两次 FlashAttention 的输出按照 softmax max 和 softmax sum 进行增量更新，支持 SBH 与 TND 排布。</span></p>
		843	+</td>
		844	+</tr>
840	</tbody>	845	</tbody>
841	</table>	846	</table>

Mop_plugin/config/op_plugin_functions.yaml+4-0

	@@ -5774,6 +5774,10 @@ custom:
5774	- func: npu_attention_update(Tensor[] lse, Tensor[] local_out, int update_type) -> (Tensor, Tensor)	5774	- func: npu_attention_update(Tensor[] lse, Tensor[] local_out, int update_type) -> (Tensor, Tensor)
5775	op_api: all_version	5775	op_api: all_version
5776		5776
		5777	+ - func: npu_ring_attention_update(Tensor prev_attn_out, Tensor prev_softmax_max, Tensor prev_softmax_sum, Tensor cur_attn_out, Tensor cur_softmax_max, Tensor cur_softmax_sum, *, Tensor? actual_seq_qlen=None, str input_layout="SBH", str input_softmax_layout="") -> (Tensor, Tensor, Tensor)
		5778	+ op_api: [v2.7, newest]
		5779	+ exposed: [v2.7, newest]
		5780	+
5777	- func: npu_alltoallv_gmm(Tensor gmm_x, Tensor gmm_weight, str hcom, int ep_world_size, int[] send_counts, int[] recv_counts, *, Tensor? send_counts_tensor=None, Tensor? recv_counts_tensor=None, Tensor? mm_x=None, Tensor? mm_weight=None, bool trans_gmm_weight=False, bool trans_mm_weight=False, bool permute_out_flag=False) -> (Tensor, Tensor, Tensor)	5781	- func: npu_alltoallv_gmm(Tensor gmm_x, Tensor gmm_weight, str hcom, int ep_world_size, int[] send_counts, int[] recv_counts, *, Tensor? send_counts_tensor=None, Tensor? recv_counts_tensor=None, Tensor? mm_x=None, Tensor? mm_weight=None, bool trans_gmm_weight=False, bool trans_mm_weight=False, bool permute_out_flag=False) -> (Tensor, Tensor, Tensor)
5778	op_api: all_version	5782	op_api: all_version
5779	exposed: all_version	5783	exposed: all_version

Aop_plugin/ops/opapi/RingAttentionUpdateKernelNpuOpApi.cpp+167-0

@@ -0,0 +1,167 @@
	1	+// Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+// All rights reserved.
	3	+//
	4	+// Licensed under the BSD 3-Clause License (the "License");
	5	+// you may not use this file except in compliance with the License.
	6	+//
	7	+// https://opensource.org/licenses/BSD-3-Clause
	8	+//
	9	+// Unless required by applicable law or agreed to in writing, software
	10	+// distributed under the License is distributed on an "AS IS" BASIS,
	11	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+// See the License for the specific language governing permissions and
	13	+// limitations under the License.
	14	+
	15	+#include <string>
	16	+
	17	+#include "op_plugin/OpApiInterface.h"
	18	+#include "op_plugin/utils/op_api_common.h"
	19	+
	20	+namespace op_api {
	21	+namespace {
	22	+constexpr int64_t SOFTMAX_LAST_DIM = 8;
	23	+
	24	+bool is_supported_attn_dtype(at::ScalarType dtype)
	25	+{
	26	+ return dtype == at::kHalf \|\| dtype == at::kFloat \|\| dtype == at::kBFloat16;
	27	+}
	28	+
	29	+void check_attention_inputs(
	30	+ const at::Tensor &prev_attn_out,
	31	+ const at::Tensor &prev_softmax_max,
	32	+ const at::Tensor &prev_softmax_sum,
	33	+ const at::Tensor &cur_attn_out,
	34	+ const at::Tensor &cur_softmax_max,
	35	+ const at::Tensor &cur_softmax_sum,
	36	+ const c10::optional<at::Tensor> &actual_seq_qlen,
	37	+ c10::string_view input_layout,
	38	+ c10::string_view input_softmax_layout)
	39	+{
	40	+ std::string input_layout_str = std::string(input_layout);
	41	+ std::string input_softmax_layout_str = std::string(input_softmax_layout);
	42	+
	43	+ TORCH_CHECK(
	44	+ input_layout_str == "SBH" \|\| input_layout_str == "TND",
	45	+ "input_layout only supports 'SBH' or 'TND', but got ", input_layout_str,
	46	+ OPS_ERROR(ErrCode::VALUE));
	47	+ TORCH_CHECK(
	48	+ input_softmax_layout_str.empty() \|\| input_softmax_layout_str == "SBH" \|\| input_softmax_layout_str == "TND",
	49	+ "input_softmax_layout only supports '', 'SBH' or 'TND', but got ", input_softmax_layout_str,
	50	+ OPS_ERROR(ErrCode::VALUE));
	51	+
	52	+ TORCH_CHECK(
	53	+ prev_attn_out.sizes() == cur_attn_out.sizes(),
	54	+ "prev_attn_out and cur_attn_out must have the same shape.",
	55	+ OPS_ERROR(ErrCode::VALUE));
	56	+ TORCH_CHECK(
	57	+ prev_attn_out.scalar_type() == cur_attn_out.scalar_type(),
	58	+ "prev_attn_out and cur_attn_out must have the same dtype.",
	59	+ OPS_ERROR(ErrCode::VALUE));
	60	+ TORCH_CHECK(
	61	+ is_supported_attn_dtype(prev_attn_out.scalar_type()),
	62	+ "attn_out dtype only supports float16, float32 or bfloat16, but got ",
	63	+ prev_attn_out.scalar_type(),
	64	+ OPS_ERROR(ErrCode::TYPE));
	65	+
	66	+ TORCH_CHECK(
	67	+ prev_softmax_max.sizes() == prev_softmax_sum.sizes() &&
	68	+ prev_softmax_max.sizes() == cur_softmax_max.sizes() &&
	69	+ prev_softmax_max.sizes() == cur_softmax_sum.sizes(),
	70	+ "softmax tensors must keep the same shape.",
	71	+ OPS_ERROR(ErrCode::VALUE));
	72	+ TORCH_CHECK(
	73	+ prev_softmax_max.scalar_type() == at::kFloat &&
	74	+ prev_softmax_sum.scalar_type() == at::kFloat &&
	75	+ cur_softmax_max.scalar_type() == at::kFloat &&
	76	+ cur_softmax_sum.scalar_type() == at::kFloat,
	77	+ "softmax tensors must use float32.",
	78	+ OPS_ERROR(ErrCode::TYPE));
	79	+ TORCH_CHECK(
	80	+ prev_softmax_max.dim() > 0 && prev_softmax_max.size(prev_softmax_max.dim() - 1) == SOFTMAX_LAST_DIM,
	81	+ "softmax tensors require the last dimension to be 8.",
	82	+ OPS_ERROR(ErrCode::VALUE));
	83	+
	84	+ TORCH_CHECK(
	85	+ prev_attn_out.dim() == 3,
	86	+ "prev_attn_out must be a 3D tensor, but got ", prev_attn_out.dim(), "D.",
	87	+ OPS_ERROR(ErrCode::VALUE));
	88	+ TORCH_CHECK(
	89	+ cur_attn_out.dim() == 3,
	90	+ "cur_attn_out must be a 3D tensor, but got ", cur_attn_out.dim(), "D.",
	91	+ OPS_ERROR(ErrCode::VALUE));
	92	+
	93	+ if (input_layout_str == "SBH") {
	94	+ TORCH_CHECK(
	95	+ prev_softmax_max.dim() == 4,
	96	+ "softmax tensors must be 4D when input_layout is 'SBH'.",
	97	+ OPS_ERROR(ErrCode::VALUE));
	98	+ } else {
	99	+ TORCH_CHECK(
	100	+ actual_seq_qlen.has_value() && actual_seq_qlen.value().defined(),
	101	+ "actual_seq_qlen is required when input_layout is 'TND'.",
	102	+ OPS_ERROR(ErrCode::VALUE));
	103	+ TORCH_CHECK(
	104	+ prev_softmax_max.dim() == 3,
	105	+ "softmax tensors must be 3D when input_layout is 'TND'.",
	106	+ OPS_ERROR(ErrCode::VALUE));
	107	+ }
	108	+
	109	+ if (actual_seq_qlen.has_value() && actual_seq_qlen.value().defined()) {
	110	+ TORCH_CHECK(
	111	+ actual_seq_qlen.value().scalar_type() == at::kLong,
	112	+ "actual_seq_qlen must use int64 dtype.",
	113	+ OPS_ERROR(ErrCode::TYPE));
	114	+ }
	115	+}
	116	+} // namespace
	117	+
	118	+using npu_preparation = at_npu::native::OpPreparation;
	119	+
	120	+std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_ring_attention_update(
	121	+ const at::Tensor &prev_attn_out,
	122	+ const at::Tensor &prev_softmax_max,
	123	+ const at::Tensor &prev_softmax_sum,
	124	+ const at::Tensor &cur_attn_out,
	125	+ const at::Tensor &cur_softmax_max,
	126	+ const at::Tensor &cur_softmax_sum,
	127	+ const c10::optional<at::Tensor> &actual_seq_qlen,
	128	+ c10::string_view input_layout,
	129	+ c10::string_view input_softmax_layout)
	130	+{
	131	+ check_attention_inputs(
	132	+ prev_attn_out,
	133	+ prev_softmax_max,
	134	+ prev_softmax_sum,
	135	+ cur_attn_out,
	136	+ cur_softmax_max,
	137	+ cur_softmax_sum,
	138	+ actual_seq_qlen,
	139	+ input_layout,
	140	+ input_softmax_layout);
	141	+
	142	+ const at::Tensor &actual_seq_qlen_tensor = actual_seq_qlen.value_or(at::Tensor());
	143	+ char input_layout_ptr = const_cast<char >(input_layout.data());
	144	+ char input_softmax_layout_ptr = const_cast<char >(input_softmax_layout.data());
	145	+
	146	+ at::Tensor attn_out = npu_preparation::apply_tensor_without_format(prev_attn_out);
	147	+ at::Tensor softmax_max = npu_preparation::apply_tensor_without_format(prev_softmax_max);
	148	+ at::Tensor softmax_sum = npu_preparation::apply_tensor_without_format(prev_softmax_sum);
	149	+
	150	+ EXEC_NPU_CMD(
	151	+ aclnnRingAttentionUpdateV2,
	152	+ prev_attn_out,
	153	+ prev_softmax_max,
	154	+ prev_softmax_sum,
	155	+ cur_attn_out,
	156	+ cur_softmax_max,
	157	+ cur_softmax_sum,
	158	+ actual_seq_qlen_tensor,
	159	+ input_layout_ptr,
	160	+ input_softmax_layout_ptr,
	161	+ attn_out,
	162	+ softmax_max,
	163	+ softmax_sum);
	164	+
	165	+ return std::make_tuple(attn_out, softmax_max, softmax_sum);
	166	+}
	167	+} // namespace op_api

Mop_plugin/python/meta/_meta_registrations.py+73-0

	@@ -6111,6 +6111,79 @@ def npu_attention_update_meta(lse, local_out, update_type):
6111	torch.empty(ref_lse.size(), dtype=ref_lse.dtype, device=ref_lse.device))	6111	torch.empty(ref_lse.size(), dtype=ref_lse.dtype, device=ref_lse.device))
6112		6112
6113		6113
		6114	+@impl(m, "npu_ring_attention_update")
		6115	+def npu_ring_attention_update_meta(prev_attn_out, prev_softmax_max, prev_softmax_sum, cur_attn_out,
		6116	+ cur_softmax_max, cur_softmax_sum, *, actual_seq_qlen=None,
		6117	+ input_layout="SBH", input_softmax_layout=""):
		6118	+ torch._check(
		6119	+ input_layout in ("SBH", "TND"),
		6120	+ lambda: f"input_layout only supports 'SBH' or 'TND', but got {input_layout}." + ops_error(ErrCode.VALUE),
		6121	+ )
		6122	+ torch._check(
		6123	+ input_softmax_layout in ("", "SBH", "TND"),
		6124	+ lambda: "input_softmax_layout only supports '', 'SBH' or 'TND'." + ops_error(ErrCode.VALUE),
		6125	+ )
		6126	+ torch._check(
		6127	+ prev_attn_out.shape == cur_attn_out.shape,
		6128	+ lambda: "prev_attn_out and cur_attn_out must have the same shape." + ops_error(ErrCode.VALUE),
		6129	+ )
		6130	+ torch._check(
		6131	+ prev_attn_out.dtype == cur_attn_out.dtype,
		6132	+ lambda: "prev_attn_out and cur_attn_out must have the same dtype." + ops_error(ErrCode.VALUE),
		6133	+ )
		6134	+ torch._check(
		6135	+ prev_attn_out.dtype in (torch.float16, torch.float32, torch.bfloat16),
		6136	+ lambda: "attn_out dtype only supports float16, float32 or bfloat16." + ops_error(ErrCode.TYPE),
		6137	+ )
		6138	+ torch._check(
		6139	+ prev_softmax_max.shape == prev_softmax_sum.shape ==
		6140	+ cur_softmax_max.shape == cur_softmax_sum.shape,
		6141	+ lambda: "softmax tensors must keep the same shape." + ops_error(ErrCode.VALUE),
		6142	+ )
		6143	+ torch._check(
		6144	+ prev_softmax_max.dtype == torch.float32 and
		6145	+ prev_softmax_sum.dtype == torch.float32 and
		6146	+ cur_softmax_max.dtype == torch.float32 and
		6147	+ cur_softmax_sum.dtype == torch.float32,
		6148	+ lambda: "softmax tensors must use float32." + ops_error(ErrCode.TYPE),
		6149	+ )
		6150	+ torch._check(
		6151	+ prev_softmax_max.dim() > 0 and prev_softmax_max.size(-1) == 8,
		6152	+ lambda: "softmax tensors require the last dimension to be 8." + ops_error(ErrCode.VALUE),
		6153	+ )
		6154	+ torch._check(
		6155	+ prev_attn_out.dim() == 3 and cur_attn_out.dim() == 3,
		6156	+ lambda: "attention tensors must be 3D." + ops_error(ErrCode.VALUE),
		6157	+ )
		6158	+
		6159	+ if input_layout == "SBH":
		6160	+ torch._check(
		6161	+ prev_softmax_max.dim() == 4,
		6162	+ lambda: "softmax tensors must be 4D when input_layout is 'SBH'." + ops_error(ErrCode.VALUE),
		6163	+ )
		6164	+ else:
		6165	+ torch._check(
		6166	+ actual_seq_qlen is not None,
		6167	+ lambda: "actual_seq_qlen is required when input_layout is 'TND'." + ops_error(ErrCode.VALUE),
		6168	+ )
		6169	+ torch._check(
		6170	+ prev_softmax_max.dim() == 3,
		6171	+ lambda: "softmax tensors must be 3D when input_layout is 'TND'." + ops_error(ErrCode.VALUE),
		6172	+ )
		6173	+
		6174	+ if actual_seq_qlen is not None:
		6175	+ torch._check(
		6176	+ actual_seq_qlen.dtype == torch.int64,
		6177	+ lambda: "actual_seq_qlen must use int64 dtype." + ops_error(ErrCode.TYPE),
		6178	+ )
		6179	+
		6180	+ return (
		6181	+ torch.empty_like(prev_attn_out, device='meta'),
		6182	+ torch.empty_like(prev_softmax_max, device='meta'),
		6183	+ torch.empty_like(prev_softmax_sum, device='meta'),
		6184	+ )
		6185	+
		6186	+
6114	@impl(m, "npu_mrope")	6187	@impl(m, "npu_mrope")
6115	def npu_mrope_meta(positions, query, key, cos_sin_cache, head_size, *, mrope_section=None, rotary_mode='half', cache_mode='default'):	6188	def npu_mrope_meta(positions, query, key, cos_sin_cache, head_size, *, mrope_section=None, rotary_mode='half', cache_mode='default'):
6116	return (torch.empty_like(query), torch.empty_like(key))	6189	return (torch.empty_like(query), torch.empty_like(key))

Mtest/allowlist_for_publicAPI.json+20-13

	@@ -22,7 +22,8 @@
22	"npu_fused_causal_conv1d",	22	"npu_fused_causal_conv1d",
23	"npu_masked_causal_conv1d",	23	"npu_masked_causal_conv1d",
24	"npu_masked_causal_conv1d_backward",	24	"npu_masked_causal_conv1d_backward",
25	- "npu_apply_rotary_pos_emb"	25	+ "npu_apply_rotary_pos_emb",
		26	+ "npu_ring_attention_update"
26	],	27	],
27	"v2.8": [	28	"v2.8": [
28	"npu_gelu_mul",	29	"npu_gelu_mul",
	@@ -33,7 +34,8 @@
33	"npu_fused_causal_conv1d",	34	"npu_fused_causal_conv1d",
34	"npu_masked_causal_conv1d",	35	"npu_masked_causal_conv1d",
35	"npu_masked_causal_conv1d_backward",	36	"npu_masked_causal_conv1d_backward",
36	- "npu_apply_rotary_pos_emb"	37	+ "npu_apply_rotary_pos_emb",
		38	+ "npu_ring_attention_update"
37	],	39	],
38	"v2.9": [	40	"v2.9": [
39	"npu_gelu_mul",	41	"npu_gelu_mul",
	@@ -44,7 +46,8 @@
44	"npu_fused_causal_conv1d",	46	"npu_fused_causal_conv1d",
45	"npu_masked_causal_conv1d",	47	"npu_masked_causal_conv1d",
46	"npu_masked_causal_conv1d_backward",	48	"npu_masked_causal_conv1d_backward",
47	- "npu_apply_rotary_pos_emb"	49	+ "npu_apply_rotary_pos_emb",
		50	+ "npu_ring_attention_update"
48	],	51	],
49	"v2.10": [	52	"v2.10": [
50	"npu_gelu_mul",	53	"npu_gelu_mul",
	@@ -55,7 +58,8 @@
55	"npu_fused_causal_conv1d",	58	"npu_fused_causal_conv1d",
56	"npu_masked_causal_conv1d",	59	"npu_masked_causal_conv1d",
57	"npu_masked_causal_conv1d_backward",	60	"npu_masked_causal_conv1d_backward",
58	- "npu_apply_rotary_pos_emb"	61	+ "npu_apply_rotary_pos_emb",
		62	+ "npu_ring_attention_update"
59	],	63	],
60	"v2.11": [	64	"v2.11": [
61	"npu_gelu_mul",	65	"npu_gelu_mul",
	@@ -66,7 +70,8 @@
66	"npu_fused_causal_conv1d",	70	"npu_fused_causal_conv1d",
67	"npu_masked_causal_conv1d",	71	"npu_masked_causal_conv1d",
68	"npu_masked_causal_conv1d_backward",	72	"npu_masked_causal_conv1d_backward",
69	- "npu_apply_rotary_pos_emb"	73	+ "npu_apply_rotary_pos_emb",
		74	+ "npu_ring_attention_update"
70	],	75	],
71	"v2.12": [	76	"v2.12": [
72	"npu_gelu_mul",	77	"npu_gelu_mul",
	@@ -77,18 +82,20 @@
77	"npu_fused_causal_conv1d",	82	"npu_fused_causal_conv1d",
78	"npu_masked_causal_conv1d",	83	"npu_masked_causal_conv1d",
79	"npu_masked_causal_conv1d_backward",	84	"npu_masked_causal_conv1d_backward",
80	- "npu_apply_rotary_pos_emb"	85	+ "npu_apply_rotary_pos_emb",
		86	+ "npu_ring_attention_update"
81	],	87	],
82	"v2.13": [	88	"v2.13": [
83	- "npu_gelu_mul",	89	+ "npu_gelu_mul",
84	- "npu_clipped_swiglu",	90	+ "npu_clipped_swiglu",
85	- "npu_sim_exponential_",	91	+ "npu_sim_exponential_",
86	- "npu_fused_floyd_attention",	92	+ "npu_fused_floyd_attention",
87	- "npu_fused_floyd_attention_backward",	93	+ "npu_fused_floyd_attention_backward",
88	- "npu_fused_causal_conv1d",	94	+ "npu_fused_causal_conv1d",
89	"npu_masked_causal_conv1d",	95	"npu_masked_causal_conv1d",
90	"npu_masked_causal_conv1d_backward",	96	"npu_masked_causal_conv1d_backward",
91	- "npu_apply_rotary_pos_emb"	97	+ "npu_apply_rotary_pos_emb",
		98	+ "npu_ring_attention_update"
92	],	99	],
93	"all_version": [	100	"all_version": [
94	"_npu_dropout",	101	"_npu_dropout",

Mtest/core_tests/test_fake_tensor.py+42-0

	@@ -3181,6 +3181,48 @@ class TestAttentionUpdate(TestCase):
3181	self.assertTrue(out.dtype == dtype)	3181	self.assertTrue(out.dtype == dtype)
3182		3182
3183		3183
		3184	+class TestRingAttentionUpdate(TestCase):
		3185	+ def test_npu_ring_attention_update_meta_sbh(self):
		3186	+ with FakeTensorMode():
		3187	+ prev_attn_out = torch.randn((4, 2, 32), dtype=torch.float16).npu()
		3188	+ cur_attn_out = torch.randn((4, 2, 32), dtype=torch.float16).npu()
		3189	+ prev_softmax_max = torch.randn((2, 2, 4, 8), dtype=torch.float32).abs().npu()
		3190	+ prev_softmax_sum = torch.randn((2, 2, 4, 8), dtype=torch.float32).abs().npu()
		3191	+ cur_softmax_max = torch.randn((2, 2, 4, 8), dtype=torch.float32).abs().npu()
		3192	+ cur_softmax_sum = torch.randn((2, 2, 4, 8), dtype=torch.float32).abs().npu()
		3193	+
		3194	+ attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
		3195	+ prev_attn_out, prev_softmax_max, prev_softmax_sum,
		3196	+ cur_attn_out, cur_softmax_max, cur_softmax_sum)
		3197	+ self.assertEqual(attn_out.shape, prev_attn_out.shape)
		3198	+ self.assertEqual(attn_out.dtype, prev_attn_out.dtype)
		3199	+ self.assertEqual(attn_out.device.type, "npu")
		3200	+ self.assertEqual(softmax_max.shape, prev_softmax_max.shape)
		3201	+ self.assertEqual(softmax_max.dtype, torch.float32)
		3202	+ self.assertEqual(softmax_sum.shape, prev_softmax_sum.shape)
		3203	+ self.assertEqual(softmax_sum.dtype, torch.float32)
		3204	+
		3205	+ def test_npu_ring_attention_update_meta_tnd(self):
		3206	+ with FakeTensorMode():
		3207	+ prev_attn_out = torch.randn((5, 2, 64), dtype=torch.bfloat16).npu()
		3208	+ cur_attn_out = torch.randn((5, 2, 64), dtype=torch.bfloat16).npu()
		3209	+ prev_softmax_max = torch.randn((5, 2, 8), dtype=torch.float32).abs().npu()
		3210	+ prev_softmax_sum = torch.randn((5, 2, 8), dtype=torch.float32).abs().npu()
		3211	+ cur_softmax_max = torch.randn((5, 2, 8), dtype=torch.float32).abs().npu()
		3212	+ cur_softmax_sum = torch.randn((5, 2, 8), dtype=torch.float32).abs().npu()
		3213	+ actual_seq_qlen = torch.tensor([2, 5], dtype=torch.int64).npu()
		3214	+
		3215	+ attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
		3216	+ prev_attn_out, prev_softmax_max, prev_softmax_sum,
		3217	+ cur_attn_out, cur_softmax_max, cur_softmax_sum,
		3218	+ actual_seq_qlen=actual_seq_qlen, input_layout="TND", input_softmax_layout="TND")
		3219	+ self.assertEqual(attn_out.shape, prev_attn_out.shape)
		3220	+ self.assertEqual(attn_out.dtype, prev_attn_out.dtype)
		3221	+ self.assertEqual(attn_out.device.type, "npu")
		3222	+ self.assertEqual(softmax_max.shape, prev_softmax_max.shape)
		3223	+ self.assertEqual(softmax_sum.shape, prev_softmax_sum.shape)
		3224	+
		3225	+
3184	class TestAntiQuant(TestCase):	3226	class TestAntiQuant(TestCase):
3185	@unittest.skipIf(torch.__version__ < '2.1.0',	3227	@unittest.skipIf(torch.__version__ < '2.1.0',
3186	"OP `AntiQuant` is supported on torch v2.1 and above, skip this test for torch version below 2.1")	3228	"OP `AntiQuant` is supported on torch v2.1 and above, skip this test for torch version below 2.1")

Mtest/core_tests/torch_npu_OpApi_schema_all.json+10-1

	@@ -197,6 +197,12 @@
197	"op_api: torch_npu.npu_alloc_float_status(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {	197	"op_api: torch_npu.npu_alloc_float_status(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
198	"version": ["v2.8", "newest"]	198	"version": ["v2.8", "newest"]
199	},	199	},
		200	+ "op_api: torch_npu.npu_ring_attention_update(args, *kwargs)": {
		201	+ "version": ["v2.7"]
		202	+ },
		203	+ "op_api: torch_npu.npu_ring_attention_update(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
		204	+ "version": ["v2.8", "newest"]
		205	+ },
200	"op_api: torch_npu._npu_dropout(args, *kwargs)": {	206	"op_api: torch_npu._npu_dropout(args, *kwargs)": {
201	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]	207	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]
202	},	208	},
	@@ -1210,6 +1216,9 @@
1210	"func: npu_attention_update(Tensor[] lse, Tensor[] local_out, int update_type) -> (Tensor, Tensor)": {	1216	"func: npu_attention_update(Tensor[] lse, Tensor[] local_out, int update_type) -> (Tensor, Tensor)": {
1211	"version": ["all_version"]	1217	"version": ["all_version"]
1212	},	1218	},
		1219	+ "func: npu_ring_attention_update(Tensor prev_attn_out, Tensor prev_softmax_max, Tensor prev_softmax_sum, Tensor cur_attn_out, Tensor cur_softmax_max, Tensor cur_softmax_sum, *, Tensor? actual_seq_qlen=None, str input_layout=\"SBH\", str input_softmax_layout=\"\") -> (Tensor, Tensor, Tensor)": {
		1220	+ "version": ["v2.7", "newest"]
		1221	+ },
1213	"func: npu_add_layer_norm_backward(Tensor? dy_opt, Tensor x1, Tensor x2, Tensor rstd, Tensor mean, Tensor gamma, Tensor? dsum_opt) -> (Tensor, Tensor, Tensor, Tensor)": {	1222	"func: npu_add_layer_norm_backward(Tensor? dy_opt, Tensor x1, Tensor x2, Tensor rstd, Tensor mean, Tensor gamma, Tensor? dsum_opt) -> (Tensor, Tensor, Tensor, Tensor)": {
1214	"version": ["all_version"]	1223	"version": ["all_version"]
1215	},	1224	},
	@@ -5015,4 +5024,4 @@
5015	"func: _fused_sdp_choice(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> int": {	5024	"func: _fused_sdp_choice(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> int": {
5016	"version": ["v2.7", "newest"]	5025	"version": ["v2.7", "newest"]
5017	}	5026	}
5018	-}	5027	+}

Atest/test_custom_ops/test_npu_ring_attention_update.py+135-0

@@ -0,0 +1,135 @@
	1	+import torch
	2	+import torch_npu
	3	+
	4	+from torch_npu.testing.testcase import TestCase, run_tests
	5	+from torch_npu.testing.common_utils import SupportedDevices
	6	+
	7	+
	8	+def _repeat_last_dim_8(x):
	9	+ return x.unsqueeze(-1).repeat_interleave(8, dim=-1).contiguous()
	10	+
	11	+
	12	+def _ring_attention_update_golden(
	13	+ prev_attn_out,
	14	+ prev_softmax_max,
	15	+ prev_softmax_sum,
	16	+ cur_attn_out,
	17	+ cur_softmax_max,
	18	+ cur_softmax_sum,
	19	+ input_layout,
	20	+):
	21	+ prev_attn = prev_attn_out.to(torch.float32).cpu()
	22	+ cur_attn = cur_attn_out.to(torch.float32).cpu()
	23	+ prev_max = prev_softmax_max[..., 0].to(torch.float32).cpu()
	24	+ prev_sum = prev_softmax_sum[..., 0].to(torch.float32).cpu()
	25	+ cur_max = cur_softmax_max[..., 0].to(torch.float32).cpu()
	26	+ cur_sum = cur_softmax_sum[..., 0].to(torch.float32).cpu()
	27	+
	28	+ softmax_max = torch.maximum(prev_max, cur_max)
	29	+ prev_scale = torch.exp(prev_max - softmax_max)
	30	+ cur_scale = torch.exp(cur_max - softmax_max)
	31	+ softmax_sum = prev_sum * prev_scale + cur_sum * cur_scale
	32	+
	33	+ if input_layout == "SBH":
	34	+ seq_len, batch_size, hidden_size = prev_attn.shape
	35	+ head_num = prev_max.shape[1]
	36	+ head_dim = hidden_size // head_num
	37	+ prev_attn = prev_attn.reshape(seq_len, batch_size, head_num, head_dim)
	38	+ cur_attn = cur_attn.reshape(seq_len, batch_size, head_num, head_dim)
	39	+ prev_factor = (prev_scale * prev_sum / softmax_sum).permute(2, 0, 1).unsqueeze(-1)
	40	+ cur_factor = (cur_scale * cur_sum / softmax_sum).permute(2, 0, 1).unsqueeze(-1)
	41	+ attn_out = prev_attn * prev_factor + cur_attn * cur_factor
	42	+ attn_out = attn_out.reshape(seq_len, batch_size, hidden_size)
	43	+ else:
	44	+ prev_factor = (prev_scale * prev_sum / softmax_sum).unsqueeze(-1)
	45	+ cur_factor = (cur_scale * cur_sum / softmax_sum).unsqueeze(-1)
	46	+ attn_out = prev_attn * prev_factor + cur_attn * cur_factor
	47	+
	48	+ return (
	49	+ attn_out.to(prev_attn_out.dtype),
	50	+ _repeat_last_dim_8(softmax_max),
	51	+ _repeat_last_dim_8(softmax_sum),
	52	+ )
	53	+
	54	+
	55	+class TestNpuRingAttentionUpdate(TestCase):
	56	+ def _run_sbh_case(self, dtype):
	57	+ seq_len, batch_size, head_num, head_dim = 4, 2, 2, 16
	58	+ hidden_size = head_num * head_dim
	59	+
	60	+ prev_attn_out = torch.randn(seq_len, batch_size, hidden_size, dtype=dtype, device="npu")
	61	+ cur_attn_out = torch.randn(seq_len, batch_size, hidden_size, dtype=dtype, device="npu")
	62	+ prev_softmax_max = _repeat_last_dim_8(torch.rand(batch_size, head_num, seq_len, dtype=torch.float32, device="npu") + 0.2)
	63	+ prev_softmax_sum = _repeat_last_dim_8(torch.rand(batch_size, head_num, seq_len, dtype=torch.float32, device="npu") + 0.5)
	64	+ cur_softmax_max = _repeat_last_dim_8(torch.rand(batch_size, head_num, seq_len, dtype=torch.float32, device="npu") + 0.3)
	65	+ cur_softmax_sum = _repeat_last_dim_8(torch.rand(batch_size, head_num, seq_len, dtype=torch.float32, device="npu") + 0.4)
	66	+
	67	+ attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
	68	+ prev_attn_out,
	69	+ prev_softmax_max,
	70	+ prev_softmax_sum,
	71	+ cur_attn_out,
	72	+ cur_softmax_max,
	73	+ cur_softmax_sum,
	74	+ )
	75	+ golden = _ring_attention_update_golden(
	76	+ prev_attn_out,
	77	+ prev_softmax_max,
	78	+ prev_softmax_sum,
	79	+ cur_attn_out,
	80	+ cur_softmax_max,
	81	+ cur_softmax_sum,
	82	+ input_layout="SBH",
	83	+ )
	84	+ self.assertRtolEqual(attn_out.cpu(), golden[0], prec16=0.005)
	85	+ self.assertRtolEqual(softmax_max.cpu(), golden[1], prec16=0.005)
	86	+ self.assertRtolEqual(softmax_sum.cpu(), golden[2], prec16=0.005)
	87	+
	88	+ def _run_tnd_case(self, dtype):
	89	+ total_tokens, head_num, head_dim = 5, 2, 64
	90	+ actual_seq_qlen = torch.tensor([2, total_tokens], dtype=torch.int64, device="npu")
	91	+
	92	+ prev_attn_out = torch.randn(total_tokens, head_num, head_dim, dtype=dtype, device="npu")
	93	+ cur_attn_out = torch.randn(total_tokens, head_num, head_dim, dtype=dtype, device="npu")
	94	+ prev_softmax_max = _repeat_last_dim_8(torch.rand(total_tokens, head_num, dtype=torch.float32, device="npu") + 0.2)
	95	+ prev_softmax_sum = _repeat_last_dim_8(torch.rand(total_tokens, head_num, dtype=torch.float32, device="npu") + 0.5)
	96	+ cur_softmax_max = _repeat_last_dim_8(torch.rand(total_tokens, head_num, dtype=torch.float32, device="npu") + 0.3)
	97	+ cur_softmax_sum = _repeat_last_dim_8(torch.rand(total_tokens, head_num, dtype=torch.float32, device="npu") + 0.4)
	98	+
	99	+ attn_out, softmax_max, softmax_sum = torch_npu.npu_ring_attention_update(
	100	+ prev_attn_out,
	101	+ prev_softmax_max,
	102	+ prev_softmax_sum,
	103	+ cur_attn_out,
	104	+ cur_softmax_max,
	105	+ cur_softmax_sum,
	106	+ actual_seq_qlen=actual_seq_qlen,
	107	+ input_layout="TND",
	108	+ input_softmax_layout="TND",
	109	+ )
	110	+ golden = _ring_attention_update_golden(
	111	+ prev_attn_out,
	112	+ prev_softmax_max,
	113	+ prev_softmax_sum,
	114	+ cur_attn_out,
	115	+ cur_softmax_max,
	116	+ cur_softmax_sum,
	117	+ input_layout="TND",
	118	+ )
	119	+ self.assertRtolEqual(attn_out.cpu(), golden[0], prec16=0.005)
	120	+ self.assertRtolEqual(softmax_max.cpu(), golden[1], prec16=0.005)
	121	+ self.assertRtolEqual(softmax_sum.cpu(), golden[2], prec16=0.005)
	122	+
	123	+ @SupportedDevices(["Ascend910B", "Ascend910_93", "Ascend950"])
	124	+ def test_npu_ring_attention_update_sbh(self):
	125	+ for dtype in (torch.float16, torch.float32, torch.bfloat16):
	126	+ self._run_sbh_case(dtype)
	127	+
	128	+ @SupportedDevices(["Ascend910B", "Ascend910_93", "Ascend950"])
	129	+ def test_npu_ring_attention_update_tnd(self):
	130	+ for dtype in (torch.float16, torch.float32, torch.bfloat16):
	131	+ self._run_tnd_case(dtype)
	132	+
	133	+
	134	+if __name__ == "__main__":
	135	+ run_tests()