Star1663
3352
Code
Code
Issues463
Pull Requests677
Actions
Discussion
Wiki
Analysis
Star1663
3352
Merged
docs: add English translations for markdown files #15372
jenny_xiangAdd on 4 days ago
docs: add English translations for markdown files #15372
Merged
jenny_xiangAdd on 4 days ago
Merged fromjenny_xiang/ops-transformerdocs/add-english-translationsintocann/ops-transformermaster
Total: 277 files changed+54964-783
A.gitcode/PULL_REQUEST_TEMPLATE.en-US.md+27-0
@@ -0,0 +1,27 @@
1+## Description
2+<!--Describe your changes in detail here, including the reason for the changes and the method used.-->
3+ 
4+## Related Issues
5+<!-- If this PR is intended to resolve a specific issue, provide the issue link here. For example: Related Issue #000-->
6+<!-- If this PR is intended to resolve a specific problem ticket, describe the problem ticket number here.-->
7+ 
8+## Tests
9+<!--Describe the tests performed to verify your changes. Include but are not limited to secondary smoke tests, operator generalization tests, and so on.-->
10+ 
11+## Documentation Updates
12+<!--If this PR contains documentation updates, point them out here. For example: updated the README.md file.-->
13+ 
14+## Type Labels
15+<!-- [x] means selected -->
16+- [ ] 🐛 Bug fix
17+- [ ] ✨ New feature
18+- [ ] ⚡ Performance optimization
19+- [ ] ♻️ Refactoring
20+- [ ] 🧪 Test
21+- [ ] 📦 Build/CI
22+- [ ] 🔧 Configuration change
23+- [ ] 📝 Documentation update
24+- [ ] ⬆️ Dependency upgrade
25+- [ ] 🔒 Security fix
26+- [ ] 🧹 Code cleanup
27+- [ ] ❓ Other, please describe:
ACHANGELOG_en.md+51-0
@@ -0,0 +1,51 @@
1+# CHANGELOG
2+ 
3+> This document records the important changes of each version. Versions are listed in reverse chronological order.
4+ 
5+## 8.5.0-beta.1
6+ 
7+Release date: 2025-12-30
8+ 
9+The first Beta version 8.5.0-beta.1 of the ops-transformer operators is now available. This version introduces multiple new features, bug fixes, and performance improvements, and is still in the testing stage.
10+We sincerely welcome community feedback to further improve the stability and functional completeness of ops-transformer.
11+For usage instructions, refer to the [official documentation](https://atomgit.com/cann/ops-transformer/blob/master/README.md).
12+ 
13+### 🔗 Version Address
14+ 
15+[CANN 8.5.0-beta 1](https://ascend.devcloud.huaweicloud.com/cann/run/software/8.5.0-beta.1/)
16+ 
17+```
18+The version directory is described as follows:
19+├── aarch64 # CPU is of the ARM type
20+│ ├── ops # ops operator package directory, used to archive operator sub-packages
21+│ ├── ...
22+├── x86_64 # CPU is of the X86 type
23+│ ├── ops # ops operator package directory, used to archive operator sub-packages
24+│ ├── ...
25+```
26+ 
27+### 📌 Version Compatibility
28+ 
29+**Compatibility between the ops-transformer sub-package and related components and CANN versions**
30+ 
31+| **CANN sub-package version** | **Compatible CANN version** |
32+|:----------------------------------|---------------------|
33+| cann-ops-transformer 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
34+| cann-ops-math 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
35+| cann-ops-nn 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
36+| cann-ops-cv 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
37+| cann-hccl 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
38+| cann-hixl 8.5.0-beta.1 | CANN 8.5.0-beta.1 |
39+ 
40+### 🚀 Key Features
41+ 
42+- [Engineering capability] Support for transformer-type ONNX operator plugins. ([#539](https://atomgit.com/cann/ops-transformer/pull/539))
43+- [Operator implementation] Some inference operators now support KirinX90. ([#609](https://atomgit.com/cann/ops-transformer/pull/609))
44+- [Documentation optimization] Added QUICK_START, the offline compilation mode, and improved the development guides for aicore/graph modes. ([#612](https://atomgit.com/cann/ops-transformer/pull/612), [#629](https://atomgit.com/cann/ops-transformer/pull/629), [#342](https://atomgit.com/cann/ops-transformer/pull/342))
45+- [Documentation optimization] Optimized the new operator contribution process in the contribution guide. ([#384](https://atomgit.com/cann/ops-transformer/pull/384))
46+ 
47+### 🐛 Bug Fixes
48+ 
49+- Fixed the issue of supporting a unified torch.dist.group for the mc2 communicator. ([Issue47](https://atomgit.com/cann/ops-transformer/issues/47))
50+- Fixed the execution and invocation issue of the add_example sample operator. ([Issue174](https://atomgit.com/cann/ops-transformer/issues/174))
51+- Fixed the issue that the install_deps.sh script does not support the openEuler system. ([Issue255](https://atomgit.com/cann/ops-transformer/issues/255))
Aattention/README_en.md+401-0
@@ -0,0 +1,401 @@
1+English | [简体中文](./README.md)
2+ 
3+## 1 attention layer Operator Classification Principles
4+ 
5+attention layer operators are divided into two major categories by whether Softmax is used: Softmax Attention (Softmax-based attention) and Linear Attention (linear attention).
6+ 
7+| Major category | Subcategory | Core features | Typical networks | Corresponding operators |
8+| :--- | :--- | :--- | :--- | :--- |
9+| Softmax Attention | GQA | Standard MHA/GQA dense attention. Executes Softmax(QK^T/√d)V on all KV with no sparse filtering. Covers training forward/backward and inference full/incremental scenarios | Llama, Qwen, GPT | flash_attention_score, flash_attention_score_grad, flash_attn, fused_infer_attention_score, quant_flash_attn |
10+| Softmax Attention | MLA | Multi-head Latent Attention. K/V share the same low-dimensional latent representation. Generates query/query_rope through downsampling+RmsNorm+upsampling+RoPE preprocessing and updates the KV cache | DeepSeek-V2/V3.1 | fused_infer_attention_score, mla_prolog, mla_prolog_v2, mla_prolog_v3 |
11+| Softmax Attention | Sparse | Sparse Flash Attention / Sparse Flash MLA. First uses the Lightning Indexer to select the Top-K key KV subset, and then performs attention on the selected KV. Includes variants such as KV compression (Compressor) and dual-path sparsity | DeepSeek-V3.2/V4 (SMLA/QLI) | lightning_indexer, lightning_indexer_grad, quant_lightning_indexer, sparse_flash_attention, sparse_flash_attention_grad, kv_quant_sparse_flash_attention, compressor, compressor_grad, quant_compressor, lightning_indexer_v2, quant_lightning_indexer_v2, sparse_flash_mla, sparse_flash_mla_grad, sparse_flash_mla_softmax_l1_norm, quant_sparse_flash_mla, mixed_quant_sparse_flash_mla |
12+| Linear Attention | GDN | Gated Delta Rule based on the Gated DeltaNet paper. Recursively updates states through gating. Provides two computation modes: Recurrent (Decode step-by-step recursion) and Chunk (Prefill chunk-parallel) | Qwen3.5 | recurrent_gated_delta_rule, chunk_gated_delta_rule |
13+| Linear Attention | KDA | Kernel Delta Attention. Implements kernelized recursion state updates of the delta rule based on the blockwise kernel method. Completes gate accumulation, QK coefficient computation, state recursion updates, and attention output through chunk splitting, and also outputs the final state and optional backward intermediates | Kimi-K3 | chunk_kda_fwd |
14+ 
15+---
16+ 
17+## 2 External Interfaces
18+ 
19+The following table lists all operators of the attention layer. Operator names are consistent with operator directory names.
20+ 
21+| Operator name | Operator description | Sample |
22+| :--- | :--- | :--- |
23+| [flash_attention_score](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_fusion_attention.md) | FlashAttention forward computation in the training scenario. Executes Dropout(Softmax(Mask(scale·Q·K^T+pse)))·V. Supports pse/mask/dropout/RoPE | [test_aclnn_flash_attention_score.cpp](./flash_attention_score/examples/test_aclnn_flash_attention_score.cpp) |
24+| [flash_attention_score_grad](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_flash_attention_score_grad.md) | Backward gradient computation of FlashAttentionScore. Computes the gradients of Q/K/V | [test_aclnn_flash_attention_score_grad_v4_fp16.cpp](./flash_attention_score_grad/examples/test_aclnn_flash_attention_score_grad_v4_fp16.cpp) |
25+| [flash_attn](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_fusion_attention.md) + [flash_attn_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_flash_attn_metadata.md) | Unified training/inference FlashAttention interface. Supports BSND/BNSD/TND layouts, paged attention, variable length, and the sliding window. The metadata operator pre-computes the tiling split scheme for it | [test_flash_attn.py](./flash_attn/tests/pytests/test_flash_attn.py) |
26+| [fused_infer_attention_score](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_fused_infer_attention_score.md) | Fused inference attention. Uniformly supports the full (Prompt) and incremental (Incre) scenarios. Supports quantization/dequantization/paging/shared prefix/RoPE | [test_aclnn_fused_infer_attention_score_v5.cpp](./fused_infer_attention_score/examples/arch35/test_aclnn_fused_infer_attention_score_v5.cpp) |
27+| [incre_flash_attention](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_incre_flash_attention.md) | Incremental attention for autoregressive inference. The S axis of query is fixed to 1, and KV is appended to the historical state. No longer supported on 950; flash_attn is recommended | [test_aclnn_incre_flash_attention_v4.cpp](./incre_flash_attention/examples/test_aclnn_incre_flash_attention_v4.cpp) |
28+| [prompt_flash_attention](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_prompt_flash_attention.md) | FlashAttention for the full inference Prefill stage. Supports sparse optimization and INT8 quantization. No longer supported on 950; flash_attn is recommended | [test_aclnn_prompt_flash_attention_v3.cpp](./prompt_flash_attention/examples/test_aclnn_prompt_flash_attention_v3.cpp) |
29+| [quant_flash_attn](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_flash_attn.md) + [quant_flash_attn_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_flash_attn.md) | MXFP8 quantized FlashAttention. Q/K/V are FP8_E4M3, softmax is computed in FP32, and the output is BF16. The metadata operator pre-computes the tiling split scheme for it. quant_flash_attn is recommended for the quantization scenarios of prompt_flash_attention / incre_flash_attention | [test_quant_flash_attn_func_rdv.py](./quant_flash_attn/tests/pytest/qfa_mxfp8_test/test_quant_flash_attn_func_rdv.py) |
30+| [lightning_indexer](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_lightning_indexer.md) | Selects the Top-K sparse position index of each token based on Q@K correlation + ReLU + weight W, serving as the preprocessing of SparseFlashAttention | [test_aclnn_lightning_indexer.cpp](./lightning_indexer/examples/test_aclnn_lightning_indexer.cpp) |
31+| [lightning_indexer_grad](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_lightning_indexer_grad.md) | Backward operator of lightning_indexer. Computes the gradients of query/key/weights | [test_aclnn_lightning_indexer_grad.cpp](./lightning_indexer_grad/examples/test_aclnn_lightning_indexer_grad.cpp) |
32+| [quant_lightning_indexer](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_lightning_indexer.md) | Quantized Lightning Indexer. query/key are INT8/FP8/HIF8 quantized inputs, with 8-bit storage and 8-bit computation | [test_aclnn_quant_lightning_indexer.cpp](./quant_lightning_indexer/examples/test_aclnn_quant_lightning_indexer.cpp) |
33+| [sparse_flash_attention](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_attention.md) | Sparse attention for large-sequence inference. Selects key KV through sparse_indices and performs Softmax(QK^T/√d)V | [test_aclnn_sparse_flash_attention_v2.cpp](./sparse_flash_attention/examples/test_aclnn_sparse_flash_attention_v2.cpp) |
34+| [sparse_flash_attention_grad](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_attention_grad.md) | Backward operator of sparse_flash_attention. Rearranges KV according to sparseIndices and then computes dQ/dK/dV | [test_aclnn_sparse_flash_attention_grad.cpp](./sparse_flash_attention_grad/examples/test_aclnn_sparse_flash_attention_grad.cpp) |
35+| [kv_quant_sparse_flash_attention](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_kv_quant_sparse_flash_attention.md) | Sparse attention supporting Per-Token-Head-Tile-128 quantized KV input. Only the IR graph mode operator, with no aclnn C interface | [test_npu_kv_quant_sparse_flash_attention.py](./kv_quant_sparse_flash_attention/examples/test_npu_kv_quant_sparse_flash_attention.py) |
36+| [mla_prolog](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mla_prolog.md) | MLA preprocessing: downsampling + RmsNorm + upsampling + RoPE. Generates query/query_rope and updates kv_cache/kr_cache | [test_aclnn_mla_prolog_nq_bsh.cpp](./mla_prolog/examples/test_aclnn_mla_prolog_nq_bsh.cpp) |
37+| [mla_prolog_v2](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mla_prolog_v2.md) | MLA preprocessing v2. Adds full quantization scenario support (both token_x/weight can be quantized) and outputs dequant_scale_q_nope | [test_aclnn_mla_prolog_v2_weight_nz_fq.cpp](./mla_prolog_v2/examples/test_aclnn_mla_prolog_v2_weight_nz_fq.cpp) |
38+| [mla_prolog_v3](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mla_prolog_v3.md) | MLA preprocessing v3. Adds scale correction factors (αq/αkv), query_norm output, and per-token/per-group quantization modes | [test_aclnn_mla_prolog_v3_weight_nz.cpp](./mla_prolog_v3/examples/test_aclnn_mla_prolog_v3_weight_nz.cpp) |
39+| [compressor](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_compressor.md) | KV compression preprocessing. Compresses the KV cache of every 4 or 128 tokens into 1 token and outputs cmp_kv | [test_aclnn_compressor.cpp](./compressor/examples/test_aclnn_compressor.cpp) |
40+| [compressor_grad](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_compressor_grad.md) | Backward operator of compressor. Computes the gradients of X, W^KV/W^Gate, and Ape | [test_aclnn_compressor_grad.cpp](./compressor_grad/examples/test_aclnn_compressor_grad.cpp) |
41+| [quant_compressor](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_compressor.md) | Quantized Compressor. Inputs x/W^KV/W^Gate are HIFLOAT8, with hardware-native Matmul followed by dequantization | [test_aclnn_quant_compressor.cpp](./quant_compressor/examples/arch35/test_aclnn_quant_compressor.cpp) |
42+| [lightning_indexer_v2](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_lightning_indexer_v2.md) + [lightning_indexer_v2_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_lightning_indexer_v2_metadata.md) | V2 sparse index selection. Supports maskMode/cmpRatio and selects sparse positions based on Q@K correlation + Top-K. The metadata operator generates the load balancing scheme for it | [test_aclnn_lightning_indexer_v2.cpp](./lightning_indexer_v2/examples/test_aclnn_lightning_indexer_v2.cpp) |
43+| [quant_lightning_indexer_v2](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_lightning_indexer_v2.md) + [quant_lightning_indexer_v2_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_lightning_indexer_v2_metadata.md) | Quantized V2 sparse index. query/key are INT8/FP8/HIF8 quantized inputs. The metadata operator generates the load balancing scheme for it | [test_aclnn_quant_lightning_indexer_v2.cpp](./quant_lightning_indexer_v2/examples/test_aclnn_quant_lightning_indexer_v2.cpp) |
44+| [sparse_flash_mla](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla.md) + [sparse_flash_mla_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla_metadata.md) | MLA sparse attention main operator. Supports dual-path (ori original KV + cmp compressed KV) sparse attention and SWA/CSA/HCA masks. The metadata operator generates the load balancing scheme for it | [test_aclnn_sparse_flash_mla.cpp](./sparse_flash_mla/examples/test_aclnn_sparse_flash_mla.cpp) |
45+| [sparse_flash_mla_grad](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla_grad.md) + [sparse_flash_mla_grad_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla_grad_metadata.md) | Backward operator of sparse_flash_mla in the training scenario. The metadata operator generates a task list containing Batch/Head/Q/K block indices for it | [test_aclnn_sparse_flash_mla_grad.cpp](./sparse_flash_mla_grad/examples/test_aclnn_sparse_flash_mla_grad.cpp) |
46+| [sparse_flash_mla_softmax_l1_norm](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla_softmax_l1_norm.md) + [sparse_flash_mla_softmax_l1_norm_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_sparse_flash_mla_softmax_l1_norm_metadata.md) | Computes the Softmax L1Norm result of the sparse_flash_mla attention, paired with the KL Loss Grad backward. The metadata operator generates the preceding core-splitting scheme for it | [test_aclnn_sparse_flash_mla_softmax_l1_norm.cpp](./sparse_flash_mla_softmax_l1_norm/examples/test_aclnn_sparse_flash_mla_softmax_l1_norm.cpp) |
47+| [quant_sparse_flash_mla](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_sparse_flash_mla.md) + [quant_sparse_flash_mla_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_quant_sparse_flash_mla_metadata.md) | Fully quantized MLA sparse attention. Q/KV are HIFLOAT8 per-tensor quantized. The metadata operator generates the load balancing scheme for it | [test_aclnn_quant_sparse_flash_mla.cpp](./quant_sparse_flash_mla/examples/test_aclnn_quant_sparse_flash_mla.cpp) |
48+| [mixed_quant_sparse_flash_mla](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mixed_quant_sparse_flash_mla.md) + [mixed_quant_sparse_flash_mla_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mixed_quant_sparse_flash_mla_metadata.md) | Mixed quantized MLA sparse attention. Supports per-token-group quantized KV. The metadata operator generates the load balancing scheme for it | [test_aclnn_mixed_quant_sparse_flash_mla.cpp](./mixed_quant_sparse_flash_mla/examples/test_aclnn_mixed_quant_sparse_flash_mla.cpp) |
49+| [mixed_quant_flash_attn](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mixed_quant_flash_attn.md) + [mixed_quant_flash_attn_metadata](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_mixed_quant_flash_attn_metadata.md) | Pseudo-quantized FlashAttention. Q remains BF16/FP16 without quantization, and KV is stored in MXFP4/HiFP4 quantization and dequantized for recovery. The metadata operator pre-computes the tiling split scheme for it | [test_aclnn_mixed_quant_flash_attn_metadata.cpp](./mixed_quant_flash_attn_metadata/examples/test_aclnn_mixed_quant_flash_attn_metadata.cpp) |
50+| [recurrent_gated_delta_rule](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_recurrent_gated_delta_rule.md) | Recursively updates the hidden state S_t at each time step and outputs o_t. Suitable for the Decode stage. Supports variable step length and state indexing | [test_aclnn_recurrent_gated_delta_rule.cpp](./recurrent_gated_delta_rule/examples/test_aclnn_recurrent_gated_delta_rule.cpp) |
51+| [chunk_gated_delta_rule](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_chunk_gated_delta_rule.md) | Computes the Gated Delta Rule in parallel by sequence chunks. More efficient for long contexts than the Recurrent version. Suitable for the Prefill stage. Outputs o_t at each step and the final state S_L | [test_aclnn_chunk_gated_delta_rule.cpp](./chunk_gated_delta_rule/examples/test_aclnn_chunk_gated_delta_rule.cpp) |
52+| [chunk_kda_fwd](https://www.hiascend.com/document/detail/zh/Pytorch/latest/apiref/customapi/docs/zh/custom_APIs/torch_npu/torch_npu-npu_chunk_kda_fwd.md) | KDA chunked forward computation. Implements parallelism by kernelizing the delta rule recursion state updates through the blockwise kernel method. Suitable for the Prefill stage | [test_chunk_kda_fwd.py](./chunk_kda_fwd/tests/pytest/test_chunk_kda_fwd.py) |
53+ 
54+---
55+ 
56+## 3 Metadata Operator Mechanism
57+ 
58+### 3.1 Overview
59+ 
60+A Metadata operator is an AICPU preprocessing operator paired with a main operator (named with the `_metadata` suffix). It runs on the AI CPU on the device side, computes a load-balanced core-splitting scheme based on the input sequence metadata (such as `cuSeqlens`, `seqused`, batch_size, and max_seqlen), and outputs a fixed-shape metadata tensor for the main operator to consume.
61+ 
62+The typical invocation flow consists of two stages. The Metadata operator runs first to generate the core-splitting scheme, and the main operator then consumes the metadata to perform attention computation:
63+ 
64+```
65+Step 1: Call the Metadata operator (AICPU) Step 2: Call the main operator (AICore)
66+ 
67+ Input: sequence metadata Input: Q, K, V + metadata
68+ (cu_seqlens, seqused, ...) |
69+ | v
70+ v +-------------+
71+ +-------------+ metadata | Main |
72+ | Metadata | (core-splitting | Attention |
73+ | operator | scheme) ----------> | computing |
74+ +-------------+ +-------------+
75+ |
76+ v
77+ Output: attn_out
78+```
79+ 
80+### 3.2 Core Problem Solved: Tiling Sinking
81+ 
82+#### 3.2.1 What Is Tiling Sinking
83+ 
84+In the traditional mode, core splitting of the main operator is completed in the host-side `GetWorkspaceSize` (tiling) stage. Tiling sinking is a graph execution mode that sinks the core-splitting computation originally performed in the host-side tiling stage to the device side (AICPU). The Metadata operator is the specific implementation of tiling sinking.
85+ 
86+#### 3.2.2 Why Tiling Sinking Is Needed
87+ 
88+| Problem | Description |
89+| :--- | :--- |
90+| **The host side cannot read the values of device tensors** | In the variable-length sequence (TND) scenario, tensors such as `cuSeqlens` and `seqused` reside on the device side. The host-side tiling stage cannot obtain their actual values and therefore cannot perform precise load balancing |
91+| **Reduce host-side overhead** | Complex core-splitting logic executed on the host side blocks the launch pipeline. After the logic is sunk to the AICPU, the main operator can directly skip the splitting computation in the tiling stage, reducing access to host arrays |
92+| **Support the aclgraph graph mode** | aclgraph requires the execution graph to be determined at static compilation time, while tiling depends on runtime data. Computing the metadata in advance through AICPU operators enables the main operator to execute in graph mode |
93+ 
94+#### 3.2.3 Overall Sinking
95+ 
96+"Overall sinking" means sinking the tiling computation from the host side as a whole into AICPU operators. The main operator no longer computes the splitting scheme itself in the tiling stage but directly consumes the output result of the AICPU operator. The main operator documentation describes the metadata as "the output result of the AICPU operator for tiling sinking".
97+ 
98+### 3.3 Value Dependency and Division of Work Between the Metadata Operator and the Main Operator
99+ 
100+The output tensor of the Metadata operator directly serves as the input tensor of the main operator (the `metadataOptional` parameter). This is a device-side runtime data dependency rather than a host-side attribute dependency. In the tiling stage, the main operator reads the metadata tensor to fill the core-splitting structure, and in the kernel stage, it performs computation according to the core-splitting result. The division of work is as follows:
101+ 
102+| Dimension | Metadata operator (AICPU) | Main operator (AICore) |
103+| :--- | :--- | :--- |
104+| Execution location | Device-side AI CPU | Device-side AI Core / Vector |
105+| Function | Load-balanced core splitting and task division | Attention numerical computation |
106+| Data access | Can read the actual values of device tensors (cuSeqlens, seqused, and so on) | Executes as indicated by the metadata |
107+| Execution timing | Before the main operator (preprocessing operator) | Executes after consuming the metadata |
108+| workspace | Usually 0 | Requires a large workspace |
109+| Pipeline relationship | The two can run in parallel in a pipeline, and metadata can be reused | Consumes the preceding result |
110+ 
111+---
112+ 
113+## 4 Common Concepts
114+ 
115+### 4.1 Layout (Data Layout)
116+ 
117+In attention layer operators, the layout format of input and output tensors is specified by the `inputLayout` parameter. The meanings of the dimension letters are as follows:
118+ 
119+| Letter | Meaning |
120+| :--- | :--- |
121+| B | Batch Size, the batch size of input samples |
122+| S | Sequence Length, the length of a single sequence |
123+| T | Total Tokens, the cumulative sum of sequence lengths of all batches (T = ΣS_i) |
124+| N | Head Num, the number of heads |
125+| D | Head Dim, the hidden dimension of a single head |
126+| H | Head Size, the size of the hidden layer (H = N × D) |
127+ 
128+The layout formats are defined as follows:
129+ 
130+| Format name | shape | Characteristics | Applicable scenarios |
131+| :--- | :--- | :--- | :--- |
132+| BSH | `[B, S, H]`, H = N × D | N and D are merged into one dimension H. 3D format, the simplest | FlashAttentionScore training; PromptFlashAttention / IncreFlashAttention / FusedInferAttentionScore inference (recommended default layout) |
133+| BSND | `[B, S, N, D]` | N and D are split into independent dimensions. 4D format | Common to all training and inference scenarios |
134+| BNSD | `[B, N, S, D]` | Same as BSND, but with the N and S axes swapped | Common to all training and inference scenarios |
135+| TND | `[T, N, D]` | Concatenates the sequences of multiple batches into a 1D tensor, eliminating padding waste; must be used with `cuSeqlens` to mark batch boundaries | Variable-length sequence (VarLen) scenarios; supported for both training and inference |
136+| NTD | `[N, T, D]` | Same as TND, but with the N and T axes swapped | Full quantization scenarios of quantized attention; often used in the `NTD_TND` combination (input NTD, output TND) |
137+| BnBsH | `[BlockNum, BlockSize, H]` | PA scenario KV Cache layout, H = N×D, the simplest 3D format; unavailable when H exceeds 65535 | PageAttention KV Cache storage |
138+| PA_BSND / BnNBsD / PA_BBND | `[BlockNum, BlockSize, N, D]` or `[BlockNum, N, BlockSize, D]` | PA scenario KV Cache layout, with the D dimension stored contiguously. 4D format; performance is better than BnBsH, and this format is recommended first | PageAttention KV Cache storage |
139+| PA_NZ / NZ | `[BlockNum, N, D/D0, BlockSize, D0]` | PA scenario KV Cache layout, with the D dimension split into an NZ fractal structure of D/D0 × D0. 5D; the innermost D0=32 (FP8/INT8) or 16 (FP16/BF16/keyRope); dedicated to quantization scenarios, and **only NZ is supported** for INT8 | PageAttention KV Cache storage in quantization scenarios |
140+ 
141+> **Supplementary notes**:
142+> - **BNSD_BSND**: A special combination layout. When the input is BNSD, the output format is BSND.
143+> - Core difference between PA_BSND and PA_NZ: the D dimension of PA_BSND is stored contiguously (4D), while PA_NZ splits the D dimension into an NZ fractal (5D), which is more Cube-instruction-friendly in quantization scenarios.
144+ 
145+### 4.2 PageAttention (Paged Attention)
146+ 
147+#### 4.2.1 Principle and Design Motivation
148+ 
149+In LLM autoregressive inference, the historical KV Cache must be read every time a token is generated. The traditional KV Cache requires each batch to reserve a segment of contiguous physical memory along the S dimension (sequence length) and pre-pad according to `maxSeqLen`, resulting in memory fragmentation and waste.
150+ 
151+PagedAttention borrows the virtual memory paging mechanism of the operating system, splits the KV Cache into fixed-size blocks stored in a segment of contiguous physical memory, and establishes the mapping from logical sequence order to physical block addresses through `blockTable` (page table). Its core values:
152+ 
153+- **Eliminates memory fragmentation**: allocates blocks on demand, with no need to pre-allocate an entire segment of contiguous memory
154+- **Improves memory utilization**: different batches share the same physical block pool, supporting larger batches or longer sequences
155+- **Logically contiguous, physically discrete**: logically contiguous KV sequences can be scattered across non-contiguous physical blocks
156+ 
157+#### 4.2.2 Working Mechanism
158+ 
159+PagedAttention is enabled only when `blockTable` exists and is valid, and meanwhile key/value are arranged in contiguous memory according to the indices in `blockTable`. After PagedAttention is enabled, the `inputLayout` parameter of key/value becomes invalid, and the layout of the KV Cache is determined by its own shape dimensions.
160+ 
161+**blockTable lookup process** (taking `B=2`, `blockSize=128`, `actualSeqLengthsKv=[256, 512]` as an example):
162+ 
163+`maxBlockNumPerSeq = ceil(512/128) = 4`, blockTable shape = `[2, 4]`:
164+ 
165+```
166+blockTable = [
167+ [ 3, 7, 1, 0], # batch 0: logical block 0 → physical block 3, logical block 1 → physical block 7 (actually only 2 blocks are used)
168+ [ 5, 2, 9, 4], # batch 1: logical block 0 → physical block 5, logical block 1 → physical block 2, logical block 2 → physical block 9, logical block 3 → physical block 4
169+]
170+```
171+ 
172+When accessing tokens 0~255 of batch 0:
173+- token 0~127 → logical block 0 → table lookup gives block_id=3 → read physical `kvCache[3]`
174+- token 128~255 → logical block 1 → table lookup gives block_id=7 → read physical `kvCache[7]`
175+ 
176+Block ids can be allocated out of order and can share the physical pool across batches. Users need to ensure the validity of block ids themselves (not verified on the host side).
177+ 
178+#### 4.2.3 blockTable
179+ 
180+| Item | Content |
181+| :--- | :--- |
182+| Data type | INT32 |
183+| shape | `[B, maxBlockNumPerSeq]` |
184+| First dimension | Equals the number of batches B |
185+| Second dimension | Not smaller than the number of blocks corresponding to the maximum KV sequence length, that is, `ceil(max_KV_S / blockSize)` |
186+| Element meaning | block id, pointing to the n-th block in the KV Cache physical pool |
187+| Constraint | Users must ensure the validity of block ids themselves; not verified on the host side |
188+ 
189+#### 4.2.4 blockSize
190+ 
191+| Item | Content |
192+| :--- | :--- |
193+| Meaning | The maximum number of tokens in each block; a user-defined parameter |
194+| General alignment | FP16/BF16 requires alignment to 16; INT8/HIFLOAT8/FP8_E4M3 requires alignment to 32; INT4/FLOAT4_E2M1 requires alignment to 64 |
195+| Maximum | 1024 (512 in some earlier versions) |
196+| Recommended value | 128 (best compatibility; the mandatory value in most full quantization scenarios) |
197+| Performance impact | Increasing blockSize brings certain performance benefits (reducing table lookups and increasing the data volume per transfer), but increases the memory waste of the tail block |
198+ 
199+#### 4.2.5 Comparison Between PA Mode and Normal Mode
200+ 
201+| Dimension | Normal mode | PA mode |
202+| :--- | :--- | :--- |
203+| KV physical layout | Contiguous along the S axis | Scattered by block, mapped through blockTable |
204+| inputLayout (KV) | Effective | Invalid |
205+| blockTable | Not required | Required |
206+| actualSeqLengthsKv | Optional | Required |
207+| tensorlist | Supported | Not supported |
208+| Left padding | Supported | Not supported |
209+| Shared prefix | Supported | Not supported |
210+| Unequal D | Supported | Not supported |
211+| KV Cache layout | BSH/BSND/BNSD/TND and so on | BnBsH / BnNBsD / NZ |
212+ 
213+#### 4.2.6 Performance Characteristics
214+ 
215+- **Throughput improvement vs. single-run latency decrease**: PA supports larger batches by improving memory utilization, thereby increasing throughput. However, the indirect addressing of blockTable introduces extra table lookup overhead and non-contiguous access, slightly decreasing single-run inference latency
216+- **KV Cache layout affects performance**: BnNBsD performs better than BnBsH (placing the N axis first facilitates fetching contiguous D data by head). BnNBsD is recommended first
217+- **BnBsH limitation**: when KV_N × D exceeds 65535, an error is reported due to hardware instruction constraints; switch to the BnNBsD format
218+ 
219+### 4.3 Mask Mode
220+ 
221+Mask Mode (`sparseMode` / `mask_mode`) controls the masking strategy in attention computation. Parameter names differ across operators but the semantics are the same.
222+ 
223+#### 4.3.1 Value Definitions
224+ 
225+| Value | Name | Meaning | attenMask requirement | preTokens/nextTokens |
226+| :---: | :--- | :--- | :--- | :--- |
227+| 0 | defaultMask | Default mode; if no mask is passed, all is computed; if a mask is passed, causal or band is applied based on preTokens/nextTokens | Optional | Effective |
228+| 3 | rightDownCausal | Lower triangular scenario divided by the bottom-right vertex | Compressed lower triangular matrix (2048×2048) | Not effective |
229+| 4 | band | Computes the part between preTokens and nextTokens, starting from the bottom-right corner | Compressed lower triangular matrix (2048×2048) | Effective |
230+ 
231+#### 4.3.2 rightDownCausal (mode=3)
232+ 
233+Bottom-right vertex causal mask. Taking the bottom-right corner of the attention matrix as the vertex, the lower triangular part (including the diagonal) is retained and the upper triangle is masked. That is, query token i can only attend to key tokens 0~i and cannot see future tokens. Applicable to autoregressive training and inference of standard causal language models (GPT/Llama, and so on).
234+ 
235+```
236+ key: 0 1 2 3 4
237+ ┌───────────────┐
238+ q 0 │ 1 0 0 0 0 │
239+ q 1 │ 1 1 0 0 0 │
240+ q 2 │ 1 1 1 0 0 │
241+ q 3 │ 1 1 1 1 0 │
242+ q 4 │ 1 1 1 1 1 │
243+ └───────────────┘
244+ 1 = computed 0 = masked
245+```
246+ 
247+The passed attenMask is a compressed lower triangular matrix (2048×2048), and preTokens/nextTokens do not take effect.
248+ 
249+#### 4.3.3 band (mode=4)
250+ 
251+Band mask. Starting from the bottom-right corner, a band-shaped region is defined through `preTokens` and `nextTokens`. Query token i can only attend to key tokens within the range `[i - preTokens + 1, i + nextTokens]`. Applicable to the sliding window attention (Sliding Window Attention) scenario.
252+ 
253+Taking `preTokens=3, nextTokens=0` as an example:
254+ 
255+```
256+ key: 0 1 2 3 4
257+ ┌───────────────┐
258+ q 0 │ 1 0 0 0 0 │
259+ q 1 │ 1 1 0 0 0 │
260+ q 2 │ 1 1 1 0 0 │
261+ q 3 │ 0 1 1 1 0 │
262+ q 4 │ 0 0 1 1 1 │
263+ └───────────────┘
264+ 1 = computed 0 = masked
265+ window width = preTokens = 3
266+```
267+ 
268+The passed attenMask is a compressed lower triangular matrix (2048×2048), and `preTokens`/`nextTokens` take effect.
269+ 
270+### 4.4 cuSeqlens & seqused
271+ 
272+These two parameters describe the actual valid length of each batch in variable-length sequences. They are used differently in different layout scenarios: the TND layout concatenates the sequences of multiple batches into a 1D tensor, requiring `cuSeqlens` to mark batch boundaries; in non-TND layouts (BSND/BNSD/BSH, and so on), each batch has an independent S dimension, requiring `seqused` to mark the actual valid length.
273+ 
274+| Item | cuSeqlens | seqused |
275+| :--- | :--- | :--- |
276+| Meaning | Cumulative prefix sum, marking the boundaries of each batch; each element is the accumulated number of valid tokens of the current batch and preceding batches | The number of tokens that each batch actually participates in computation |
277+| shape | `(B+1,)` | `(B,)` |
278+| Data type | INT32 | INT32 |
279+| First element | Fixed to 0 | No requirement |
280+| Numeric constraint | Monotonically non-decreasing; the last element equals T (the total number of tokens) | Not greater than S in the shape |
281+| Numeric relationship | `cuSeqlens[i+1] - cuSeqlens[i] = seqused[i]` | Can be obtained by differencing cuSeqlens |
282+| TND layout | Required (data is compactly arranged by batch, and boundaries need to be marked) | Optional (the total length is expressed by cuSeqlens, and seqused can explicitly override the valid length) |
283+| Non-TND layout | Forbidden to pass | Optional (used to mark the valid length of each batch) |
284+ 
285+#### 4.4.1 Equivalent Padding Effect in the TND Layout
286+ 
287+The TND layout compactly concatenates the sequences of multiple batches into a 1D tensor, requiring no padding. `cuSeqlens` marks the boundaries of each batch, and `seqused` marks the valid length of each batch. The two work together to achieve the equivalent padding effect.
288+ 
289+Taking `B=3` and the valid lengths of batches `2, 1, 3` as an example:
290+ 
291+```
292+TND (compact concatenation, zero waste): BSND (requires padding to maxLen=3):
293+ 
294+┌────┬────┬────┬────┬────┬────┐ ┌────┬────┬─────┐
295+│ t0 │ t1 │ t2 │ t3 │ t4 │ t5 │ │ t0 │ t1 │ pad │ batch0
296+└────┴────┴────┴────┴────┴────┘ ├────┼────┼─────┤
297+ b0: 2 b1:1 b2: 3 │ t2 │pad │ pad │ batch1
298+ ├────┼────┼─────┤
299+cuSeqlens = [0, 2, 3, 6] │ t3 │ t4 │ t5 │ batch2
300+seqused = [2, 1, 3] └────┴────┴─────┘
301+ occupies 9 slots, wastes 3 slots
302+occupies 6 slots, zero waste
303+```
304+ 
305+### 4.5 Sparse Index Parameters (topk / sparse_indices)
306+ 
307+Sparse attention operators select the key KV subset through TopK and perform attention computation only on the selected KV. `topk` controls the number of selections, and `sparse_indices` records the selected position indices. The two form the production-consumption link of Lightning Indexer → Sparse Flash Attention.
308+ 
309+#### 4.5.1 topk (Number of TopK Selections)
310+ 
311+The core idea of sparse attention is not to perform attention computation on all KV. Instead, the most relevant KV positions for each query token are filtered out first through correlation scoring, and `Softmax(QK^T)V` is executed only on this part of KV, greatly reducing the computation load. This filtering process is completed in the Lightning Indexer operator.
312+ 
313+`topk` is the quantity parameter that controls how many to select: the larger the topk, the more KV participates in computation, and the closer the attention result is to the precision of dense attention, but the computation and memory access overhead also increase; the smaller the topk, the faster the computation, but key information may be missed, resulting in precision degradation. Weigh performance against precision based on the model precision requirements and sequence length. `topk` directly determines the size of the last dimension of the downstream `sparse_indices` output.
314+ 
315+Taking the sequence length `KV_S=8192` as an example:
316+ 
317+```
318+topk=2048: selects 25% of the KV, computation reduced by 75%
319+topk=512: selects 6% of the KV, computation reduced by 94%
320+```
321+ 
322+#### 4.5.2 sparse_indices (Sparse Position Indices)
323+ 
324+`sparse_indices` is the output of the Lightning Indexer and the input of Sparse Flash Attention / Sparse Flash MLA, recording the KV position indices that each query token should attend to. It passes the topk filtering result to the attention main operator in the form of an INT32 index array. The main operator discretely gathers the KV at the corresponding positions based on it and then performs attention computation, instead of performing dense computation on all KV.
325+ 
326+In the index array, valid positions store the actual block indices of the KV, and invalid positions are filled with -1. The consumer side convention is that valid values are arranged in the first half and invalid values in the second half, so that the main operator only needs to process the first `topk` valid entries and skip the invalid part at the tail.
327+ 
328+**Production-consumption link**:
329+ 
330+```
331+lightning_indexer(query, key, weights, topk) → sparse_indices (output)
332+ ↓
333+sparse_flash_attention(query, key, value, sparse_indices) → attn_out (input)
334+```
335+ 
336+Lightning Indexer computes `ReLU(Q@K^T)` to obtain correlation scores, multiplies them by weights and aggregates along the group, and selects and outputs the first `topk` indices through TopK. Sparse Flash Attention discretely gathers the KV accordingly and then executes `Softmax(QK^T)@V`. The entire link decouples selecting which KV and performing attention on the selected KV into two independent operators, facilitating separate optimization and flexible combination.
337+ 
338+### 4.6 Quantization Scenario Concepts
339+ 
340+According to whether Q/K/V are quantized and the quantization granularity, attention layer operators fall into the non-quantized, pseudo-quantized, and fully quantized scenarios.
341+ 
342+#### 4.6.1 Quantization Type Classification
343+ 
344+| Quantization type | Q | K/V | Dequantization parameters | Description |
345+| :--- | :--- | :--- | :--- | :--- |
346+| Non-quantized | FP16/BF16 | FP16/BF16 | None | Q/K/V are all high-precision and directly participate in computation |
347+| Pseudo-quantized | FP16/BF16 | INT8/FP8/HIF8 and so on | `antiquantScale` / `antiquantOffset` | Only K/V is quantized for storage and dequantized within the operator before computation; Q and the output remain high-precision |
348+| Fully quantized | INT8/FP8/HIF8 and so on | INT8/FP8/HIF8 and so on | `dequantScaleQuery` + `antiquantScale` | Q/K/V are all quantized inputs, all dequantized within the operator before computation |
349+ 
350+#### 4.6.2 Dequantization Parameter System
351+ 
352+| Parameter | Function | Applicable scenario |
353+| :--- | :--- | :--- |
354+| `antiquantScale` / `antiquantOffset` | Dequantization factor and offset for K/V pseudo-quantization | Pseudo-quantization (symmetric/asymmetric) |
355+| `keyAntiquantScale` / `valueAntiquantScale` | Separate dequantization factors for K/V (K/V can use different scales) | Pseudo-quantization/full quantization (K/V separate mode) |
356+| `dequantScaleQuery` | Factor for dequantizing Q | **Dedicated to full quantization** |
357+| `deqScale1` / `quantScale1` | Dequantizes the QK result / quantizes P | INT8 online quantization (Q high-precision scenario) |
358+| `deqScale2` / `quantScale2` / `quantOffset2` | Dequantizes the PV result / quantizes the output and offset | INT8 online quantization (output INT8 scenario) |
359+ 
360+#### 4.6.3 GQA Full Quantization (FP8)
361+ 
362+In the GQA full quantization scenario, Q/K/V are all in FLOAT8_E4M3FN format, and high-precision computation is restored within the operator through FLOAT32 dequantization factors.
363+ 
364+| Item | Content |
365+| :--- | :--- |
366+| Q/K/V data type | FLOAT8_E4M3FN |
367+| Dequantization factor dtype | FLOAT32 |
368+| Quantization granularity | `queryQuantMode=3` (per-token plus per-head), `keyAntiquantMode=3` (per-token plus per-head), `valueAntiquantMode=2` (per-tensor plus per-head) |
369+| HeadDim constraint | Only 128 is supported |
370+| blockSize constraint | Only 128 is supported |
371+| Typical operator | fused_infer_attention_score |
372+ 
373+#### 4.6.4 MXFP8 Full Quantization
374+ 
375+MXFP8 (Microscaling FP8) is a quantization format combining FP8_E4M3FN data with FLOAT8_E8M0 scaling factors, where every 32 elements share one E8M0 scale.
376+ 
377+| Item | Content |
378+| :--- | :--- |
379+| Q/K/V data type | FLOAT8_E4M3FN |
380+| Scaling factor dtype | FLOAT8_E8M0 |
381+| Quantization granularity | `queryQuantMode/keyAntiquantMode=6` (per-token-group), `valueAntiquantMode=8` (per-channel-group) |
382+| scale shape | `(1, B, N, >=KV_S, D/32)`, every 32 D elements share one E8M0 scale |
383+| HeadDim constraint | Only 64 or 128 is supported |
384+| blockSize constraint | Supports 64/128/256/512/1024 |
385+| Typical operator | quant_flash_attn, fused_infer_attention_score |
386+ 
387+### 4.7 Batch Consistency
388+ 
389+Batch consistency (deterministic computation) refers to the ability of an operator to keep computation results consistent when executing under different batch sizes and different batch combinations. By default, the computation results of some operators (such as compressor / quant_compressor) may be related to the position of the input batch in the overall data and the batch division method. That is, the same group of data may produce slight numerical differences under different batch divisions.
390+ 
391+Configuring `ACL_OPT_DETERMINISTIC=3` through `aclrtSetSysParamOpt()` enables the batch consistency mode. After enabling, the operator internally adopts a deterministic computation strategy, ensuring that computation results are independent of the batch size and position.
392+ 
393+| Item | Content |
394+| :--- | :--- |
395+| Meaning | Computation results are independent of the batch size and the position of the batch in the data |
396+| Enabling method | `aclrtSetSysParamOpt(ACL_OPT_DETERMINISTIC, 3)` |
397+| Applicable operators | compressor, quant_compressor, and so on |
398+| Default state | Disabled (non-deterministic) |
399+| Cost | Enabling may cause certain performance degradation |
400+ 
401+---
Aattention/attention_update/README_en.md+93-0
@@ -0,0 +1,93 @@
1+# AttentionUpdate
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+|Product | Supported |
8+|:----------------------------|:-----------:|
9+|<term>Ascend 950PR&950DT products</term>| √ |
10+|<term>Atlas A3 products</term>| √ |
11+|<term>Atlas A2 products</term>| √ |
12+|<term>Atlas 200I/500 A2 inference products</term>| × |
13+|<term>Atlas inference products</term>| × |
14+|<term>Atlas training products</term>| × |
15+ 
16+## Function Description
17+ 
18++ Operator function: Updates the intermediate results lse and localOut output by the PA operators in each SP domain into the global results.
19++ Computation formula: Inputs $lse_i$ and $O_i$, and output $O$.
20+ 
21+$$
22+lse_{max} = \text{max}lse_i
23+$$
24+ 
25+$$
26+lse = \sum_i \text{exp}(lse_i - lse_{max})
27+$$
28+ 
29+$$
30+lse_m = lse_{max} + \text{log}(lse)
31+$$
32+ 
33+$$
34+O = \sum_i O_i \cdot \text{exp}(lse_i - lse_m)
35+$$
36+ 
37+## Parameter Description
38+ 
39+<table class="tg"><thead>
40+ <tr>
41+ <th class="tg-0pky">Parameter name</th>
42+ <th class="tg-0pky">Input/Output/Attribute</th>
43+ <th class="tg-0pky">Description</th>
44+ <th class="tg-0pky">Data type</th>
45+ <th class="tg-0pky">Data format</th>
46+ </tr></thead>
47+<tbody>
48+ <tr>
49+ <td class="tg-0pky">lse<sub>i</sub></td>
50+ <td class="tg-0pky">Input</td>
51+ <td class="tg-0pky">Local lse of each SP domain.</td>
52+ <td class="tg-0pky">FLOAT32</td>
53+ <td class="tg-0pky">ND</td>
54+ </tr>
55+ <tr>
56+ <td class="tg-0pky">O<sub>i</sub></td>
57+ <td class="tg-0pky">Input</td>
58+ <td class="tg-0pky">Local attentionout of each SP domain.</td>
59+ <td class="tg-0pky">FLOAT32, FLOAT16, BFLOAT16</td>
60+ <td class="tg-0pky">ND</td>
61+ </tr>
62+ <tr>
63+ <td class="tg-0pky">lse<sub>m</sub></td>
64+ <td class="tg-0pky">Output</td>
65+ <td class="tg-0pky">Updated global lse.</td>
66+ <td class="tg-0pky">FLOAT32</td>
67+ <td class="tg-0pky">ND</td>
68+ </tr>
69+ <tr>
70+ <td class="tg-0pky">O</td>
71+ <td class="tg-0pky">Output</td>
72+ <td class="tg-0pky">Updated global attentionout.</td>
73+ <td class="tg-0pky">FLOAT32, FLOAT16, BFLOAT16</td>
74+ <td class="tg-0pky">ND</td>
75+ </tr>
76+</tbody></table>
77+ 
78+## Constraint Description
79+ 
80+* For <term>Atlas A2 products</term> and <term>Atlas A3 products</term>: O<sub>i</sub> and O of the FLOAT32, FLOAT16, and BFLOAT16 types are supported.
81+* For <term>Ascend 950PR&950DT products</term>: O<sub>i</sub> and O of the FLOAT32, FLOAT16, and BFLOAT16 types are supported, and O<sub>i</sub> and O must have the same data type.
82+* The value range of the sequence parallelism degree sp:
83+ * <term>Atlas A2 products</term> and <term>Atlas A3 products</term>: [1, 128].
84+ * <term>Ascend 950PR&950DT products</term>: [1, 16].
85+* headDim ranges from [8, 512] and must be a multiple of 8.
86+* localout supports non-contiguous tensors.
87+* Empty input with T=0 is supported. In this case, the output is empty and no computation is performed.
88+ 
89+## Invocation Description
90+ 
91+| Invocation method | Sample code | Description |
92+| ----------- | ------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- |
93+| aclnn interface | [test_aclnn_attention_update](./examples/test_aclnn_attention_update.cpp) | Invokes the AttentionUpdate operator through the [aclnnAttentionUpdate](docs/aclnnAttentionUpdate_en.md) interface. |
Aattention/attention_worker_combine/README_en.md+102-0
@@ -0,0 +1,102 @@
1+# AttentionWorkerCombine
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+| Product | Supported |
8+| :---- | :----: |
9+|<term>Ascend 950PR&950DT products</term>| √ |
10+|<term>Atlas A3 products</term>| √ |
11+|<term>Atlas A2 products</term>| √ |
12+|<term>Atlas 200I/500 A2 inference products</term>| × |
13+|<term>Atlas inference products</term>| × |
14+|<term>Atlas training products</term>| × |
15+ 
16+## Function Description
17+ 
18++ Operator function: Fuses the attention token data processed by multiple compute units, weights the results with expert weights, outputs the final attention fusion result, and updates the layer ID.
19+ 
20+## Parameter Description
21+ 
22+<table class="tg"><thead>
23+ <tr>
24+ <th class="tg-0pky">Parameter name</th>
25+ <th class="tg-0pky">Input/Output/Attribute</th>
26+ <th class="tg-0pky">Description</th>
27+ <th class="tg-0pky">Data type</th>
28+ <th class="tg-0pky">Data format</th>
29+ </tr></thead>
30+<tbody>
31+ <tr>
32+ <td class="tg-0pky">schedule_context</td>
33+ <td class="tg-0pky">Input</td>
34+ <td class="tg-0pky">Contains scheduling context information.</td>
35+ <td class="tg-0pky">INT8</td>
36+ <td class="tg-0pky">ND</td>
37+ </tr>
38+ <tr>
39+ <td class="tg-0pky">expert_scales</td>
40+ <td class="tg-0pky">Input</td>
41+ <td class="tg-0pky">Expert weights. The type is FLOAT8_E8M0 in the MXFP dequantization scenario.</td>
42+ <td class="tg-0pky">FLOAT, FLOAT8_E8M0</td>
43+ <td class="tg-0pky">ND</td>
44+ </tr>
45+ <tr>
46+ <td class="tg-0pky">layer_id</td>
47+ <td class="tg-0pky">Input</td>
48+ <td class="tg-0pky">Current model layer ID.</td>
49+ <td class="tg-0pky">INT32</td>
50+ <td class="tg-0pky">ND</td>
51+ </tr>
52+ <tr>
53+ <td class="tg-0pky">y</td>
54+ <td class="tg-0pky">Output</td>
55+ <td class="tg-0pky">Final attention merge result.</td>
56+ <td class="tg-0pky">FLOAT16, BFLOAT16</td>
57+ <td class="tg-0pky">ND</td>
58+ </tr>
59+ <tr>
60+ <td class="tg-0pky">next_layer_id</td>
61+ <td class="tg-0pky">Output</td>
62+ <td class="tg-0pky">ID of the next layer to process.</td>
63+ <td class="tg-0pky">INT32</td>
64+ <td class="tg-0pky">ND</td>
65+ </tr>
66+ <tr>
67+ <td class="tg-0pky">hidden_size</td>
68+ <td class="tg-0pky">Attribute</td>
69+ <td class="tg-0pky">Hidden dimension size of token_data, used to determine the size of the second dimension of the output y. Required attribute.</td>
70+ <td class="tg-0pky">Int</td>
71+ <td class="tg-0pky">-</td>
72+ </tr>
73+ <tr>
74+ <td class="tg-0pky">token_dtype</td>
75+ <td class="tg-0pky">Attribute</td>
76+ <td class="tg-0pky">Specifies the original precision type of the token data in schedule_context. 0 indicates FLOAT16, and 1 indicates BFLOAT16; <term>Ascend 950PR&950DT products</term> supports MXFP type dequantization. 2 indicates MXFP8_E5M2, 3 indicates MXFP8_E4M3FN, and 4 indicates MXFP4_E2M1.</td>
77+ <td class="tg-0pky">Int</td>
78+ <td class="tg-0pky">-</td>
79+ </tr>
80+ <tr>
81+ <td class="tg-0pky">need_schedule</td>
82+ <td class="tg-0pky">Attribute</td>
83+ <td class="tg-0pky">Specifies whether to wait until the token data is filled before execution. 0 indicates not waiting, and 1 indicates waiting.</td>
84+ <td class="tg-0pky">Int</td>
85+ <td class="tg-0pky">-</td>
86+ </tr>
87+</tbody></table>
88+ 
89+## Constraint Description
90+ 
91+* In the non-quantized scenario, expert_scales is a FLOAT 2D tensor, [BatchSize, K], where K ranges from [1, 64].
92+* In the MXFP dequantization scenario, expert_scales is an E8M0 3D tensor with the shape [BatchSize, K+1, S], where S = Ceil(H/32), and then S is rounded up to an even number.
93+* In the MXFP dequantization scenario, token_dtype=2/3/4 is supported; R/H must be positive, K ranges from [1, 64], and need_schedule is 0 or 1.
94+* y is a 2D tensor, [BatchSize, HiddenSize]. That is, the second dimension is determined by the attribute hidden_size.
95+ 
96+## Invocation Description
97+ 
98+| Invocation method | Sample code | Description |
99+| ---- | ---- | ---- |
100+| Graph mode invocation | [test_geir_attention_worker_combine.cpp](examples/test_geir_attention_worker_combine.cpp) | Invokes the AttentionWorkerCombine operator by constructing the graph through the [operator IR](op_graph/attention_worker_combine_proto.h). |
101+| aclnn interface | [test_aclnn_attention_worker_combine.cpp](examples/test_aclnn_attention_worker_combine.cpp) | Invokes the AttentionWorkerCombine operator through the [aclnnAttentionWorkerCombine](docs/aclnnAttentionWorkerCombine_en.md) interface. |
102+| PyTorch interface | [torchapi_attention_worker_combine.md](docs/torchapi_attention_worker_combine.md) | Invoked through `cann_ops_transformer.attention_worker_combine`, supporting the non-quantized and MXFP dequantization scenarios. |
Aattention/attention_worker_scheduler/README_en.md+103-0
@@ -0,0 +1,103 @@
1+# AttentionWorkerScheduler
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+| Product | Supported |
8+| :----------------------------------------------------------- | :------: |
9+|<term>Ascend 950PR&950DT products</term>| √ |
10+| <term>Atlas A3 products</term> | √ |
11+| <term>Atlas A2 products</term> | √ |
12+| <term>Atlas 200I/500 A2 inference products</term> | × |
13+| <term>Atlas inference products</term> | √ |
14+| <term>Atlas training products</term> | √ |
15+ 
16+## Function Description
17+ 
18+- Operator function: A data scanning operator on the Attention side in the scenario where Attention and FFN are deployed separately. This operator receives the output data from the FFNToAttention operator and scans the data step by step to ensure that the data is ready.
19+ 
20+ **This operator is not recommended for standalone use. It is recommended to use it together with the FFNToAttention and AttentionWorkerCombine operators to form a complete workflow.**
21+ 
22+ 1. Receives the data sent by the FFNToAttention operator. The data is stored in the ScheduleContext structure. The structure contains the CommonArea, ControlArea, AttentionArea, and FfnArea fields. This interface involves CommonArea (used to store configuration information such as session_num, micro_batch_num, micro_batch_size, and selected_expert_num), ControlArea (used by upper-layer processes to control whether to exit), and the AttentionArea field (responsible for managing the core data buffers and status information required during operator computation, where token_info_buf stores input-related data information).
23+ 
24+ 2. Reads the flag information stored in token_info_buf in the ScheduleContext.AttentionArea field to check whether the communication data is ready.
25+ 
26+ 3. After all data is ready, the data can be used by the AttentionWorkerCombine operator.
27+ 
28+- Computation formula:
29+ 
30+$$
31+\text{Initialize:} \quad \text{ready\_count} = 0, \quad \text{flag\_num} = \text{micro\_batch\_size} \times \text{selected\_expert\_num}
32+$$
33+ 
34+$$
35+\text{Check if run\_flag is 0:}
36+\quad \text{if run\_flag} = 0, \quad \text{exit and log}
37+$$
38+ 
39+$$
40+\text{Loop:} \quad \text{while run\_flag} \neq 0:
41+\quad \text{ready\_count} = \sum\_{i=1}^{\text{flag\_num}} \mathbf{1}\_{\{ \text{flag}[i] = 1 \}}; \quad \text{if ready\_count} = \text{flag\_num}, \quad \text{break}
42+$$
43+ 
44+$$
45+\text{Reset flags:}
46+\quad \text{flag}[i] = 0 \quad \text{for} \quad i = 1, 2, \dots, \text{flag\_num}
47+$$
48+ 
49+$$
50+\text{Set micro\_batch\_id:} \quad \text{micro\_batch\_id} = (\text{micro\_batch\_id} + 1) \% \text{micro\_batch\_num}
51+$$
52+ 
53+Note: micro_batch_size, selected_expert_num, run_flag, and micro_batch_id are parameters of the input ScheduleContext structure. The structure information is demonstrated in the invocation example.
54+ 
55+## Parameter Description
56+ 
57+<table style="undefined;table-layout: fixed; width: 1005px"><colgroup>
58+ <col style="width: 170px">
59+ <col style="width: 170px">
60+ <col style="width: 352px">
61+ <col style="width: 213px">
62+ <col style="width: 100px">
63+ </colgroup>
64+ <thead>
65+ <tr>
66+ <th>Parameter name</th>
67+ <th>Input/Output/Attribute</th>
68+ <th>Description</th>
69+ <th>Data type</th>
70+ <th>Data format</th>
71+ </tr></thead>
72+ <tbody>
73+ <tr>
74+ <td>schedule_context</td>
75+ <td>Input</td>
76+ <td>Data to be processed received on the Attention side, indicating the input ScheduleContext information. Empty tensors are not supported, and the shape is fixed to (1024).</td>
77+ <td>INT8</td>
78+ <td>ND</td>
79+ </tr>
80+ <tr>
81+ <td>schedule_context</td>
82+ <td>Output</td>
83+ <td>Processed ScheduleContext information, with the same shape, dtype, and format as the input.</td>
84+ <td>INT8</td>
85+ <td>ND</td>
86+ </tr>
87+ </tbody>
88+ </table>
89+ 
90+## Constraint Description
91+ 
92+- The input `schedule_context` is a 1D INT8 tensor whose data contains at least `sizeof(ScheduleContext)` bytes (the current structure is 1024 bytes).
93+- `schedule_mode` must be `1` (Attention); `micro_batch_num`, `micro_batch_size`, and `selected_expert_num` must all be greater than 0.
94+- `token_info_buf` and `token_data_buf` must both be non-empty addresses; `token_info_buf_size` must be at least `micro_batch_num * (sizeof(AttentionDataDesc) + sizeof(int32_t) * micro_batch_size * selected_expert_num)` bytes.
95+- `micro_batch_id` must be in the range `[0, micro_batch_num)`.
96+- When `run_flag` is `0`, the operator exits; when it is a non-zero value, the operator waits until all flags of the current batch are ready and then continues processing.
97+- If an integer multiplication overflows when computing the required buffer length, the operator returns an internal error.
98+ 
99+## Invocation Description
100+ 
101+| Invocation method | Sample code | Description |
102+| ----------- | ------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- |
103+| aclnn interface | [test_aclnn_inplace_attention_worker_scheduler](./examples/test_aclnn_inplace_attention_worker_scheduler.cpp) | Invokes the AttentionWorkScheduler operator through the [aclnnInplaceAttentionWorkerScheduler](docs/aclnnInplaceAttentionWorkerScheduler_en.md) interface. |
Aattention/block_attn_res_prepare/README_en.md+186-0
@@ -0,0 +1,186 @@
1+# BlockAttnResPrepare
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+| Product | Supported |
8+| ---- | :------: |
9+| <term>Ascend 950PR&950DT products</term> | √ |
10+| <term>Atlas A3 products</term> | × |
11+| <term>Atlas A2 products</term> | × |
12+| <term>Atlas 200I/500 A2 inference products</term> | × |
13+| <term>Atlas inference products</term> | × |
14+| <term>Atlas training products</term> | × |
15+ 
16+## Function Description
17+ 
18+- Operator function: Completes the first stage of the two-stage computation of Attention Residuals historical residual attention. It computes the inter-block attention of all $S$ target slots in parallel at one time and returns the softmax weighted numerator and the statistics $(O,M,L)$ for the next stage to fuse the current residual.
19+ 
20+- Computation formula:
21+ 
22+ $$
23+ \operatorname{block\_attn\_res\_prepare}
24+ \left(V,Q,\operatorname{valid\_blocks};\epsilon\right)
25+ \longrightarrow (O,M,L)
26+ $$
27+ 
28+ Computes the RMS normalization factor for the $n$-th valid historical source of token $t$:
29+ 
30+ $$
31+ R_{t,n}=\sqrt{\frac{1}{D}\sum_{d=0}^{D-1}V_{t,n,d}^{2}+\epsilon}
32+ $$
33+ 
34+ Computes the logit for each target slot, token, and valid historical source:
35+ 
36+ $$
37+ Z_{s,t,n}=R_{t,n}^{-1}\sum_{d=0}^{D-1}Q_{s,d}V_{t,n,d}
38+ $$
39+ 
40+ Let $N_v$ be the number of valid historical sources specified by `valid_blocks`. Then:
41+ 
42+ $$
43+ M_{s,t}=\max_{0\le n<N_v}Z_{s,t,n}
44+ $$
45+ 
46+ $$
47+ E_{s,t,n}=\exp\left(Z_{s,t,n}-M_{s,t}\right)
48+ $$
49+ 
50+ $$
51+ L_{s,t}=\sum_{n=0}^{N_v-1}E_{s,t,n}
52+ $$
53+ 
54+ $$
55+ O_{s,t,d}=\sum_{n=0}^{N_v-1}E_{s,t,n}V_{t,n,d}
56+ $$
57+ 
58+ When `valid_blocks[0] == 0`, the preceding softmax computation is not performed, and the following values are returned directly:
59+ 
60+ $$
61+ O=\mathbf{0}_{S\times T\times D},\qquad
62+ M=m_{\min}\mathbf{1}_{S\times T},\qquad
63+ L=\mathbf{0}_{S\times T},\qquad
64+ m_{\min}=-3.4028234663852886\times10^{38}
65+ $$
66+ 
67+ That is, `numerator` is an all-0 tensor with the shape `[S, T, D]`, `logit_max` is a tensor with the shape `[S, T]` whose elements are all the minimum finite FLOAT32 value $m_{\min}$, and `exp_sum` is an all-0 tensor with the shape `[S, T]`. This result indicates the empty state of online softmax.
68+ 
69+ where:
70+ 
71+ - $V\in\mathbb{R}^{T\times N\times D}$ denotes the input `block_res`; $V_{t,n,d}$ denotes the value of the $d$-th hidden feature of the $n$-th historical residual block of the $t$-th token.
72+ - $Q\in\mathbb{R}^{S\times D}$ denotes the input `pseudo_query`; each target slot corresponds to a query vector, and $Q_{s,d}$ denotes the value of the $d$-th hidden feature of the $s$-th target slot.
73+ - $T$ denotes the number of tokens, $N$ denotes the number of historical residual blocks that dimension 1 of `block_res` can hold, $S$ denotes the number of target slots, and $D$ denotes HiddenSize.
74+ - $N_v=\min(\texttt{valid\_blocks[0]},N)$ denotes the number of historical residual blocks that actually participate in the computation.
75+ - $R$ denotes the RMS normalization factor, $Z$ denotes the normalized attention logit, and $E$ denotes the exponential value after max-shift.
76+ - $O$, $M$, and $L$ denote the softmax weighted numerator, the maximum logit, and the exponential sum, corresponding to the outputs `numerator`, `logit_max`, and `exp_sum`, respectively.
77+ - $\epsilon$ denotes the numerical stability term of RMS normalization, corresponding to the attribute `eps`.
78+ 
79+## Parameter Description
80+ 
81+<table style="undefined;table-layout: fixed; width: 980px"><colgroup>
82+ <col style="width: 120px">
83+ <col style="width: 150px">
84+ <col style="width: 490px">
85+ <col style="width: 120px">
86+ <col style="width: 100px">
87+ </colgroup>
88+ <thead>
89+ <tr>
90+ <th>Parameter name</th>
91+ <th>Input/Output/Attribute</th>
92+ <th>Description</th>
93+ <th>Data type</th>
94+ <th>Data format</th>
95+ </tr>
96+ </thead>
97+ <tbody>
98+ <tr>
99+ <td>block_res</td>
100+ <td>Input</td>
101+ <td>Historical residual blocks, corresponding to V in the formula, with the shape [T, N, D].</td>
102+ <td>FLOAT</td>
103+ <td>ND</td>
104+ </tr>
105+ <tr>
106+ <td>valid_blocks</td>
107+ <td>Input</td>
108+ <td>Number of currently valid historical sources, with the shape [1]. When the value is 0, numerator returns an all-0 tensor, logit_max returns a tensor whose elements are all the minimum finite FLOAT32 value, and exp_sum returns an all-0 tensor. When the value is greater than N, it is processed as N.</td>
109+ <td>UINT64</td>
110+ <td>ND</td>
111+ </tr>
112+ <tr>
113+ <td>pseudo_query</td>
114+ <td>Input</td>
115+ <td>Pseudo query, corresponding to Q in the formula. Each target slot corresponds to a query vector, with the shape [S, D].</td>
116+ <td>FLOAT</td>
117+ <td>ND</td>
118+ </tr>
119+ <tr>
120+ <td>numerator</td>
121+ <td>Output</td>
122+ <td>Softmax weighted numerator, corresponding to O in the formula, with the shape [S, T, D].</td>
123+ <td>FLOAT</td>
124+ <td>ND</td>
125+ </tr>
126+ <tr>
127+ <td>logit_max</td>
128+ <td>Output</td>
129+ <td>Maximum logit of valid historical sources, corresponding to M in the formula, with the shape [S, T].</td>
130+ <td>FLOAT</td>
131+ <td>ND</td>
132+ </tr>
133+ <tr>
134+ <td>exp_sum</td>
135+ <td>Output</td>
136+ <td>Exponential sum after max-shift, corresponding to L in the formula, with the shape [S, T].</td>
137+ <td>FLOAT</td>
138+ <td>ND</td>
139+ </tr>
140+ <tr>
141+ <td>eps</td>
142+ <td>Attribute</td>
143+ <td>Numerical stability term of RMS normalization. Optional. The default value is 1e-6, and it must be a finite positive number.</td>
144+ <td>FLOAT</td>
145+ <td>-</td>
146+ </tr>
147+ </tbody>
148+</table>
149+ 
150+## Constraint Description
151+ 
152+- `T >= 0`, `S >= 0`, `1 <= N <= 64`, and `1 <= D <= 8192`.
153+- The shape of `block_res` is `[T, N, D]`, and the shape of `pseudo_query` is `[S, D]`. The shapes of `numerator`, `logit_max`, and `exp_sum` are `[S, T, D]`, `[S, T]`, and `[S, T]`, respectively.
154+- The shape of `valid_blocks` must be `[1]`. When `valid_blocks[0] == 0`, `numerator` returns an all-0 tensor, `logit_max` returns a tensor whose elements are all the minimum finite FLOAT32 value, and `exp_sum` returns an all-0 tensor. When `valid_blocks[0] > N`, it is processed as `N`.
155+- All input and output tensors must be contiguous tensors. Non-contiguous tensors are not supported.
156+- `eps` must be a finite positive number.
157+- When `T == 0` or `S == 0`, empty outputs of the corresponding shapes are returned.
158+ 
159+## Invocation Description
160+ 
161+<table><thead>
162+ <tr>
163+ <th>Invocation method</th>
164+ <th>Invocation sample</th>
165+ <th>Description</th>
166+ </tr></thead>
167+<tbody>
168+ <tr>
169+ <td>aclnn invocation</td>
170+ <td><a href="./examples/arch35/test_aclnn_block_attn_res_prepare.cpp">test_aclnn_block_attn_res_prepare</a></td>
171+ <td>Invokes the operator through <a href="./docs/aclnnBlockAttnResPrepare_en.md">aclnnBlockAttnResPrepare_en.md</a>, and compiles the operator block_attn_res_prepare. For the detailed operator compilation and running methods, see <a href="../../docs/zh/invocation/quick_op_invocation.md">Operator Invocation</a>.</td>
172+ </tr>
173+ <tr>
174+ <td>PyTorch API</td>
175+ <td>-</td>
176+ <td>Invokes the block_attn_res_prepare operator through the <a href="../../attention/block_attn_res_prepare/docs/torchapi_block_attn_res_prepare_en.md">block_attn_res_prepare</a> interface.</td>
177+ </tr>
178+</tbody>
179+</table>
180+ 
181+```bash
182+# Execute in the root directory of the ops-transformer repository
183+bash build.sh --pkg --soc=ascend950 --ops=block_attn_res_prepare
184+(cd build_out && ./cann-ops-transformer-*.run)
185+bash build.sh --run_example block_attn_res_prepare eager cust --soc=ascend950 --vendor_name=custom
186+```
Aattention/block_attn_res_update/README_en.md+207-0
@@ -0,0 +1,207 @@
1+# BlockAttnResUpdate
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support Status
6+ 
7+| Product | Supported |
8+| ---- | :----:|
9+|<term>Ascend 950PR&950DT products</term>|√|
10+|<term>Atlas A3 products</term>|×|
11+|<term>Atlas A2 products</term>|×|
12+|<term>Atlas 200I/500 A2 inference products</term>|×|
13+|<term>Atlas inference products</term>|×|
14+|<term>Atlas training products</term>|×|
15+ 
16+## Description
17+ 
18+- Operator function: Performs the phase 2 update of BlockAttnRes. The operator first updates the local attention result partial_block with the incremental delta, then computes the normalized correlation score between the updated result and pseudo_query, and generates the attention output h by combining the existing online Softmax state numerator, logit_max, and exp_sum.
19+ 
20+- Computation formula:
21+ 
22+For each token $t$, first update the local attention result:
23+ 
24+$$
25+p_t = partial\_block_t + \operatorname{float}(delta_t)
26+$$
27+ 
28+Compute the RMS and the normalized correlation score of the updated result:
29+ 
30+$$
31+rms_t = \sqrt{\frac{1}{D}\sum_{j=0}^{D-1}p_{t,j}^{2} + eps}
32+$$
33+ 
34+$$
35+score_t = \frac{\sum_{j=0}^{D-1}p_{t,j} \cdot pseudo\_query_j}{rms_t}
36+$$
37+ 
38+Fuse the existing state and the current result in the online Softmax manner:
39+ 
40+$$
41+max_t = \max(logit\_max_t, score_t)
42+$$
43+ 
44+$$
45+alpha_t = \exp(logit\_max_t-max_t), \qquad
46+beta_t = \exp(score_t-max_t)
47+$$
48+ 
49+$$
50+denominator_t = exp\_sum_t \cdot alpha_t + beta_t
51+$$
52+ 
53+The final outputs are:
54+ 
55+$$
56+partial\_block\_out_t = p_t
57+$$
58+ 
59+$$
60+h_t = \operatorname{bfloat16}\left(
61+numerator_t \cdot \frac{alpha_t}{denominator_t} +
62+p_t \cdot \frac{beta_t}{denominator_t}\right)
63+$$
64+ 
65+Where $T$ denotes the number of tokens, $D$ denotes the hidden dimension, $t \in [0,T)$, and $j \in [0,D)$.
66+ 
67+For the non-empty history set $\mathcal{H}_t$ of token $t$, let the historical logits be $s_{t,i}$ and the corresponding $D$-dimensional values be
68+$v_{t,i}$. The input online Softmax state must satisfy:
69+ 
70+$$
71+logit\_max_t = \max_{i \in \mathcal{H}_t}s_{t,i}
72+$$
73+ 
74+$$
75+exp\_sum_t = \sum_{i \in \mathcal{H}_t}\exp(s_{t,i}-logit\_max_t) > 0
76+$$
77+ 
78+$$
79+numerator_t = \sum_{i \in \mathcal{H}_t}\exp(s_{t,i}-logit\_max_t)v_{t,i}
80+$$
81+ 
82+This operator only reads the historical state and fuses the current $p_t$ with it. It does not update numerator, logit_max, or exp_sum.
83+ 
84+## Parameter Description
85+ 
86+<table style="undefined;table-layout: fixed; width: 980px"><colgroup>
87+ <col style="width: 100px">
88+ <col style="width: 150px">
89+ <col style="width: 280px">
90+ <col style="width: 330px">
91+ <col style="width: 120px">
92+ </colgroup>
93+ <thead>
94+ <tr>
95+ <th>Parameter</th>
96+ <th>Input/Output/Attribute</th>
97+ <th>Description</th>
98+ <th>Data Type</th>
99+ <th>Data Format</th>
100+ </tr></thead>
101+ <tbody>
102+ <tr>
103+ <td>partial_block</td>
104+ <td>Input</td>
105+ <td>Local attention result to be updated, corresponding to partial_block in the formula, with shape $(T, D)$; reuses storage with the output of the same name.</td>
106+ <td>FLOAT</td>
107+ <td>ND</td>
108+ </tr>
109+ <tr>
110+ <td>delta</td>
111+ <td>Input</td>
112+ <td>Increment of the local attention result, corresponding to delta in the formula, with shape $(T, D)$.</td>
113+ <td>BFLOAT16</td>
114+ <td>ND</td>
115+ </tr>
116+ <tr>
117+ <td>pseudo_query</td>
118+ <td>Input</td>
119+ <td>Pseudo Query vector used to compute the normalized correlation score, corresponding to pseudo_query in the formula, with shape $(D)$.</td>
120+ <td>FLOAT</td>
121+ <td>ND</td>
122+ </tr>
123+ <tr>
124+ <td>numerator</td>
125+ <td>Input</td>
126+ <td>Numerator corresponding to the non-empty historical online Softmax state, corresponding to numerator in the formula, with shape $(T, D)$; read-only for this operator.</td>
127+ <td>FLOAT</td>
128+ <td>ND</td>
129+ </tr>
130+ <tr>
131+ <td>logit_max</td>
132+ <td>Input</td>
133+ <td>Finite maximum value of each token in the non-empty historical online Softmax state, corresponding to logit_max in the formula, with shape $(T)$; read-only for this operator.</td>
134+ <td>FLOAT</td>
135+ <td>ND</td>
136+ </tr>
137+ <tr>
138+ <td>exp_sum</td>
139+ <td>Input</td>
140+ <td>Finite positive exponent sum of each token in the non-empty historical online Softmax state, corresponding to exp_sum in the formula, with shape $(T)$; read-only for this operator.</td>
141+ <td>FLOAT</td>
142+ <td>ND</td>
143+ </tr>
144+ <tr>
145+ <td>partial_block</td>
146+ <td>Output</td>
147+ <td>Updated local attention result, corresponding to partial_block_out in the formula, with shape $(T, D)$; forms a reference relationship with the input of the same name.</td>
148+ <td>FLOAT</td>
149+ <td>ND</td>
150+ </tr>
151+ <tr>
152+ <td>h</td>
153+ <td>Output</td>
154+ <td>Attention result obtained by fusing the updated local result with the existing online Softmax state, corresponding to h in the formula, with shape $(T, D)$.</td>
155+ <td>BFLOAT16</td>
156+ <td>ND</td>
157+ </tr>
158+ <tr>
159+ <td>eps</td>
160+ <td>Optional attribute</td>
161+ <td>Stability term in the RMS computation. The value must be finite and greater than 0. The default value is $10^{-6}$.</td>
162+ <td>FLOAT</td>
163+ <td>-</td>
164+ </tr>
165+ </tbody></table>
166+ 
167+## Constraints
168+ 
169+- $T$ and $D$ must be known positive integers, and $D$ must be in the range $[1, 8192]$.
170+- The shapes of partial_block, delta, numerator, the output partial_block, and h must be the same, all $(T, D)$; the shapes of pseudo_query, logit_max, and exp_sum must be $(D)$, $(T)$, and $(T)$, respectively.
171+- The original format and storage format of all inputs must be ND; the StorageShape of all inputs and outputs must be exactly the same as the OriginShape.
172+- numerator, logit_max, and exp_sum are runtime read-only states in GM (Global Memory); Host Tiling does not read or verify their element values.
173+ You must ensure that the historical state of each token is non-empty: the elements in logit_max and numerator are finite values, exp_sum is a finite value strictly greater than 0,
174+ and the three tensors come from the same historical state. The empty historical state represented by `logit_max=-inf`, `exp_sum=0`, and a zero numerator is currently not supported.
175+- $T \times D$ cannot exceed the representation range of the signed 64-bit Kernel GM element offset; the value of $T$ after core partitioning cannot exceed the representation range of uint32_t.
176+- A complete row of $D$-dimensional data must fit into the Unified Buffer (UB) in the double buffering mode; otherwise, Tiling fails.
177+- partial_block is an in-place update parameter. During framework invocation, the reference relationship ensures that the input and output reuse storage; the Kernel directly reads and writes through the partial_block input address and does not access the corresponding output ABI address.
178+- The workspace size requested by Host Tiling is currently 0, and the Kernel does not access the workspace address.
179+- The RMS computation uses FP32 sqrt and division. It does not perform Newton iteration and does not perform extra processing on zero values, extreme values, NaN, or Inf in the input Tensor.
180+ 
181+## Calling Description
182+ 
183+<table><thead>
184+ <tr>
185+ <th>Invocation Method</th>
186+ <th>Invocation Sample</th>
187+ <th>Description</th>
188+ </tr></thead>
189+<tbody>
190+ <tr>
191+ <td>aclnn invocation</td>
192+ <td><a href="examples/arch35/test_aclnn_block_attn_res_update.cpp">test_aclnn_block_attn_res_update.cpp</a></td>
193+ <td>Invoked through the aclnnBlockAttnResUpdate API. For details, see the <a href="docs/aclnnBlockAttnResUpdate_en.md">API documentation</a>.</td>
194+ </tr>
195+ <tr>
196+ <td>PyTorch API</td>
197+ <td>-</td>
198+ <td>Invoked through the <a href="../../attention/block_attn_res_update/docs/torchapi_block_attn_res_update_en.md">cann_ops_transformer.block_attn_res_update</a> API.</td>
199+ </tr>
200+</tbody>
201+</table>
202+ 
203+```bash
204+## Quick Start
205+# Run in the repository root directory, assuming the environment variables are ready
206+bash build.sh --soc=${soc_version} --ops=block_attn_res_update
207+```
Aattention/block_sparse_attention/README_en.md+191-0
@@ -0,0 +1,191 @@
1+# BlockSparseAttention Operator
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Overview
6+ 
7+BlockSparseAttention is a high-performance sparse attention operator implemented based on the CATLASS template library. It supports flexible block-level sparse patterns.
8+ 
9+## Features
10+ 
11+- **Flexible sparse block pattern**: Supports custom x*y block-level sparse patterns, specified by the blockShape parameter
12+- **Multiple KV Cache layouts**: Supports the TND [T, N, D] and BNSD [B, N, S, D] KV cache layouts
13+- **Multiple Q input layouts**: Supports two Query input layouts: TND [T, N, D] and BNSD [B, N, S, D]
14+- **High-performance computing**: Based on the CATLASS template library, making full use of the Ascend A2 hardware features
15+- **Multi-precision support**: Supports FP16 and BF16 precisions
16+- **Unaligned sequences**: Supports scenarios where the sequence length is not divisible by the block size, and automatically handles boundary cases
17+ 
18+## Interface Definition
19+ 
20+### aclnnBlockSparseAttentionGetWorkspaceSize
21+ 
22+```cpp
23+aclnnStatus aclnnBlockSparseAttentionGetWorkspaceSize(
24+ const aclTensor *query, // Query input
25+ const aclTensor *key, // Key input
26+ const aclTensor *value, // Value input
27+ const aclTensor *blockSparseMaskOptional, // Sparse mask
28+ const aclTensor *attenMaskOptional, // Attention mask (not supported currently)
29+ const aclIntArray *blockShapeOptional, // Sparse block shape array
30+ const aclIntArray *actualSeqLengthsOptional, // Actual Q sequence lengths (required for the TND layout)
31+ const aclIntArray *actualSeqLengthsKvOptional, // Actual KV sequence lengths (required for the TND layout)
32+ const aclTensor *blockTableOptional, // Block table (used for PagedAttention, not supported currently)
33+ char *qInputLayout, // Query data layout (TND/BNSD)
34+ char *kvInputLayout, // Key/Value data layout (TND/BNSD)
35+ int64_t numKeyValueHeads, // Number of KV heads
36+ int64_t maskType, // Mask type
37+ double scaleValue, // Scale factor
38+ int64_t innerPrecise, // Precision level of the Softmax computation
39+ int64_t blockSize, // Block size (used for PagedAttention, not supported currently)
40+ int64_t preTokens, // Sliding window parameter (not supported currently)
41+ int64_t nextTokens, // Sliding window parameter (not supported currently)
42+ int64_t softmaxLseFlag, // Whether to output LSE
43+ aclTensor *attentionOut, // Output tensor
44+ aclTensor *softmaxLseOptional, // Softmax LSE output (optional)
45+ uint64_t *workspaceSize, // Returns the workspace size
46+ aclOpExecutor **executor); // Returns the executor
47+```
48+ 
49+### aclnnBlockSparseAttention
50+ 
51+```cpp
52+aclnnStatus aclnnBlockSparseAttention(
53+ void *workspace, // Workspace address
54+ uint64_t workspaceSize, // Workspace size
55+ aclOpExecutor *executor, // executor
56+ aclrtStream stream); // ACL stream
57+```
58+ 
59+## Parameter Description
60+ 
61+### Input Parameters
62+ 
63+- **query**: Query tensor, supports the following layouts
64+ - TND: [total_tokens, num_heads, head_dim]
65+ - BNSD: [batch, num_heads, seqlen, head_dim]
66+ 
67+- **key**: Key tensor, supports the following layouts
68+ - TND: [total_kv_tokens, num_kv_heads, head_dim]
69+ - BNSD: [batch, num_kv_heads, kv_seqlen, head_dim]
70+ 
71+- **value**: Value tensor, supports the following layouts
72+ - TND: [total_kv_tokens, num_kv_heads, head_dim]
73+ - BNSD: [batch, num_kv_heads, kv_seqlen, head_dim]
74+ 
75+- **blockSparseMaskOptional**:
76+ - Sparse mask: [batch, headNum, qBlockNum, kvBlockNum]. For TND, the block numbers are computed from the maximum values of actualSeqLengths/actualSeqLengthsKv; for BNSD/BSND, the block numbers can be computed from S in the Q/K shapes, or from their maximum values when actualSeqLengths/actualSeqLengthsKv are passed in.
77+ 
78+- **blockShapeOptional**:
79+ - blockShapeX: Block size in the Q direction
80+ - blockShapeY: Block size in the KV direction
81+ 
82+- **innerPrecise**: Controls the data types used in the online softmax and rescale computation stages. The value range is limited by the product and the input data type. For details, see the Precautions section.
83+ - 0: Both online softmax and rescale use FLOAT32.
84+ - 1: Both online softmax and rescale use FLOAT16; only FLOAT16 input is supported.
85+ - 4: Online softmax uses FLOAT16/BFLOAT16, and rescale uses FLOAT32.
86+ 
87+### Output Parameters
88+ 
89+- **attentionOut**: Attention output, with the same shape as query
90+- **softmaxLse**: Softmax log-sum-exp output (optional)
91+ 
92+## Usage Example
93+ 
94+```cpp
95+// 1. Prepare the input data
96+int64_t batch = 1;
97+int64_t qSeqlen = 256;
98+int64_t kvSeqlen = 512;
99+int64_t numHeads = 8;
100+int64_t numKvHeads = 8;
101+int64_t headDim = 128;
102+int64_t blockShapeX = 128;
103+int64_t blockShapeY = 128;
104+ 
105+// 2. Create tensors (omitted)
106+ 
107+// 3. Set blockShape
108+int64_t blockShapeData[2] = {blockShapeX, blockShapeY};
109+aclIntArray *blockShape = aclCreateIntArray(blockShapeData, 2);
110+ 
111+// 4. Call the operator
112+uint64_t workspaceSize = 0;
113+aclOpExecutor *executor = nullptr;
114+ 
115+aclnnStatus ret = aclnnBlockSparseAttentionGetWorkspaceSize(
116+ query, key, value, blockSparseMask, nullptr, blockShape, nullptr,
117+ nullptr, nullptr, "BNSD", "BNSD", numKvHeads, 0, scaleValue, 0, 128,
118+ 2147483647, 2147483647, 0, attentionOut, nullptr, &workspaceSize, &executor);
119+ 
120+if (ret == ACLNN_SUCCESS) {
121+ void *workspace = nullptr;
122+ aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
123+ 
124+ ret = aclnnBlockSparseAttention(workspace, workspaceSize, executor, stream);
125+ 
126+ aclrtFree(workspace);
127+}
128+```
129+ 
130+## Sparse Pattern Description
131+ 
132+### BlockSparseMask Sparse Pattern
133+ 
134+BlockSparseAttention uses the blockSparseMask sparse pattern and does not need the selectIdx index format or the traditional sBlockIdx computation:
135+ 
136+- **blockSparseMask**: Sparse mask with the shape `[batch, numHeads, qBlockNum, kvBlockNum]`
137+- **blockShapeX**: Basic block size in the Q direction
138+- **blockShapeY**: Basic block size in the KV direction
139+ 
140+### Example
141+ 
142+For qSeqlen=512, kvSeqlen=1024, blockShapeX=128, and blockShapeY=128:
143+ 
144+```
145+Q direction: ceil(512/128)=4 blocks [0, 1, 2, 3]
146+KV direction: ceil(1024/128)=8 blocks [0, 1, 2, 3, 4, 5, 6, 7]
147+ 
148+Each Q block can select any combination of KV blocks:
149+- KV blocks selected by Q block 0: [0, 2, 5]
150+- KV blocks selected by Q block 1: [1, 3, 6]
151+- KV blocks selected by Q block 2: [0, 4, 7]
152+- KV blocks selected by Q block 3: [2, 5, 6]
153+```
154+ 
155+## Performance Characteristics
156+ 
157+- **Memory efficiency**: Significantly reduces memory usage compared with dense attention
158+- **Computing efficiency**: Computes only the specified sparse blocks, reducing the computation workload
159+- **Flexible configuration**: Supports sparse block configurations with any blockShape
160+- **Hardware optimization**: Makes full use of the AI Core and Vector Core of Ascend A2
161+ 
162+## Calling Description
163+ 
164+| Invocation Method | Sample Code | Description |
165+| ---------- | -------- | --------------------------------------------- |
166+| PyTorch API | - | Calls the block_sparse_attention operator through the [block_sparse_attention](docs/torchapi_block_sparse_attention_en.md) API. |
167+ 
168+## Precautions
169+ 
170+1. qInputLayout currently supports only "TND" and "BNSD".
171+2. kvInputLayout currently supports only "TND" and "BNSD".
172+3. If blockShapeOptional is passed in, it must contain at least two elements [blockShapeX, blockShapeY], and the values must be greater than 0. blockShapeY must be a multiple of 16 on <term>Ascend 950PR&950DT products</term> and a multiple of 128 on <term>Atlas A2 products</term> and <term>Atlas A3 products</term>.
173+4. qSeqlen and kvSeqlen do not need to be divisible by blockShape. Unaligned scenarios are supported, and the actual number of blocks is computed by rounding up.
174+5. blockSparseMaskOptional must be passed in currently, with the shape [batch, headNum, qBlockNum, kvBlockNum]. For TND, qBlockNum = ceilDiv(max(actualSeqLengths), blockShapeX), and kvBlockNum = ceilDiv(max(actualSeqLengthsKv), blockShapeY). For BNSD/BSND, qBlockNum = ceilDiv(S in the Q shape, blockShapeX), and kvBlockNum = ceilDiv(S in the K shape, blockShapeY); if actualSeqLengths and actualSeqLengthsKv are passed in, the block numbers can also be computed from the actual lengths: qBlockNum = ceilDiv(max(actualSeqLengths), blockShapeX), and kvBlockNum = ceilDiv(max(actualSeqLengthsKv), blockShapeY).
175+6. attenMaskOptional currently supports only nullptr.
176+7. maskType currently supports only 0, which means no mask is added.
177+8. actualSeqLengthsOptional is required when qInputLayout is "TND"; actualSeqLengthsKvOptional is required when kvInputLayout is "TND".
178+9. PagedAttention is not supported currently, and blockTableOptional currently supports only nullptr.
179+10. On <term>Ascend 950PR&950DT products</term>, FLOAT16/BFLOAT16 non-quantized input supports innerPrecise set to 0 or 4, but not 1; on <term>Atlas A2 products</term> and <term>Atlas A3 products</term>, FLOAT16 input supports 0 or 1, and BFLOAT16 input supports only 0.
180+ 
181+## Compilation Description
182+ 
183+Operator compilation is integrated into the unified compilation process of ops-transformer:
184+ 
185+```bash
186+bash build.sh
187+```
188+ 
189+## Version History
190+ 
191+- v1.0.0: Initial version
Aattention/block_sparse_attention_grad/README_en.md+382-0
@@ -0,0 +1,382 @@
1+# BlockSparseAttentionGrad
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+|Product | Supported |
8+|:----------------------------|:-----------:|
9+|<term>Ascend 950PR&950DT products</term>| √ |
10+|<term>Atlas A3 training products</term>| √ |
11+|<term>Atlas A3 inference products</term>| × |
12+|<term>Atlas A2 training products</term>| √ |
13+|<term>Atlas A2 inference products</term>| × |
14+|<term>Atlas 200I/500 A2 inference products</term> | × |
15+|<term>Atlas inference products</term> | × |
16+|<term>Atlas training products</term> | × |
17+ 
18+## Function Description
19+ 
20+* Operator function: aclnnBlockSparseAttention sparse attention backward computation. This operator supports flexible block-level sparsity patterns and uses BlockSparseMask to specify the KV blocks selected by each Q block, implementing efficient sparse attention computation.
21+* Computation formula:
22+ 
23+ The sparse block size is $blockShapeX×blockShapeY$, and BlockSparseMask specifies the sparsity pattern.
24+ 
25+ blockShapeX and blockShapeY must be positive multiples of 16. BlockSparseMask defines the sparsity pattern at the logical block granularity. Each mask element controls whether a contiguous block of blockShapeX × blockShapeY participates in the computation.
26+ The forward computation formula is:
27+ 
28+ $$
29+ attentionOut=Softmax(Mask(scale⋅query⋅key_{sparse}^{T}, atten\_mask))⋅value_{sparse}
30+ $$
31+ 
32+ For ease of expression, the computation formulas are expressed using the variables $S$ and $P$:
33+ 
34+ $$
35+ S = Mask(scale⋅query⋅key_{sparse}^{T},atten\_mask)
36+ $$
37+ 
38+ $$
39+ P = SoftMax(S)
40+ $$
41+ 
42+ $$
43+ V = value_{sparse}
44+ $$
45+ 
46+ $$
47+ Out = PV
48+ $$
49+ 
50+ The backward computation formulas are:
51+ 
52+ $$
53+ softmax\_grad = softmaxGrad(dOut, attentionOut)
54+ $$
55+ 
56+ $$
57+ dP=dOut * V^T
58+ $$
59+ 
60+ $$
61+ dS = P * (dP-softmax\_grad)
62+ $$
63+ 
64+ $$
65+ dV=P^T * dOut
66+ $$
67+ 
68+ $$
69+ dQ=(dS*K)*scale
70+ $$
71+ 
72+ $$
73+ dK=(dS^T*Q)*scale
74+ $$
75+ 
76+BlockSparseAttentionGrad supports multiple interpretations of the data layout of the input dout, query, key, value, and attentionOut, which can be passed in through qInputLayout and kvInputLayout. To facilitate understanding of the specific layout formats supported later (such as BNSD and TND), this section first provides a unified description of the dimension meaning of each abbreviation letter in the layout formats:
77+ 
78+* B: Batch size of the input samples (Batch)
79+* T: Length when B and S are tightly arranged on the same axis (Total tokens)
80+* S: Sequence length of the input samples (Seq-Length)
81+* H: Hidden layer size (Head-Size)
82+* N: Number of heads (Head-Num)
83+* D: Size of the smallest unit of the hidden layer, which must satisfy D=H/N (Head-Dim)
84+ 
85+Currently supported layouts:
86+ 
87+* qInputLayout: "TND" "BNSD" "BSND"
88+* kvInputLayout: "TND" "BNSD" "BSND"
89+ 
90+- <term>Atlas A2 training products</term> and <term>Atlas A3 training products</term>:
91+"BSND" is not supported.
92+ 
93+## Parameter Description
94+ 
95+<table style="undefined;table-layout: fixed; width: 1550px">
96+<colgroup>
97+ <col style="width: 170px">
98+ <col style="width: 120px">
99+ <col style="width: 271px">
100+ <col style="width: 330px">
101+ <col style="width: 223px">
102+ <col style="width: 101px">
103+ <col style="width: 190px">
104+ <col style="width: 145px">
105+</colgroup>
106+<thead>
107+ <tr>
108+ <th>Parameter</th>
109+ <th>Input/Output</th>
110+ <th>Description</th>
111+ <th>Usage Instructions</th>
112+ <th>Data Type</th>
113+ <th>Data Format</th>
114+ <th>Dimensions (shape)</th>
115+ <th>Non-contiguous Tensor</th>
116+ </tr>
117+</thead>
118+<tbody>
119+ <tr>
120+ <td>dout (aclTensor*)</td>
121+ <td>Input</td>
122+ <td>Backward output gradient, representing the gradient information of the final output with respect to this operator.</td>
123+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalQTokens, headNum, headDim].</li><li>BNSD: [batch, headNum, maxQSeqLength, headDim].</li><li>BSND: [batch, maxQSeqLength, headNum, headDim].</li></ul></td>
124+ <td>FLOAT16, BFLOAT16</td>
125+ <td>ND</td>
126+ <td>3-4</td>
127+ <td>×</td>
128+ </tr>
129+ <tr>
130+ <td>query (aclTensor*)</td>
131+ <td>Input</td>
132+ <td>Query vector in the attention computation, that is, query in the formula.</td>
133+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalQTokens, headNum, headDim].</li><li>BNSD: [batch, headNum, maxQSeqLength, headDim].</li><li>BSND: [batch, maxQSeqLength, headNum, headDim].</li></ul></td>
134+ <td>FLOAT16, BFLOAT16</td>
135+ <td>ND</td>
136+ <td>3-4</td>
137+ <td>×</td>
138+ </tr>
139+ <tr>
140+ <td>key (aclTensor*)</td>
141+ <td>Input</td>
142+ <td>Key vector in the attention computation, that is, key in the formula.</td>
143+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalKTokens, numKeyValueHeads, headDim].</li><li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim].</li><li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim].</li></ul></td>
144+ <td>FLOAT16, BFLOAT16</td>
145+ <td>ND</td>
146+ <td>3-4</td>
147+ <td>×</td>
148+ </tr>
149+ <tr>
150+ <td>value (aclTensor*)</td>
151+ <td>Input</td>
152+ <td>Value vector in the attention computation, that is, value in the formula.</td>
153+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalVTokens, numKeyValueHeads, headDim].</li><li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim].</li><li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim].</li></ul></td>
154+ <td>FLOAT16, BFLOAT16</td>
155+ <td>ND</td>
156+ <td>3-4</td>
157+ <td>×</td>
158+ </tr>
159+ <tr>
160+ <td>attentionOut (aclTensor*)</td>
161+ <td>Input</td>
162+ <td>Output result of the forward BlockSparseAttention computation, that is, attentionOut in the formula.</td>
163+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalQTokens, headNum, headDim].</li><li>BNSD: [batch, headNum, maxQSeqLength, headDim].</li><li>BSND: [batch, maxQSeqLength, headNum, headDim].</li></ul></td>
164+ <td>FLOAT16, BFLOAT16</td>
165+ <td>ND</td>
166+ <td>3-4</td>
167+ <td>×</td>
168+ </tr>
169+ <tr>
170+ <td>softmaxLse (aclTensor*)</td>
171+ <td>Input</td>
172+ <td>Log-sum-exp intermediate result of the Softmax computation. It is used for the logarithm and exponential back-derivation of gradients in the backward computation.</td>
173+ <td>Empty tensors are not supported.<br>Supported shapes: <ul><li>TND: [totalQTokens, headNum, 1].</li><li>BNSD: [batch, headNum, maxQSeqLength, 1].</li></ul></td>
174+ <td>FLOAT</td>
175+ <td>ND</td>
176+ <td>3-4</td>
177+ <td>×</td>
178+ </tr>
179+ <tr>
180+ <td>blockSparseMaskOptional (aclTensor*)</td>
181+ <td>Input</td>
182+ <td>Block sparse mask, indicating the actual sparsity pattern. Determines which blocks actually participate in the attention computation.</td>
183+ <td>Empty tensors are not supported.<br>Optional input (required in the current version): <ul><li>The shape is [batch, headNum, ceilDiv(maxQSeqLength, blockShapeX), ceilDiv(maxKvSeqLength, blockShapeY)].</li><li>Indicates which blocks participate in the computation (value 1) and which blocks do not (value 0) after the data is divided by blocks.</li><li>If nullptr is passed in, block sparse computation is considered disabled, that is, the attention scores between all tokens are computed.</li></ul></td>
184+ <td>BOOL</td>
185+ <td>ND</td>
186+ <td>4</td>
187+ <td>×</td>
188+ </tr>
189+ <tr>
190+ <td>attenMaskOptional (aclTensor*)</td>
191+ <td>Input</td>
192+ <td>Attention mask, that is, atten_mask in the formula. It is used to mask out specific tokens that should not participate in the computation.</td>
193+ <td>Empty tensors are supported.<br>Required when maskType=1. The shape must be exactly [batch, headNum, maxBlockNum, 2], and the dtype must be INT32: <ul><li>maxBlockNum = max(ceilDiv(maxQSeqLength, blockShapeX), ceilDiv(maxKvSeqLength, blockShapeY)).</li><li>[..., 0] is actualSizeX along the Q direction, and [..., 1] is actualSizeY along the KV direction. The constraint 0 ≤ actualSize ≤ corresponding blockShape must be satisfied.</li><li>Within a block, [0, actualSize) validly participates in the computation, and [actualSize, blockShape) is considered padding. actualSize=0 indicates that the block does not participate in the computation.</li></ul>When maskType=0, this input is ignored, and nullptr must be passed in.</td>
194+ <td>BOOL, INT32</td>
195+ <td>ND</td>
196+ <td>2</td>
197+ <td>×</td>
198+ </tr>
199+ <tr>
200+ <td rowspan="3">blockShapeOptional (aclIntArray*)</td>
201+ <td rowspan="3">Input</td>
202+ <td rowspan="3">Array of sparse block shapes. It specifies the two-dimensional size (number of rows and columns) of each sparse block.</td>
203+ <td> <ul><li>When blockSparseMaskOptional is configured: if this input is configured, the operator obtains the sparse block size from it; if this input is not configured, the operator uses the default sparse block size [128,128].</li></ul></td>
204+ <td rowspan="3">INT64</td>
205+ <td rowspan="3">-</td>
206+ <td rowspan="3">1</td>
207+ <td rowspan="3">-</td>
208+ </tr>
209+ <tr>
210+ <td><ul><li>When blockSparseMaskOptional is not configured: the operator ignores this input regardless of how it is configured.</li></ul></td>
211+ </tr>
212+ <tr>
213+ <td>Element requirements when this input is configured: <ul><li>It must contain at least two elements [blockShapeX, blockShapeY].</li><li>blockShapeX: block size along the Q direction. The value must be greater than 0.</li><li>blockShapeY: block size along the KV direction. The value must be greater than 0.</li></ul></td>
214+ </tr>
215+ <tr>
216+ <td>actualSeqLengthsOptional (aclIntArray*)</td>
217+ <td>Input</td>
218+ <td>Array of actual query sequence lengths.<br>It describes the number of actually valid query tokens in each batch in variable-length sequence scenarios (that is, scenarios containing Padding data).</td>
219+ <td> Variable-length sequence scenario (when qInputLayout is "TND"): this input must be configured. Because the TND format is a one-dimensional contiguous layout, the operator relies on this array to accurately split and delimit the real boundaries of each sequence.</td>
220+ <td>INT64</td>
221+ <td>-</td>
222+ <td>1</td>
223+ <td>-</td>
224+ </tr>
225+ <tr>
226+ <td>actualSeqLengthsKvOptional (aclIntArray*)</td>
227+ <td>Input</td>
228+ <td>Array of actual key/value sequence lengths.<br>It describes the number of actually valid key/value tokens in each batch in variable-length sequence scenarios (that is, scenarios containing Padding data).</td>
229+ <td> Variable-length sequence scenario (when kvInputLayout is "TND"): this input must be configured. Because the TND format is a one-dimensional contiguous layout, the operator relies on this array to accurately split and delimit the real boundaries of each sequence.</td>
230+ <td>INT64</td>
231+ <td>-</td>
232+ <td>1</td>
233+ <td>-</td>
234+ </tr>
235+ <tr>
236+ <td>qInputLayout (char*)</td>
237+ <td>Input</td>
238+ <td>Data layout format of query. It indicates the specific arrangement of the input tensors in memory (such as contiguous or concatenated arrangement).</td>
239+ <td>Currently only "TND", "BNSD", and "BSND" are supported, and qInputLayout must be consistent with kvInputLayout.</td>
240+ <td>-</td>
241+ <td>-</td>
242+ <td>-</td>
243+ <td>-</td>
244+ </tr>
245+ <tr>
246+ <td>kvInputLayout (char*)</td>
247+ <td>Input</td>
248+ <td>Data layout format of key and value. It indicates the specific arrangement of the input tensors in memory.</td>
249+ <td>Currently only "TND", "BNSD", and "BSND" are supported, and qInputLayout must be consistent with kvInputLayout.</td>
250+ <td>-</td>
251+ <td>-</td>
252+ <td>-</td>
253+ <td>-</td>
254+ </tr>
255+ <tr>
256+ <td>numKeyValueHeads (int64_t)</td>
257+ <td>Input</td>
258+ <td>Number of attention heads of key/value. It is used to support the head number ratio mapping under the GQA (grouped query attention) mechanism.</td>
259+ <td>-</td>
260+ <td>-</td>
261+ <td>-</td>
262+ <td>-</td>
263+ <td>-</td>
264+ </tr>
265+ <tr>
266+ <td>maskType (int64_t)</td>
267+ <td>Input</td>
268+ <td>Mask type in the attention computation. It specifies which preset mask logic is used.</td>
269+ <td>Currently 0 and 1 are supported: <ul><li>0: No mask is applied. attenMaskOptional is ignored.</li><li>1: PerBlockSizeMask scenario. The computation is performed based on the block-level variable-length mask, and attenMaskOptional must be configured.</li></ul></td>
270+ <td>-</td>
271+ <td>-</td>
272+ <td>-</td>
273+ <td>-</td>
274+ </tr>
275+ <tr>
276+ <td>scaleValue (double)</td>
277+ <td>Input</td>
278+ <td>Scaling factor, that is, scale in the formula. It is used to normalize the attention scores.</td>
279+ <td>Generally set to D^-0.5.</td>
280+ <td>-</td>
281+ <td>-</td>
282+ <td>-</td>
283+ <td>-</td>
284+ </tr>
285+ <tr>
286+ <td>preTokens (int64_t)</td>
287+ <td>Input</td>
288+ <td>Number of preceding tokens included in the sliding window. It limits the number of historical tokens before the current token that participate in the attention computation.</td>
289+ <td>Used for sliding window attention scenarios. Sliding window attention is currently not supported, so only 2147483647 can be passed in.</td>
290+ <td>-</td>
291+ <td>-</td>
292+ <td>-</td>
293+ <td>-</td>
294+ </tr>
295+ <tr>
296+ <td>nextTokens (int64_t)</td>
297+ <td>Input</td>
298+ <td>Number of following tokens included in the sliding window. It limits the number of future tokens after the current token that participate in the attention computation.</td>
299+ <td>Used for sliding window attention scenarios. Sliding window attention is currently not supported, so only 2147483647 can be passed in.</td>
300+ <td>-</td>
301+ <td>-</td>
302+ <td>-</td>
303+ <td>-</td>
304+ </tr>
305+ <tr>
306+ <td>dq (aclTensor*)</td>
307+ <td>Output</td>
308+ <td>Gradient output of query, that is, dq in the formula.</td>
309+ <td>Empty tensors are not supported.<br>The data type and shape are the same as those of the input query.</td>
310+ <td>FLOAT16, BFLOAT16</td>
311+ <td>ND</td>
312+ <td>3-4</td>
313+ <td>√</td>
314+ </tr>
315+ <tr>
316+ <td>dk (aclTensor*)</td>
317+ <td>Output</td>
318+ <td>Gradient output of key, that is, dk in the formula.</td>
319+ <td>Empty tensors are not supported.<br>The data type and shape are the same as those of the input key.</td>
320+ <td>FLOAT16, BFLOAT16</td>
321+ <td>ND</td>
322+ <td>3-4</td>
323+ <td>√</td>
324+ </tr>
325+ <tr>
326+ <td>dv (aclTensor*)</td>
327+ <td>Output</td>
328+ <td>Gradient output of value, that is, dv in the formula.</td>
329+ <td>Empty tensors are not supported.<br>The data type and shape are the same as those of the input value.</td>
330+ <td>FLOAT16, BFLOAT16</td>
331+ <td>ND</td>
332+ <td>3-4</td>
333+ <td>√</td>
334+ </tr>
335+ <tr>
336+ <td>workspaceSize (uint64_t*)</td>
337+ <td>Output</td>
338+ <td>Returns the workspace size that needs to be requested on the Device.</td>
339+ <td>-</td>
340+ <td>-</td>
341+ <td>-</td>
342+ <td>-</td>
343+ <td>-</td>
344+ </tr>
345+ <tr>
346+ <td>executor (aclOpExecutor**)</td>
347+ <td>Output</td>
348+ <td>Returns the op executor, which contains the operator computation process.</td>
349+ <td>-</td>
350+ <td>-</td>
351+ <td>-</td>
352+ <td>-</td>
353+ <td>-</td>
354+ </tr>
355+</tbody>
356+</table>
357+ 
358+<ul>
359+- <term>Atlas A2 training products</term> and <term>Atlas A3 training products</term>:
360+FLOAT8_E5M2 and FLOAT8_E4M3FN are not supported.
361+</ul>
362+ 
363+<ul>
364+- <term>Ascend 950PR&950DT products</term>:
365+maskType does not support the value 1; the attenMaskOptional input is ignored, and nullptr should be passed in for this input.
366+</ul>
367+ 
368+## Constraints
369+ 
370+* When this API is used together with PyTorch, ensure that the versions of the CANN packages match those of the PyTorch packages.
371+* actualSeqLengthsOptional is required when qInputLayout is "TND"; actualSeqLengthsKvOptional is required when kvInputLayout is "TND".
372+* The layout of softmaxLse must be consistent with the layout of query. If the layout of query is "BSND", the layout "BNS1" should be passed in for softmaxLse.
373+* HeadDim must be 128.
374+* According to the input layouts supported by the operator, the size of the corresponding head dimension in the query tensor shape is denoted as N1, and the size of the corresponding head dimension in the key and value tensor shapes is denoted as N2. N1 % N2 == 0 must be satisfied. MQA, MHA, and GQA scenarios are supported.
375+* Constraints related to actualSeqLengthsOptional and actualSeqLengthsKvOptional: when qInputLayout or kvInputLayout is not "TND", these two input parameters are ignored.
376+ 
377+## Calling Instructions
378+ 
379+| Calling Method | Calling Sample | Description |
380+|----------------|-------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------|
381+| aclnn calling | [test_aclnn_block_sparse_attention_grad](examples/test_aclnn_block_sparse_attention_grad.cpp) | Calls the BlockSparseAttentionGrad operator through the [aclnnBlockSparseAttentionGrad](docs/aclnnBlockSparseAttentionGrad_en.md) API. |
382+| PyTorch API | - | Calls the block_sparse_attention_grad operator through the [torchapi_block_sparse_attention_grad](docs/torchapi_block_sparse_attention_grad.md) API. |
Aattention/bsa_select_block_mask/README_en.md+164-0
@@ -0,0 +1,164 @@
1+# BSASelectBlockMask
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+| Product | Supported |
8+| :----------------------------------------- | :------:|
9+| <term>Ascend 950PR&950DT products</term> | √ |
10+| <term>Atlas A3 training products</term> | √ |
11+| <term>Atlas A3 inference products</term> | √ |
12+| <term>Atlas A2 training products</term> | √ |
13+| <term>Atlas A2 inference products</term> | √ |
14+| <term>Atlas 200I/500 A2 inference products</term> | × |
15+| <term>Atlas inference products</term> | × |
16+| <term>Atlas training products</term> | × |
17+ 
18+## Function Description
19+ 
20+- Operator function: aclnnBSASelectBlockMask is a preceding operator of BSA (BlockSparseAttention). It dynamically generates blockSparseMask based on the content of query and key, changing the BSA call chain from "manually providing the mask" to "adaptively selecting the sparsity pattern based on the Q/K content".
21+- Computation formula:
22+ 
23+ Assume blockShape = [blockShapeX, blockShapeY], Sq is the maximum sequence length of query, and Skv is the maximum sequence length of key. The number of blocks after compression is:
24+ 
25+ $$
26+ Xblocks = \lceil Sq / blockShapeX \rceil,\quad Yblocks = \lceil Skv / blockShapeY \rceil
27+ $$
28+ 
29+ **Step1: Mean Pooling Compression**
30+ 
31+ When actualBlockLenQuery / actualBlockLenKey is null (full compression):
32+ 
33+ $$
34+ q\_compressed[b, n, x, d] = \frac{1}{blockShapeX} \sum_{i=0}^{blockShapeX-1} query[b, n, x \cdot blockShapeX + i, d]
35+ $$
36+ 
37+ $$
38+ k\_compressed[b, n, y, d] = \frac{1}{blockShapeY} \sum_{j=0}^{blockShapeY-1} key[b, n, y \cdot blockShapeY + j, d]
39+ $$
40+ 
41+ When actualBlockLenQuery / actualBlockLenKey is not null (partial compression), the mean is computed only over the first actualBlockLen tokens in each block:
42+ 
43+ $$
44+ q\_compressed[b, n, x, d] = \frac{1}{actualBlockLenQ[b,x]} \sum_{i=0}^{actualBlockLenQ[b,x]-1} query[b, n, x \cdot blockShapeX + i, d]
45+ $$
46+ 
47+ $$
48+ k\_compressed[b, n, y, d] = \frac{1}{actualBlockLenK[b,y]} \sum_{j=0}^{actualBlockLenK[b,y]-1} key[b, n, y \cdot blockShapeY + j, d]
49+ $$
50+ 
51+ **Step2a: QK Matmul**
52+ 
53+ $$
54+ score[b, n, x, y] = scale \cdot \sum_{d=0}^{D-1} q\_compressed[b, n, x, d] \cdot k\_compressed[b, n, y, d]
55+ $$
56+ 
57+ **Step2b: Softmax**
58+ 
59+ $$
60+ attn\_score[b, n, x, y] = softmax(score[b, n, x, :]) = \frac{\exp(score[b, n, x, y] - m_{final})}{l_{final}}
61+ $$
62+ 
63+ **Step2c: Post-Softmax Mean Pooling**
64+ 
65+ When postBlockShape is not null, a second round of mean pooling is performed on attn_score to generate the coarse-grained pooled_score:
66+ 
67+ $$
68+ postXBlocks = \lceil Xblocks / postBlockShapeX \rceil,\quad postYBlocks = \lceil Yblocks / postBlockShapeY \rceil
69+ $$
70+ 
71+ $$
72+ pooled\_score[b, n, px, py] = \frac{1}{|R_{px,py}|} \sum_{x \in R_x(px)} \sum_{y \in R_y(py)} attn\_score[b, n, x, y]
73+ $$
74+ 
75+ When postBlockShape is null, this step is skipped, and pooled_score = attn_score.
76+ 
77+ **Step3: TopK Selection to Generate Indices**
78+ 
79+ $$
80+ topk\_value = \text{round}(sparsity \times postXBlocks \times postYBlocks)
81+ $$
82+ 
83+ $$
84+ \mathcal{indices}= \text{TopK}\left(pooled\_score[b, n, px, py],\; topK\_value\right)
85+ $$
86+ 
87+ When postBlockShape is null, postXBlocks=Xblocks, postYBlocks=Yblocks, and pooled_score=attn_score, which is equivalent to directly performing TopK on attn_score.
88+ 
89+ Here indices is the set of indices corresponding to the topk\_value largest values in pooled\_score[b, n, px, py].
90+ 
91+ **Step4: Generate BlockSparseMask**
92+ 
93+ When postBlockShape is not null, the output is directly the **coarse-grained** mask (the shape is [B, N, postXBlocks, postYBlocks]; the second-round pooling splitting semantics apply, and no expansion is needed):
94+ 
95+ $$
96+ blockSparseMaskOut[b, n, px, py] =
97+ \begin{cases}
98+ 1 & (px, py) \in \mathcal{indices} \\
99+ 0 & (px, py) \notin \mathcal{indices}
100+ \end{cases}
101+ $$
102+ 
103+ When postBlockShape is null, the output is the fine-grained mask (the shape is [B, N, Xblocks, Yblocks]), which is generated element by element directly:
104+ 
105+ $$
106+ blockSparseMaskOut[b, n, x, y] =
107+ \begin{cases}
108+ 1 & (b, n, x, y) \in \mathcal{indices} \\
109+ 0 & (b, n, x, y) \notin \mathcal{indices}
110+ \end{cases}
111+ $$
112+ 
113+- Data layout format:
114+ 
115+ BSASelectBlockMask supports multiple interpretations of the data layout of the input query and key, which can be passed in through qInputLayout and kvInputLayout. To facilitate understanding of the specific layout formats supported later (such as BNSD and TND), this section first provides a unified description of the dimension meaning of each abbreviation letter in the layout formats:
116+ 
117+ - B: Batch size of the input samples (Batch)
118+ - T: Length when B and S are tightly arranged on the same axis (Total tokens)
119+ - S: Sequence length of the input samples (Seq-Length)
120+ - H: Hidden layer size (Head-Size)
121+ - N: Number of heads (Head-Num)
122+ - D: Size of the smallest unit of the hidden layer, which must satisfy D = H / N (Head-Dim)
123+ 
124+- Currently supported layouts:
125+ 
126+ - qInputLayout: "TND" "BNSD"
127+ - kvInputLayout: "TND" "BNSD"
128+ 
129+## Parameter Description
130+ 
131+| Parameter | Input/Output/Attribute | Description | Data Type | Data Format |
132+|-----|-----------|------|----------|----------|
133+| query | Input | Query matrix in the attention computation, that is, `query` in the formula. | FLOAT16, BFLOAT16 | ND |
134+| key | Input | Key matrix in the attention computation, that is, `key` in the formula. | The data type is the same as that of query | ND |
135+| block_shape | Input | Array of sparse block shapes. It specifies the two-dimensional size (number of rows and columns) of each sparse block, that is, `blockShape` in the formula. It must contain at least two elements `[blockShapeX, blockShapeY]`: blockShapeX is the block size along the Q direction, and blockShapeY is the block size along the KV direction. If this parameter is not configured (nullptr is passed in), the operator uses blockShapeX = 128 and blockShapeY = 128 by default. | INT64 | - |
136+| post_block_shape | Input | Array of second-round pooling block shapes. It specifies the block grouping size of the second-round pooling after Softmax (in blocks), that is, `postBlockShape` in the formula. Optional input: when `[postBlockShapeX, postBlockShapeY]` is passed in, the second-round pooling is enabled, TopK selects on the coarse-grained pooled_score, and the output mask is the coarse-grained [B, N, postXBlocks, postYBlocks] (the second-round pooling splitting semantics apply; the selection result is output directly without expansion). When nullptr is passed in, the second-round pooling is disabled (backward compatibility), TopK selects directly on attn_score, and the output mask is the fine-grained [B, N, Xblocks, Yblocks]. postBlockShapeX/postBlockShapeY are multiples of 8, indicating the number of block groups. | INT64 | - |
137+| actual_seq_lengths | Input | Actual sequence length of query for each batch, that is, the actual valid value of Sq in the formula for each batch. It describes the number of actually valid query tokens in each batch in variable-length sequence scenarios (that is, scenarios containing Padding data). | INT64 | - |
138+| actual_seq_lengths_kv | Input | Actual sequence length of key, that is, the actual valid value of Skv in the formula for each batch. It describes the number of actually valid key tokens in each batch in variable-length sequence scenarios (that is, scenarios containing Padding data). | INT64 | - |
139+| actual_block_len_query | Input | Actual valid sequence length compressed within each query block, that is, `actualBlockLenQuery` in the formula. It is used for partial compression scenarios (such as incomplete blocks at the end or compressing only valid tokens). The mean is computed only over the first actual_block_len_query tokens in each block. If this parameter is not configured (nullptr is passed in), query is fully compressed. | INT64 | - |
140+| actual_block_len_key | Input | Actual valid sequence length compressed within each key block, that is, `actualBlockLenKey` in the formula. It is used for partial compression scenarios (such as incomplete blocks at the end or compressing only valid tokens). The mean is computed only over the first actual_block_len_key tokens in each block. If this parameter is not configured (nullptr is passed in), key is fully compressed. | INT64 | - |
141+| q_input_layout | Input | Data layout format of query. It indicates the specific arrangement of the input tensors in memory. Currently only "TND" and "BNSD" are supported, and it must be consistent with kv_input_layout. | String | - |
142+| kv_input_layout | Input | Data layout format of key. It indicates the specific arrangement of the input tensors in memory. Currently only "TND" and "BNSD" are supported, and it must be consistent with q_input_layout. | String | - |
143+| num_key_value_heads | Input | Number of attention heads of key, that is, the number of heads N of the key tensor in the formula. Currently only MHA is supported, and it must be consistent with the number of heads of query. | Int | - |
144+| scale_value | Input | Scaling factor, that is, `scale` in the formula. It is used to normalize the attention scores and is generally set to 1 / sqrt(D). | Float | - |
145+| sparsity | Input | Sparsity retention ratio, that is, `sparsity` in the formula. It specifies the proportion of block positions to be retained in attn_score out of all block positions. The value range is (0.0, 1.0). | Float | - |
146+| block_sparse_mask_out | Output | Block sparse mask output, that is, `blockSparseMaskOut` in the formula. It indicates the sparsity pattern adaptively generated based on the Q/K content, and can be directly used as the blockSparseMask input of the BSA operator. The shape changes with post_block_shape: when post_block_shape is not null, the shape is [B, N, postXBlocks, postYBlocks] (coarse-grained), where postXBlocks = ceilDiv(Xblocks, postBlockShapeX) and postYBlocks = ceilDiv(Yblocks, postBlockShapeY); when post_block_shape is null, the shape is [B, N, Xblocks, Yblocks] (fine-grained). The value 1 indicates that the block participates in the attention computation, and the value 0 indicates that it does not. | INT8 | ND |
147+ 
148+## Constraints
149+ 
150+- actual_seq_lengths is required when q_input_layout is "TND"; actual_seq_lengths_kv is required when kv_input_layout is "TND".
151+- The size of the head dimension in the query tensor shape is denoted as N1, and the size of the head dimension in the key tensor shape is denoted as N2. N1 = N2 must be satisfied (only MHA is supported).
152+- headDim = 128.
153+- blockShapeX and blockShapeY in block_shape must be multiples of 8.
154+- After query and key are compressed, the corresponding Xblocks and Yblocks of query and key must satisfy Xblocks * Yblocks > 1. When post_block_shape is not null, the BNSD scenario further requires postXBlocks * postYBlocks > 1 (when the coarse-grained grid collapses to 1×1, there is no TopK selection semantics, and the host side rejects the input; the TND variable-length scenario allows it).
155+- The data types of query and key must be consistent, and only FLOAT16 and BFLOAT16 are supported.
156+- The data type of block_sparse_mask_out is INT8 (binary: 0 or 1).
157+- post_block_shape is an optional input. When nullptr is passed in, the second-round pooling is disabled (backward compatibility). When `[postBlockShapeX, postBlockShapeY]` is passed in, the second-round pooling is enabled, and postBlockShapeX/postBlockShapeY are multiples of 8, indicating the number of block groups.
158+- If actual_block_len_query / actual_block_len_key is not null, the value range of each element is [0, blockShapeX] / [0, blockShapeY]. If it is null, full compression is performed.
159+ 
160+## Calling Instructions
161+ 
162+| Calling Method | Calling Sample | Description |
163+|---------|---------|------|
164+| aclnn API | [test_aclnn_bsa_select_block_mask](examples/test_aclnn_bsa_select_block_mask.cpp) | Only MHA and scenarios with the TND\BNSD format are supported. The BSASelectBlockMask operator is called through the [aclnnBSASelectBlockMask](docs/aclnnBSASelectBlockMask_en.md) API. |
Aattention/chunk_gated_delta_rule/README_en.md+157-0
@@ -0,0 +1,157 @@
1+# ChunkGatedDeltaRule
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support
6+ 
7+|Product | Supported |
8+|:----------------------------|:-----------:|
9+|<term>Ascend 950PR&950DT products</term>| √ |
10+|<term>Atlas A3 products</term>| √ |
11+|<term>Atlas A2 products</term>| √ |
12+|<term>Atlas 200I/500 A2 inference products</term>| × |
13+|<term>Atlas inference products</term>| × |
14+|<term>Atlas training products</term>| × |
15+ 
16+## Function Description
17+ 
18+- Operator function: Performs the chunked Gated Delta Rule computation.
19+ 
20+- Computation formula:
21+ 
22+ The Gated Delta Rule (GDR) is an operator applied in recurrent neural networks and also applied in a linear attention mechanism. At each time step $t$, GDR computes the current attention output $o_t$ and the new hidden state $S_t$ based on the current inputs $q_t$, $k_t$, and $v_t$, the previous hidden state $S_{t-1}$, the decay coefficient $\alpha_t$, and the update strength $\beta_t$. The computation formulas are as follows:
23+ $$
24+ S_t := S_{t-1}(\alpha_t(I - \beta_t k_t k_t^T)) + \beta_t v_t k_t^T = \alpha_t S_{t-1} + \beta_t (v_t - \alpha_t S_{t-1}k_t)k_t^T
25+ $$
26+ $$
27+ o_t := S_t (q_t \cdot scale)
28+ $$
29+ 
30+ Here $S_{t-1},S_t \in \mathbb{R}^{D_v \times D_k}$, $q_t, k_t \in \mathbb{R}^{D_k}$, $v_t \in \mathbb{R}^{D_v}$, $\alpha_t \in \mathbb{R}$, $\beta_t \in \mathbb{R}$, and $o_t \in \mathbb{R}^{D_v}$.
31+ 
32+ Chunked Gated Delta Rule is the chunked implementation of GDR ([reference paper](https://arxiv.org/abs/2412.06464)). It divides the input sequence into chunks to achieve a certain parallelization effect. In long-context scenarios, its computation efficiency is higher than that of Recurrent Gated Delta Rule, making it suitable for the prefill stage. Given an input sequence of length L, this operator computes the output $o_t, t \in \{1, 2, .., L\}$ at each step and the final state matrix $S_L$.
33+ 
34+## Parameter Description
35+ 
36+<table style="undefined;table-layout: fixed; width: 900px"><colgroup>
37+<col style="width: 180px">
38+<col style="width: 120px">
39+<col style="width: 200px">
40+<col style="width: 300px">
41+<col style="width: 100px">
42+</colgroup>
43+<thead>
44+ <tr>
45+ <th>Parameter</th>
46+ <th>Input/Output/Attribute</th>
47+ <th>Description</th>
48+ <th>Data Type</th>
49+ <th>Data Format</th>
50+ </tr></thead>
51+<tbody>
52+ <tr>
53+ <td>query</td>
54+ <td>Input</td>
55+ <td>q in the formula.</td>
56+ <td>BFLOAT16</td>
57+ <td>ND</td>
58+ </tr>
59+ <tr>
60+ <td>key</td>
61+ <td>Input</td>
62+ <td>Input k in the formula.</td>
63+ <td>BFLOAT16</td>
64+ <td>ND</td>
65+ </tr>
66+ <tr>
67+ <td>value</td>
68+ <td>Input</td>
69+ <td>Input v in the formula.</td>
70+ <td>BFLOAT16</td>
71+ <td>ND</td>
72+ </tr>
73+ <tr>
74+ <td>beta</td>
75+ <td>Input</td>
76+ <td>β in the formula.</td>
77+ <td>BFLOAT16</td>
78+ <td>ND</td>
79+ </tr>
80+ <tr>
81+ <td>initial_state</td>
82+ <td>Input</td>
83+ <td>Initial state matrix, that is, the input S_0 in the formula.</td>
84+ <td>BFLOAT16, FLOAT32</td>
85+ <td>ND</td>
86+ </tr>
87+ <tr>
88+ <td>actual_seq_lengths</td>
89+ <td>Input</td>
90+ <td>Sequence length of each batch.</td>
91+ <td>INT32</td>
92+ <td>ND</td>
93+ </tr>
94+ <tr>
95+ <td>g</td>
96+ <td>Optional input</td>
97+ <td>Decay coefficient, α=e^g in the formula</td>
98+ <td>FLOAT32</td>
99+ <td>ND</td>
100+ </tr>
101+ <tr>
102+ <td>out</td>
103+ <td>Output</td>
104+ <td>Attention result at each step, that is, o_t in the formula.</td>
105+ <td>BFLOAT16</td>
106+ <td>ND</td>
107+ </tr>
108+ <tr>
109+ <td>final_state</td>
110+ <td>Output</td>
111+ <td>Final state matrix, that is, S_L in the formula.</td>
112+ <td>BFLOAT16, FLOAT32</td>
113+ <td>ND</td>
114+ </tr>
115+ <tr>
116+ <td>scale_value</td>
117+ <td>Optional attribute</td>
118+ <td>Scaling factor of query, that is, scale in the formula. The default value is 1.0</td>
119+ <td>FLOAT</td>
120+ <td>-</td>
121+ </tr>
122+</tbody>
123+</table>
124+ 
125+## Constraints
126+ 
127+- To facilitate understanding of the layout formats later (such as BNSD and TND), the meaning of each abbreviated dimension is uniformly described as follows:
128+ * B: Batch size of the input samples (Batch).
129+ * T: If $L_i$ is the length of the $i$-th sequence, $T=\sum_i^B L_i$ indicates the cumulative sequence length.
130+ * Nk: Number of query and key heads.
131+ * Nv: Number of value heads.
132+ * Dk: Hidden layer dimension of query and key.
133+ * Dv: Hidden layer dimension of value.
134+ 
135+- Currently only the TND layout is supported:
136+ - Shapes of query and key: $(T, Nk, Dk)$
137+ - Shapes of value and out: $(T, Nv, Dv)$
138+ - Shapes of beta and g: $(T, Nv)$
139+ - Shape of actual_seq_lengths: $(B,)$
140+ - Shapes of initial_state and final_state: $(B, Nv, Dv, Dk)$
141+ 
142+ The dimensions must satisfy the following constraints:
143+ - $0 \lt Nv \le 64, 0 \lt Nk \le 64$, and $Nv \bmod Nk = 0$
144+ - $0 \lt Dv \le 128, 0 \lt Dk \le 128$
145+ - $B \gt 0, T \gt 0$
146+ 
147+- Due to the numerical characteristics of the algorithm, the following value constraints must be satisfied; otherwise, numerical overflow and precision anomalies easily occur:
148+ - Tensor elements: $-1 \le query[i][j][k] \le 1$
149+ - Tensor elements: $-1 \le key[i][j][k] \le 1$
150+ - Tensor elements: $-1 \le g[i][j] \le 0$
151+ - Tensor elements: $0 < beta[i][j] < 1$
152+ 
153+## Calling Instructions
154+ 
155+| Calling Method | Sample Code | Description |
156+| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ |
157+| aclnn | [test_aclnn_chunk_gated_delta_rule.cpp](./examples/test_aclnn_chunk_gated_delta_rule.cpp) | Calls the aclnnChunkGatedDeltaRule operator through [aclnnChunkGatedDeltaRule](docs/aclnnChunkGatedDeltaRule_en.md) |
Aattention/chunk_gated_delta_rule/tests/assets/README_en.md+209-0
@@ -0,0 +1,209 @@
1+# TTK Test Assets for the Chunk_gated_delta_rule Operator
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Function Description
6+ 
7+Based on the ops-test-kit (TTK) test framework, this directory implements E2E and ACLNN precision verification for the Chunk_gated_delta_rule operator:
8+ 
9+- **E2E mode**: Obtains actual NPU data through direct calls to the TorchNPU Python API (`torch_npu.npu_chunk_gated_delta_rule`) or the torch.compile graph mode, and compares it with the CPU golden data
10+- **ACLNN mode**: Calls the operator through the aclnn C API (`aclnnChunkGatedDeltaRule`) and compares it with the CPU golden data
11+- **Precision comparison**: Three-way cross-check (NPU vs golden vs benchmark), supporting multiple execution modes including eager and aclgraph
12+ 
13+## Current Implementation Scope
14+ 
15+### Supported Test Modes
16+ 
17+| Mode | API | Description |
18+|------|-----|------|
19+| E2E eager | `torch_npu.npu_chunk_gated_delta_rule` | Direct torch call, including profiling |
20+| E2E aclgraph | Same as above | torch.compile reduce-overhead mode (npugraph_ex backend) |
21+| ACLNN | `aclnnChunkGatedDeltaRule` | aclnn C API call |
22+ 
23+> **Note**: The const graph (torchair GE backend) is currently not supported. The operator is not registered in the GE graph engine, and execution reports `Interface npu_chunk_gated_delta_rule is not supported in GE graph mode`. The aclgraph mode uses the `npugraph_ex` backend, captures graphs at the aclnn level, and can run properly.
24+ 
25+### Precision Standard
26+ 
27+- The tolerance declares the community standard `cross_check` (bfloat16) with level `L1`.
28+- Comparison determination conditions (`impl/compare.py`): three-way cross-check. Four ratios (NPU error / benchmark error) are computed for mare/mere/rmse/smra, with thresholds `CV_MAX_RE=5.0` / `CV_AVER_RE=1.5` / `CV_RMSE=1.5` / `CV_SMALL_VAL=2.0` respectively. The result is Pass only when all ratios are below the thresholds. The floor of the ratio denominator is `2^-8` (`_ERR_THRESHOLD`).
29+ 
30+### Environment Configuration
31+ 
32+#### Prerequisites
33+ 
34+1. TorchNPU installation package download path (replace it with the latest version in time): [TorchNPU Installation Tutorial](https://atomgit.com/Ascend/pytorch)
35+2. Complete the environment installation and environment variable configuration. For details, refer to [ops-transformer](../../../../README_en.md).
36+3. ops-test-kit test framework (TTK): [ops-test-kit](https://atomgit.com/cann/ops-test-kit)
37+ 
38+## File Structure
39+ 
40+```
41+attention/chunk_gated_delta_rule/tests/assets/
42+├── convert_rdv_to_csv.py # Script for converting all RDV cases to TTK CSV (E2E + ACLNN)
43+├── spec.py # TestSpec: registers API + golden/inputs/tolerance/compare (e2e + aclnn)
44+└── impl/
45+ ├── golden.py # golden adapter (e2e + aclnn), loads CPU golden + third-party benchmark
46+ ├── inputs.py # inputs adapter (e2e + aclnn), fills actual_seq_lengths + applies constraint clipping
47+ └── compare.py # Three-way cross-check precision comparison (output + state dual outputs)
48+```
49+ 
50+### File Responsibilities
51+ 
52+#### spec.py
53+ 
54+- Defines the `ChunkGatedDeltaRuleSpec` class (e2e), which declares `golden`, `customize_inputs`, `tolerance`, and `compare`.
55+- Defines the `AclnnChunkGatedDeltaRuleSpec` class (aclnn), which declares `golden`, `customize_inputs`, `tolerance`, and `compare`.
56+- Registers through the `__spec__` dict:
57+ - `torch_npu.npu_chunk_gated_delta_rule` → `ChunkGatedDeltaRuleSpec`
58+ - `aclnnChunkGatedDeltaRule` → `AclnnChunkGatedDeltaRuleSpec`
59+- `compare` is a custom three-way comparison. The golden computation also computes the benchmark result and passes it to compare through `_GOLDEN_CONTEXT`.
60+- Dynamically loads the three modules under `impl/` to avoid hardcoded paths.
61+ 
62+#### impl/golden.py
63+ 
64+- Loads the CPU golden implementation and the third-party benchmark implementation through `importlib` without re-implementing them.
65+- **e2e adapter** `cpu_chunk_gated_delta_rule`: uses the `*args, **kwargs` signature, is compatible with positional and keyword arguments according to `_PARAM_ORDER`, and returns `[out, final_state]`.
66+- **aclnn adapter** `aclnn_chunk_gated_delta_rule_golden`: aligns the parameter order with the `aclnnChunkGatedDeltaRuleGetWorkspaceSize` function signature, returns `[out_golden, finalState_golden]`, and keeps the order consistent with `output_tensor_indexes`.
67+- The golden and benchmark results are computed in the same function `_compute_and_store`, and the benchmark result is stored in `_GOLDEN_CONTEXT` for compare to use.
68+- **prepare persistence** `_save_bench`: when golden data is dumped with `--dump in,golden`, the benchmark data is simultaneously saved to the sibling directory `<case_name>.bench/` in the manual-data directory (the TTK case directory schema is fixed and does not allow extra files); the format is consistent with the TTK golden bin: raw bytes, with dtype/shape encoded in the file name (`bench_<i>_<dtype>__shape_<DxN>.bin`, reusing `dump_to_file`/`load_numpy_data`).
69+- **replay loading** `resolve_bench`: during full bin replay, the golden function is skipped and `_GOLDEN_CONTEXT` is empty. In this case, the benchmark is loaded from the `.bench/` directory, consistent with the behavior of loading golden from the bin, and neither is recomputed. For legacy data without a `.bench` directory, the fallback is to recompute from inputs restored from the bin through `compare_context` (the replay-safe input context injected by TTK), with a warning in the log. Both paths are cached by `testcase_name` to prevent data mixing across cases.
70+ 
71+#### impl/inputs.py
72+ 
73+- Fills the int32 tensor that cannot be generated randomly and applies the operator numerical constraints:
74+ - `actual_seq_lengths`: evenly distributes T across B batches.
75+ - q/k: L2 normalization (the operator requires pre-normalized inputs).
76+ - `g`: clipped to `[-1, 0]`.
77+ - `beta`: clipped to `(0, 1)`.
78+ 
79+#### impl/compare.py
80+ 
81+- Three-way cross-check: NPU output vs golden output vs benchmark output, supporting dual outputs of output + finalState.
82+- Computes four ratios (NPU error / benchmark error) for mare/mere/rmse/smra. The thresholds and floor values are exactly the same as those of `compare_cv` on the operator side.
83+- Returns a dict structure containing `pass`/`precision`/`error_info`/`metrics` (including the four ratios and the raw error values).
84+ 
85+### convert_rdv_to_csv.py
86+ 
87+- Converts all RDV cases (38 redline cases + 34 STC cases, doubled after FP32 state and non-contiguous state extension) to TTK CSV.
88+- Covers combinations of contiguous/non-contiguous state and bf16/fp32 state dtypes. Non-contiguous state is described through `tensor_storage_shapes`/`tensor_view_strides`/`tensor_view_offsets`.
89+- Generates two CSV files: one for E2E mode (`chunk_gated_delta_rule_rdv.csv`) and one for ACLNN mode (`aclnn_chunk_gated_delta_rule_rdv.csv`).
90+ 
91+```bash
92+cd attention/chunk_gated_delta_rule/tests/assets
93+python3 convert_rdv_to_csv.py
94+```
95+ 
96+## Usage
97+ 
98+Before use, prepare the CSV case files for E2E and ACLNN modes and define the path variables:
99+ 
100+```bash
101+TTK_DIR=<ops-test-kit path>
102+ASSETS=attention/chunk_gated_delta_rule/tests/assets
103+CSV_E2E=<E2E CSV case path>
104+CSV_ACLNN=<ACLNN CSV case path>
105+```
106+ 
107+### E2E Mode (eager + aclgraph)
108+ 
109+```bash
110+cd $TTK_DIR
111+python3 -m ttk e2e \
112+ -i $CSV_E2E \
113+ --plugin $ASSETS \
114+ --aclgraph \
115+ -o <result output path>
116+```
117+ 
118+### ACLNN Mode
119+ 
120+```bash
121+cd $TTK_DIR
122+python3 -m ttk aclnn \
123+ -i $CSV_ACLNN \
124+ --plugin $ASSETS \
125+ -o <result output path>
126+```
127+ 
128+### Offline Data Dump to Bin and Replay from Specified Bin (Two-Stage Execution)
129+ 
130+Splits input generation/CPU golden generation and device execution into two steps: the prepare stage only generates and saves input/golden to bin files (without running the device), and the replay stage restores data from the bin files and then runs the device + compare. This applies to scenarios where the preparation machine and the execution machine are separate, or where fixed inputs need to be reused for repeated runs.
131+ 
132+```bash
133+MANUAL_DIR=/tmp/cgdr_manual # Root directory for storing bin data
134+OUT=<result output CSV path>
135+```
136+ 
137+#### 1. prepare (dump bin)
138+ 
139+Full data (input + golden):
140+ 
141+```bash
142+# E2E
143+cd $TTK_DIR
144+python3 -m ttk e2e \
145+ -i $CSV_E2E --plugin $ASSETS \
146+ --no-prof --dump in,golden --dump-format bin \
147+ --manual-data-dirs $MANUAL_DIR --pc 1 \
148+ -o $OUT
149+ 
150+# ACLNN
151+python3 -m ttk aclnn \
152+ -i $CSV_ACLNN --plugin $ASSETS \
153+ --no-prof --dump in,golden --dump-format bin \
154+ --manual-data-dirs $MANUAL_DIR --pc 1 \
155+ -o $OUT
156+```
157+ 
158+Prepare input only (golden is recomputed during replay):
159+ 
160+```bash
161+python3 -m ttk e2e \
162+ -i $CSV_E2E --plugin $ASSETS \
163+ --no-prof --dump in --dump-format bin \
164+ --manual-data-dirs $MANUAL_DIR --pc 1 \
165+ -o $OUT
166+```
167+ 
168+#### 2. replay (run from specified bin)
169+ 
170+```bash
171+# E2E
172+python3 -m ttk e2e \
173+ -i $CSV_E2E --plugin $ASSETS \
174+ --manual-data-dirs $MANUAL_DIR --pc 1 \
175+ -o $OUT
176+ 
177+# ACLNN
178+python3 -m ttk aclnn \
179+ -i $CSV_ACLNN --plugin $ASSETS \
180+ --manual-data-dirs $MANUAL_DIR --pc 1 \
181+ -o $OUT
182+```
183+ 
184+- The success status of prepare is `MANUAL_DATA_PREPARED`. When replay hits, random input generation is skipped, and input (and golden) are restored from the bin. The log shows `OnLoadManualData: loaded prepared input/scalar data from $MANUAL_DIR/...`.
185+- **The three-way benchmark is persisted together with golden and is not recomputed during replay**: during prepare with `--dump in,golden`, the benchmark is saved to `$MANUAL_DIR/<case_name>.bench/bench_<i>_<dtype>__shape_<DxN>.bin`; during replay, golden is loaded from the bin and the benchmark is loaded from that directory, with consistent behavior. For legacy data without a `.bench` directory, the fallback is recomputation from bin inputs (with a log warning); in input-only mode and direct mode, the benchmark is computed as usual during the golden stage.
186+- `--pc 1` runs serially in a single process. A single process is recommended for this operator.
187+- Both stages must use the same CSV and assets. After modifying the input shape/view or attributes, run prepare again.
188+- **`--plugin` is required**: both stages depend on spec.py to provide the golden/inputs/compare adapters.
189+- **`--no-prof` uses double dashes**: writing `-no-prof` reports `unrecognized arguments: -no-prof`.
190+ 
191+### Key Parameters
192+ 
193+| Parameter | Function | Value for This Operator |
194+|------|------|-----------|
195+| `-o FILE` | Output result CSV (including time consumption columns) | Recommended when profiling is enabled |
196+| `--aclgraph` | Tests the E2E aclgraph mode (reduce-overhead) | Recommended; const graph is not supported |
197+| `-c` | Tests the E2E const graph mode (GE backend) | Not supported; the operator is not registered in GE |
198+| `--pc N` | Number of concurrent processes | `--pc 1` is recommended for this operator |
199+| `--no-prof --dump in[,golden]` | Dumps bin data in the prepare stage without running the device | For details, refer to the "Offline Data Dump to Bin" section |
200+| `--manual-data-dirs DIR` | Restores bin data from DIR during replay | Keep DIR consistent with the one used in prepare |
201+ 
202+### CSV Format Validation Only
203+ 
204+```bash
205+# E2E
206+python3 -m ttk e2e -i $CSV_E2E --validate
207+# ACLNN
208+python3 -m ttk aclnn -i $CSV_ACLNN --plugin $ASSETS --validate
209+```
Aattention/chunk_gated_delta_rule/tests/pytest/README_en.md+171-0
@@ -0,0 +1,171 @@
1+# Chunk_gated_delta_rule Operator Test Framework
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Function Description
6+ 
7+Based on the pytest test framework, this directory implements functional verification for the Chunk_gated_delta_rule operator:
8+ 
9+- **CPU side**: Reproduces the operator function to generate golden data
10+- **NPU side**: Obtains actual data through direct operator calls using TorchNPU
11+- **Precision comparison**: Compares the precision of CPU and NPU results to verify the operator function
12+ 
13+## Current Implementation Scope
14+ 
15+### Parameter Restrictions
16+ 
17+- batch_size greater than 0 is supported.
18+- The seqlen sequence length is supported.
19+- NK and NV head counts are supported, and NV must be a multiple of NK.
20+- DK and DV no greater than 128 are supported.
21+- data_type BF16 is supported.
22+ 
23+### Environment Configuration
24+ 
25+#### Prerequisites
26+ 
27+1. TorchNPU installation package download path (replace it with the latest version in time): [TorchNPU Installation Tutorial](https://atomgit.com/Ascend/pytorch)
28+2. Complete the environment installation and environment variable configuration. For details, refer to [ops-transformer](../../../../README_en.md).
29+ 
30+#### custom Package Invocation
31+ 
32+Custom package invocation is supported.
33+ 
34+## File Structure
35+ 
36+### pytest File Structure Description
37+ 
38+- test_run.sh # Execution script
39+- conftest.py # pytest hook: records parameters/results/seed for each case, writes CSV at session end
40+- chunk_gated_delta_rule_golden.py # CPU-side golden implementation of the operator
41+- chunk_gated_delta_rule_main.py # Precision comparison between CPU golden and NPU results, and main test logic
42+- pytest.ini # Creates test markers for CI single-operator and graph modes
43+ 
44+Single-case testing:
45+ 
46+- test_chunk_gated_delta_rule_single.py # Main program for testing single-case runs
47+- chunk_gated_delta_rule_operator_single.py # CPU-side operator logic implementation to obtain golden and direct NPU operator calls
48+- test_chunk_gated_delta_rule_paramset.py # Single-case input parameter configuration
49+- test_chunk_gated_delta_rule_paramset_rdv.py # RDV test input parameter configuration
50+ 
51+## Usage
52+ 
53+Run the commands in the pytest directory:
54+ 
55+### Running Test Cases
56+ 
57+#### Single-Case Debugging
58+ 
59+1. Manually configure the parameters in test_chunk_gated_delta_rule_paramset.py.
60+ 
61+2. Run the command:
62+ 
63+``` bash
64+bash test_run.sh single
65+```
66+ 
67+#### RDV Testing
68+ 
69+1. Manually configure the parameters in test_chunk_gated_delta_rule_paramset_rdv.py.
70+ 
71+2. Run the command:
72+ 
73+``` bash
74+bash test_run.sh rdv
75+```
76+ 
77+#### Random Case Testing
78+ 
79+Randomly generates N cases and executes them (including CPU golden precision comparison). You can fix the random seed with the `RANDOM_SEED` environment variable for reproduction (if not specified, the seed is automatically generated and recorded in the CSV):
80+ 
81+``` bash
82+bash test_run.sh random 100
83+```
84+ 
85+#### Random Case Testing (NPU Only)
86+ 
87+Randomly generates N cases. Setting `SKIP_GOLDEN=1` skips CPU golden computation and precision comparison and executes the NPU operator only, which speeds up execution:
88+ 
89+``` bash
90+bash test_run.sh random_npu 100
91+```
92+ 
93+#### Random Case Generation Rules
94+ 
95+Both random/random_npu modes use the same random parameter generator (`_generate_random_param_dict`), which randomly generates from 0 within the operator constraints and does not depend on the single/rdv parameter pools. The generation rules for the input parameters of each case are as follows:
96+ 
97+| API Input Parameter | Random Rule | Constraints/Description |
98+|----------|---------|-----------|
99+| B | choice([1,2,4,8,16,32,64,128]) | batch size |
100+| seqlen | choice([1,3,7,32,64,100,128,200,256,300,512,1000,1024,2048,4096,5000,8192,10000,16384,32768,65535]) | Sequence length; includes non-chunk-aligned values (such as 100, 300, 5000, 65535) to test the partial chunk path |
101+| seqlen (variable length) | With B>1, a list is generated with a 30% probability, and each batch has an independent random seqlen | Covers the actual_seq_lengths variable-length path |
102+| nk | randint(1,64) | Number of key heads |
103+| nv | nk × randint(1,64//nk) | Nv>=Nk and Nv%Nk==0 |
104+| dk | randint(1,min(128,budget)) | key dimension, constrained by the state element upper limit |
105+| dv | randint(1,min(128,budget//dk)) | value dimension, constrained by the state element upper limit |
106+| chunk_size | Fixed 64 | Hardcoded in operator tiling; golden must be consistent with NPU |
107+| data_type | Fixed bfloat16 | |
108+| state_data_type | choice([bfloat16,float32]) | State data type |
109+| has_g | choice([True,False]) | Gating is enabled with a 50% probability |
110+| is_contiguous | choice([True,False]) | Non-contiguous with a 50% probability |
111+| query_datarange | Fixed [-1,1] | q is L2-normalized |
112+| key_datarange | Fixed [-1,1] | k is L2-normalized |
113+| value_datarange | choice([-10,10], [-1,1]) | Random data range |
114+| gamma_datarange | choice([-1,0], [-0.5,0], [-0.1,0], [-1,-0.5]) | Document constraint [-1,0]; g decays through exp |
115+| beta_datarange | Fixed [0,1] | |
116+| state_datarange | Fixed [-10,10] | |
117+ 
118+**Shape constraints**: B>0, seqlen>0, 0<Nk<=64, 0<Nv<=64 with Nv>=Nk and Nv%Nk==0, 0<Dk<=128, 0<Dv<=128.
119+ 
120+**Memory constraints**:
121+- Dk×Dv is constrained by the state element count upper limit `_STATE_ELEM_CAP=2.0B` (`budget = STATE_ELEM_CAP // (B × Nv)`)
122+- T×max(Nk×Dk, Nv×Dv) is constrained by the QKV tensor element upper limit `_QKV_ELEM_CAP=1.0B` (`max_seqlen = QKV_ELEM_CAP // (B × qkv_per_token)`)
123+- The two constraints jointly prevent OOM in a single process
124+ 
125+**Random seed mechanism**:
126+- `RANDOM_SEED` controls the shape/parameter sequence (one seed corresponds to a deterministic set of N case parameters)
127+- `TORCH_SEED` controls tensor values (independent for each case; conftest automatically generates and records it in the CSV)
128+- If `RANDOM_SEED` is not set, it is automatically generated, written back to `os.environ`, and written to the CSV by conftest
129+- Reproduction: `RANDOM_SEED=<seed> bash test_run.sh random N`
130+ 
131+### Result Output and Reproduction
132+ 
133+After execution, all modes output to the `output/` directory (gitignored):
134+ 
135+- `run_<timestamp>.log`: complete execution log (displayed on screen through tee)
136+- `result_<timestamp>.csv`: per-case result table, with one case per row
137+ 
138+| Column | Description |
139+|----|------|
140+| random_seed | Random shape sequence seed (random mode; single/rdv use a fixed parameter set and have no such value) |
141+| seed | Tensor value seed of this case (recorded independently for each case) |
142+| test_name | Case name |
143+| test_mode | single/rdv/random |
144+| check_type | precision=with golden precision comparison / execution_only=NPU execution only |
145+| model | Execution mode (direct torch call/aclgraph) |
146+| status | pytest execution result (PASSED/FAILED/SKIPPED) |
147+| B...is_continue | All input parameters of this case |
148+| errmsg | Failure details (truncated to 2000 characters) |
149+| durations | Operator time consumption (prof mode only) |
150+ 
151+Reproducing failed cases:
152+ 
153+``` bash
154+# Batch reproduction (same shape sequence): take random_seed from the CSV
155+RANDOM_SEED=<random_seed> bash test_run.sh random N
156+ 
157+# Single-case value-level reproduction (same shape + same tensor values): take the input parameters and seed from the CSV
158+TORCH_SEED=<tensor_seed> bash test_run.sh random 1
159+```
160+ 
161+### Environment Variable Summary
162+ 
163+| Variable | Function | Applicable Modes |
164+|------|------|---------|
165+| RANDOM_SEED | Fixes the random shape sequence seed (if not set, it is automatically generated and recorded in the CSV) | random/random_npu |
166+| TORCH_SEED | Fixes the tensor value seed (if not set, it is automatically generated for each case and recorded in the CSV) | All |
167+| RANDOM_CASE_COUNT | Number of random cases (passed through by test_run.sh) | random series |
168+| SKIP_GOLDEN | =1 skips CPU golden and precision comparison, and executes NPU only | random_npu |
169+| CSV_FILE | Specifies the CSV output path (automatically set by test_run.sh) | All |
170+| CSV_APPEND | =1 appends to the CSV | All |
171+| USE_GRAPH | =true enables the aclgraph mode | All |
Aattention/chunk_kda_fwd/README_en.md+174-0
@@ -0,0 +1,174 @@
1+# ChunkKdaFwd
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Product Support Status
6+ 
7+| Product | Supported |
8+| :--- | :---: |
9+| <term>Ascend 950PR&950DT products</term> | √ |
10+| <term>Atlas A3 products</term> | √ |
11+| <term>Atlas A2 products</term> | √ |
12+| <term>Atlas 200I/500 A2 inference products</term> | × |
13+| <term>Atlas inference products</term> | × |
14+| <term>Atlas training products</term> | × |
15+ 
16+## Function Description
17+ 
18+- Operator function: Performs chunked forward KDA computation without CP splitting, and computes the attention output, the final state, and optional intermediate quantities for backward computation.
19+ 
20+- Computation formula:
21+ 
22+ Each sequence is divided into $M$ chunks with $C=chunk\_size$. Taking the c-th chunk as an example, $i$ and $j$ are token indices within the chunk, and $l_c$ is the index of the last valid token in the chunk. The formulas omit batch and head indices. When $HV>H$, the scenario is grouped value attention (GVA): each group of $HV/H$ consecutive Value heads shares one Query/Key head. When head numbering starts from 0, the Value head numbered $h_v$ corresponds to the Query/Key head numbered $\lfloor h_v/(HV/H)\rfloor$.
23+ 
24+ Let $x_{c,i,d}=g_{c,i,d}+dt\_bias_d$. If `dt_bias` is not passed, set $dt\_bias_d=0$. The gate after activation is:
25+ 
26+ $$
27+ \gamma_{c,i,d}=
28+ \begin{cases}
29+ g_{c,i,d}, & use\_gate\_in\_kernel=false,\\
30+ -\exp(A_{log})\operatorname{softplus}(x_{c,i,d}),
31+ & use\_gate\_in\_kernel=true,\ safe\_gate=false,\\
32+ lower\_bound\operatorname{sigmoid}(\exp(A_{log})x_{c,i,d}),
33+ & use\_gate\_in\_kernel=true,\ safe\_gate=true.
34+ \end{cases}
35+ $$
36+ 
37+ `gk` is the chunk-local log2 cumulative gate:
38+ 
39+ $$
40+ gk_{c,i,d}=\frac{1}{\ln 2}\sum_{t=0}^{i}\gamma_{c,t,d}.
41+ $$
42+ 
43+ `qg` and `kg` are the gate-scaled Query and Key respectively:
44+ 
45+ $$
46+ qg_{c,i,d}=q_{c,i,d}2^{gk_{c,i,d}},\qquad
47+ kg_{c,i,d}=k_{c,i,d}2^{gk_{c,l_c,d}-gk_{c,i,d}}.
48+ $$
49+ 
50+ For tokens $i$ and $j$ in the same chunk, `Aqk` is the lower-triangular Query-Key coefficient matrix including the diagonal:
51+ 
52+ $$
53+ Aqk_{c,i,j}=\mathbb{1}_{j\le i}\cdot scale
54+ \sum_d q_{c,i,d}k_{c,j,d}2^{gk_{c,i,d}-gk_{c,j,d}}.
55+ $$
56+ 
57+ The strictly lower-triangular matrix $L_c$ and `Akk` are:
58+ 
59+ $$
60+ L_{c,i,j}=\mathbb{1}_{j<i}\cdot\beta_{c,i}
61+ \sum_d k_{c,i,d}k_{c,j,d}2^{gk_{c,i,d}-gk_{c,j,d}},\qquad
62+ Akk_c=(I+L_c)^{-1}.
63+ $$
64+ 
65+ Let $w^{seed}$ and $u^{seed}$ be:
66+ 
67+ $$
68+ w^{seed}_{c,i,d}=\beta_{c,i}k_{c,i,d}2^{gk_{c,i,d}},\qquad
69+ u^{seed}_{c,i,e}=\beta_{c,i}v_{c,i,e}.
70+ $$
71+ 
72+ `w` and `u` are:
73+ 
74+ $$
75+ w_c=Akk_c\,w^{seed}_c,\qquad u_c=Akk_c\,u^{seed}_c.
76+ $$
77+ 
78+ Let $h_c$ be the state before the computation of the c-th chunk. If `initial_state` is passed, $h_0=initial\_state$; otherwise, $h_0=0$. `v_new` and the state of the next chunk are:
79+ 
80+ $$
81+ v^{new}_c=u_c-w_c\,h_c,
82+ $$
83+ 
84+ $$
85+ h_{c+1}=2^{gk_{c,l_c}}\odot h_c+kg_c^T\,v^{new}_c,
86+ $$
87+ 
88+ where $2^{gk_{c,l_c}}\odot h_c$ denotes scaling the state matrix along the $K$ dimension. The `h` output stores the input state of each chunk, and `final_state` is the state after the last chunk is updated:
89+ 
90+ $$
91+ h[c]=h_c,\qquad final\_state=h_M.
92+ $$
93+ 
94+ `attn_out` is:
95+ 
96+ $$
97+ attn\_out_c=scale\cdot qg_c\,h_c+Aqk_c\,v^{new}_c.
98+ $$
99+ 
100+ The 12th return value of the Python interface is the pass-through of the input object:
101+ 
102+ $$
103+ initial\_state\_out=initial\_state.
104+ $$
105+ 
106+ The optional output parameters only determine whether the corresponding intermediate quantities are returned to the caller, and do not change the computation semantics described above.
107+ 
108+## Parameter Description
109+ 
110+This document uses the following dimension symbols:
111+ 
112+- B: batch size of input samples.
113+- S: sequence length of rank-4 inputs.
114+- T: total number of tokens of rank-3 inputs.
115+- H: number of Query and Key heads.
116+- HV: number of Value and gate heads.
117+- K: head dim of Query and Key.
118+- V: head dim of Value.
119+- N: number of logical sequences.
120+- C: chunk size.
121+- NC: number of chunks per sequence ($\lceil S/C\rceil$) when `cu_seqlens` is not passed; total number of chunks across all logical sequences when `cu_seqlens` is passed.
122+ 
123+| Parameter Name | Input/Output/Attribute | Description | Data Type | Data Format |
124+| :--- | :--- | :--- | :--- | :--- |
125+| q | Input | Query in the formula. `BSND` is (B, S, H, K), `BNSD` is (B, H, S, K), `TND` is (T, H, K), `NTD` is (H, T, K). | FLOAT16, BFLOAT16 | ND |
126+| k | Input | Key in the formula. The shape and data type must be the same as those of `q`. | FLOAT16, BFLOAT16 | ND |
127+| v | Input | Value in the formula. `BSND` is (B, S, HV, V), `BNSD` is (B, HV, S, V), `TND` is (T, HV, V), `NTD` is (HV, T, V). The data type must be the same as that of `q`. | FLOAT16, BFLOAT16 | ND |
128+| g | Input | Raw gate or activated natural-log gate. `BSND` is (B, S, HV, K), `BNSD` is (B, HV, S, K), `TND` is (T, HV, K), `NTD` is (HV, T, K). | FLOAT, BFLOAT16 | ND |
129+| beta | Input | Delta coefficient in the formula. `BSND` is (B, S, HV), `BNSD` is (B, HV, S), `TND` is (T, HV), `NTD` is (HV, T). | FLOAT, BFLOAT16 | ND |
130+| a_log | Optional input | Gate decay parameter with shape (HV). It must be passed when `use_gate_in_kernel=true`. | FLOAT | ND |
131+| dt_bias | Optional input | Gate bias with shape (HV×K). | FLOAT | ND |
132+| initial_state | Optional input | Initial state. The shape is (N, HV, K, V) when `state_v_first=false`, and (N, HV, V, K) otherwise. | FLOAT | ND |
133+| cu_seqlens | Optional input | Cumulative lengths of variable-length sequences, with shape (N+1). The first element is 0, the last element is T or S, and elements are monotonically non-decreasing. | INT64 | ND |
134+| chunk_indices | Optional input | Chunk indices with shape (2×NC), stored as `(seq_id, chunk_id)`, in sequence-major canonical order. When this parameter is passed, `cu_seqlens` must also be passed. | INT64 | ND |
135+| attn_out | Output | Attention output. rank-4 inputs always output the BSND format (B, S, HV, V); rank-3 inputs always output the TND format (T, HV, V). | FLOAT16, BFLOAT16 | ND |
136+| final_state | Optional output | Final state. The shape is (N, HV, K, V) when `state_v_first=false`, and (N, HV, V, K) otherwise. | FLOAT | ND |
137+| gk | Optional output | Chunk-local log2 cumulative gate. rank-4 inputs: (B, HV, S, K); rank-3 inputs: (HV, T, K). | FLOAT | ND |
138+| Aqk | Output | Query-Key coefficient matrix within a chunk. rank-4 inputs: (B, HV, S, C); rank-3 inputs: (HV, T, C). | FLOAT16, BFLOAT16 | ND |
139+| Akk | Output | Lower-triangular inversion result within a chunk. rank-4 inputs: (B, HV, S, C); rank-3 inputs: (HV, T, C). | FLOAT16, BFLOAT16 | ND |
140+| w | Optional output | W intermediate quantity for backward computation. rank-4 inputs: (B, HV, S, K); rank-3 inputs: (HV, T, K). | FLOAT16, BFLOAT16 | ND |
141+| u | Optional output | U intermediate quantity for backward computation. rank-4 inputs: (B, HV, S, V); rank-3 inputs: (HV, T, V). | FLOAT16, BFLOAT16 | ND |
142+| qg | Optional output | Gate-scaled Query for backward computation. rank-4 inputs: (B, HV, S, K); rank-3 inputs: (HV, T, K). | FLOAT16, BFLOAT16 | ND |
143+| kg | Optional output | Gate-scaled Key for backward computation. rank-4 inputs: (B, HV, S, K); rank-3 inputs: (HV, T, K). | FLOAT16, BFLOAT16 | ND |
144+| v_new | Optional output | Value intermediate quantity for backward computation. rank-4 inputs: (B, HV, S, V); rank-3 inputs: (HV, T, V). | FLOAT16, BFLOAT16 | ND |
145+| h | Optional output | Exposed chunk state. rank-4 inputs: (B, NC, HV, K, V); rank-3 inputs: (NC, HV, K, V). The last two dimensions are swapped when `state_v_first=true`. | FLOAT16, BFLOAT16 | ND |
146+| layout | Optional attribute | Input layout. `BSND`, `BNSD`, `TND`, and `NTD` are supported, with the default value `BSND`. This attribute only describes the input layout. | STRING | - |
147+| scale | Required attribute | Query scale coefficient. Usually $K^{-0.5}$. | FLOAT | - |
148+| chunk_size | Required attribute | Chunk size. 64 and 128 are supported, with the default value 64. | INT | - |
149+| safe_gate | Required attribute | Whether to use the bounded gate. The default value is false. | BOOL | - |
150+| lower_bound | Optional attribute | Lower bound of the bounded gate. When `safe_gate=true` and `use_gate_in_kernel=true`, the value range is [-5, 0), and the default value is -5.0. | FLOAT | - |
151+| use_gate_in_kernel | Required attribute | Whether to compute the activation from the raw gate inside the kernel. The default value is false. | BOOL | - |
152+| state_v_first | Optional attribute | Whether the last two dimensions of the state tensor are arranged as (V, K). The default value is false. | BOOL | - |
153+ 
154+## Constraints
155+ 
156+- H and HV must be positive integers and satisfy HV≥H, HV≤128, and HV % H=0.
157+- K and V are in the range [16, 256] and must be multiples of 16.
158+- Only 64 and 128 are supported for `chunk_size`.
159+- The data types of `q`, `k`, and `v` must be the same.
160+- When `use_gate_in_kernel=true`, `a_log` must be passed, and `dt_bias` does not need to be passed.
161+- When `safe_gate=true` and `use_gate_in_kernel=true`, the value range of `lower_bound` is [-5, 0).
162+- rank-4 variable-length inputs require B=1.
163+- Variable-length inputs support at most 1024 logical sequences.
164+- `layout` describes only the input. `attn_out` is always sequence-major; intermediate outputs for subsequent backward computation are always head-major; `final_state` and the exposed `h` are output as specified in the parameter table.
165+- Deterministic computation is supported by default.
166+ 
167+## Invocation Description
168+ 
169+| Invocation Mode | Sample Code | Description |
170+| :--- | :--- | :--- |
171+| aclnn API | [test_aclnn_chunk_kda_fwd.cpp](./examples/test_aclnn_chunk_kda_fwd.cpp) | Calls the ChunkKdaFwd operator through [aclnnChunkKdaFwd](docs/aclnnChunkKdaFwd_en.md). For compilation and running steps, see [Invocation Sample Description](./examples/README_en.md). |
172+| PyTorch API | [torchapi_chunk_kda_fwd](./docs/torchapi_chunk_kda_fwd.md) | Calls the ChunkKdaFwd operator through `cann_ops_transformer.ops.chunk_kda_fwd`. |
173+ 
174+For stage division, platform templates, and pipeline design, refer to [ChunkKdaFwd Operator Design Introduction](./docs/ChunkKdaFwd_design_en.md).
Rattention/chunk_kda_fwd/docs/ChunkKdaFwd算子设计介绍.md→attention/chunk_kda_fwd/docs/ChunkKdaFwd_design.md+0-0
File renamed without changes.
Aattention/compressor/tests/pytest/README_en.md+259-0
@@ -0,0 +1,259 @@
1+# Compressor Operator Test Framework
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Function Description
6+ 
7+Based on the pytest test framework, this framework implements functional verification of the Compressor operator:
8+ 
9+- **CPU side**: Reproduces the operator functionality to generate golden data
10+- **NPU side**: Directly calls the operator through torch_npu to obtain actual data
11+- **Accuracy comparison**: Compares the accuracy of CPU and NPU results to verify operator functionality
12+- **Dual-mode execution isolation**: Supports two batch modes: direct multi-process pytest execution and process isolation at the shell layer
13+- **Performance collection**: Supports mounting msprof to collect operator performance data and output a summary
14+ 
15+## Current Implementation Scope
16+ 
17+### Parameter Restrictions
18+ 
19+- Supports D of 128/512.
20+- Supports H from 1K to 10K, aligned to 512.
21+- Supports cmp_ratio of 2, 4, 8, 16, 32, 64, and 128.
22+- When cmp_ratio is 4/128, the following three cases are supported:
23+ - C4A: D=512, coff=2, cmp_ratio=4;
24+ - C4Li: D=128, coff=2, cmp_ratio=4;
25+ - C128A: D=512, coff=1, cmp_ratio=128.
26+ 
27+### Environment Configuration
28+ 
29+#### Prerequisites
30+ 
31+1. Ensure that torch_npu is the latest version
32+ 
33+2. Ensure that the compressor operator package and the torch extension have been built and installed (when the binary does not contain the compressor operator):
34+ 
35+``` bash
36+# ① Build the compressor operator run package at the root directory of the ops-transformer repository (--vendor_name can be customized to avoid overwriting existing vendors packages)
37+# ${CANN_PATH}: Root installation path of CANNToolkit (for example, /usr/local/Ascend/ascend-toolkit)
38+# ${SOC}: Target chip (for example, ascend910b)
39+source ${CANN_PATH}/set_env.sh
40+bash build.sh --pkg --soc=${SOC} --ops=compressor --vendor_name=cmpbit -j16
41+# Artifact: build_out/custom_opp_*.run. Install it to the CANN vendors directory (confirm interactively)
42+bash build_out/custom_opp_*.run --install-path=${CANN_PATH}/latest
43+ 
44+# ② Build and install the torch extension wheel (full package or single operator)
45+bash build.sh --torch_extension --ops=compressor --vendor_name=cmpbit
46+pip3 install --force-reinstall --no-deps build_out/cann_ops_transformer_cmpbit-*.whl
47+```
48+ 
49+3. Before running the tests, source the CANN environment variables and the operator vendor package environment (`${CANN_PATH}` is the root installation path of CANNToolkit):
50+ 
51+``` bash
52+source ${CANN_PATH}/set_env.sh
53+# You also need to source the vendor package environment script that provides the Compressor operator:
54+source ${CANN_PATH}/latest/vendors/<vendor>_transformer/bin/set_env.bash
55+```
56+ 
57+#### custom Package Invocation
58+ 
59+The custom package invocation is supported.
60+ 
61+## File Structure
62+ 
63+### pytest File Structure Description
64+ 
65+- test_run.sh # Execution script. Supports three commands: single/bitlevel/batch (bitlevel supports only A2/A3)
66+- batch_isolated_run.sh # Batch isolated execution script (process isolation at the shell layer + msprof performance collection)
67+- compressor_golden.py # Golden implementation of the operator on the CPU side and accuracy comparison
68+- collect_perf_data.py # msprof performance data collection and summarization
69+- pytest.ini # Creates test markers
70+ 
71+Single test case:
72+ 
73+- test_compressor_single.py # Main program for running a single test case
74+- compressor_operator_single.py # CPU-side operator logic implementation for obtaining golden data and directly calling the NPU operator
75+- test_compressor_paramset.py # Input parameter configuration for a single test case
76+ 
77+Bit-level consistency test:
78+ 
79+- test_compressor_bitlevel.py # Main program for bit-level consistency tests (pair/position test cases)
80+- compressor_bitlevel_golden.py # Deterministic data construction + direct NPU invocation + bit-level comparison support library
81+- test_compressor_bitlevel_paramset.py # Parameter configuration for bit-level consistency test cases
82+ 
83+Batch test:
84+ 
85+- test_compressor_batch.py # Main program for batch test execution and Excel file generation for saving results (isolation mode supported)
86+- ./batch/compressor_pt_loadprocess.py # Reads pt files and calls the operator to obtain NPU output
87+- ./batch/compressor_pt_save.py # Reads Excel files and generates test case pt files in batches
88+- ./batch/replace_path.py # Placeholder replacement for test_compressor_batch.py
89+ 
90+## Usage
91+ 
92+Run the commands in the pytest directory:
93+ 
94+### Running Test Cases
95+ 
96+#### Single Test Case Debugging
97+ 
98+1. Manually configure the ENABLED_PARAMS parameter in test_compressor_paramset.py
99+ 
100+2. Run the command:
101+ 
102+``` bash
103+bash test_run.sh single
104+```
105+ 
106+#### Bit-Level Consistency Debugging
107+ 
108+1. Manually configure the BITLEVEL_TEST_PARAMS parameter in test_compressor_bitlevel_paramset.py
109+ 
110+2. Run the command:
111+ 
112+``` bash
113+bash test_run.sh bitlevel
114+# Or directly
115+python3 -m pytest -rA -s test_compressor_bitlevel.py -v -m ci -W ignore::UserWarning -W ignore::DeprecationWarning
116+```
117+ 
118+3. Environment variables:
119+ 
120+| Variable | Default Value | Description |
121+|---|---|---|
122+| `CMP_TEST_DEVICE_ID` | `0` | Specifies the ID of the NPU device to run on |
123+| `CMP_BITLEVEL_STRICT` | `1` | `1`: The test case fails if bit-level results are inconsistent; `0`: Downgrades failure to a warning plus threshold check (used to distinguish bit-level tail differences from functional errors) |
124+ 
125+4. Test case description (see the comments at the head of test_compressor_bitlevel_paramset.py for details):
126+ 
127+| Test Case | Purpose |
128+|---|---|
129+| B=32/33, S=4 | tokenSize=128/132, FULL_LOAD → PERF template switching, cross-template bit-level consistency |
130+| B=96/97, S=4 | PERF dynamic D partitioning changes under 24 AIC cores, coff=2, and D=512 (dBaseSize 32↔64, hIdxStart changes) |
131+| B=31/32 | tokenSize=124/128, both FULL_LOAD, control group with the same template |
132+| Same B, with the target sample moved to the first, middle, or last position | Batch position consistency |
133+ 
134+Note: The `dyn` test cases in the `PerfDsplitCliff` group automatically derive the dynamic D partitioning cliff boundary pairs based on the number of AIC cores on the current device
135+(formula: `B_dyn = 128 * (AIC // (D/64)) / S`; B=96/97 on a 24-core device, and B=64/65 on the 910B3 device with 20 cores).
136+Duplicates with fixed test cases are automatically removed. Each test case performs both CPU golden accuracy verification and NPU result bit-level comparison.
137+ 
138+5. Baseline results:
139+ 
140+**(Problem record, L0 behavior)**: At the time of measurement, the operator did not yet include the batch consistency modification (deterministic level =0). This record documents the problem symptoms that this modification resolves.
141+Environment: Atlas 910B3, 20 AIC cores, CANN 9.1.0-beta.1.
142+ 
143+| Test Case | Bit-Level Result | Description |
144+|---|---|---|
145+| `FullLoadCtrl-pair_31_32` | Consistent | Both use the FULL_LOAD template; the control group passes |
146+| `FullLoadToPerf-pair_32_33` | **Inconsistent** | cmp_kv: 7/16384 elements differ by 1 BF16 ULP; state update: about 0.001% of elements differ in FP32 mantissa (100% within the threshold) |
147+| `PerfDsplitCliffDyn-pair_64_65` | **Inconsistent** | The cliff mechanism changes hIdxStart: cmp_kv: 13/32768 elements differ by 1 BF16 ULP; state update: about 1.3% of elements differ in FP32 mantissa (100% within the threshold) |
148+| `PerfDsplitCliff-pair_96_97` | Consistent | Both sides use dBaseSize=64 under 20 cores, without crossing the cliff |
149+| `PositionConsistency-pos_B32/33/96/97` | Consistent | Results are bit-level consistent when the target sample is at the first, middle, or last position |
150+ 
151+Conclusion (before the fix): **Bit-level inconsistencies exist both across templates (FULL_LOAD vs PERF) and across the PERF dynamic D partitioning cliff (dBaseSize 32↔64, which changes hIdxStart)**.
152+All differences are at the level of 1 BF16 ULP / FP32 accumulation tail difference (all CPU golden accuracy checks pass, with no functional errors);
153+batch position consistency is satisfied.
154+ 
155+**Measured after the fix (level=3)**: `bash test_run.sh bitlevel` (deterministic level =3) with the default configuration (`CMP_BITLEVEL_STRICT=1`).
156+Environment: Atlas 910B3, 20 AIC cores, CANN 9.2.0; 7/7 test cases pass.
157+ 
158+| Test Case | Bit-Level Result |
159+|---|---|
160+| `FullLoadCtrl-pair_31_32` | Consistent |
161+| `FullLoadToPerf-pair_32_33` | Consistent |
162+| `PerfDsplitCliff-pair_64_65` | Consistent |
163+| `PositionConsistency-pos_B32/33/96/97` | Consistent |
164+ 
165+Conclusion: **At level=3, both sides are bit-level consistent across templates and across the PERF dynamic D partitioning cliff**; batch position consistency is maintained.
166+In addition, on the generalized batch splitting suite (13 groups of production shapes and 388 splitting policies), all cases at level=3 are bit-level consistent.
167+Setting `CMP_BITLEVEL_STRICT=0` downgrades bit-level failures to warnings to distinguish tail differences from functional errors (the default is 1, that is, strict bit-level assertion).
168+ 
169+#### Batch Generation and Execution of Test Cases
170+ 
171+##### Option A: Batch Execution Using test_run.sh
172+ 
173+1. Place the test case Excel file in the excel directory
174+ 
175+2. Set the test case Excel file path to read (PATH1) and the pt file storage path (PATH2) in test_run.sh
176+ 
177+3. Run the command:
178+ 
179+``` bash
180+bash test_run.sh batch
181+```
182+ 
183+##### Option B: Manual Step-by-Step Execution
184+ 
185+1. Generate pt files:
186+ 
187+``` bash
188+python3 batch/compressor_pt_save.py "excel/*.xlsx" pt_path
189+```
190+ 
191+2. Replace the test script path:
192+ 
193+``` bash
194+python3 batch/replace_path.py test_compressor_batch.py pt_path
195+```
196+ 
197+3. Run the test:
198+ 
199+``` bash
200+python3 -m pytest -rA -s test_compressor_batch.py -v -m ci -W ignore::UserWarning -W ignore::DeprecationWarning
201+```
202+ 
203+4. Restore the test script:
204+ 
205+``` bash
206+cp test_compressor_batch.py.bak test_compressor_batch.py
207+```
208+ 
209+##### Option C: Batch Isolated Execution (Recommended for Performance Collection)
210+ 
211+This option starts a separate pytest process for each test case, achieving complete inter-process isolation and preventing a crashed test case from affecting other test cases.
212+Test case acquisition logic: The framework first reads the `Testcase_Name` column of the Excel file in the `./excel` directory and maps it to `.pt` file paths. If no Excel file exists or the `Testcase_Name` column is missing, the framework automatically falls back to scanning all `.pt` files in the test case directory.
213+ 
214+``` bash
215+bash batch_isolated_run.sh ./pt_path 0 # Do not collect performance data
216+bash batch_isolated_run.sh ./pt_path 1 # Collect performance data (mount msprof)
217+```
218+ 
219+## Excel Test Case Table Format
220+ 
221+Sheet1 of `excel/test_cases.xlsx` must contain the following columns:
222+ 
223+| Column Name | Type | Example | Description |
224+|---|---|---|---|
225+| Testcase_Name | str | `Prefill0` | Test case name, also used as the name of the generated .pt file |
226+| batch_size | int | `1` | |
227+| hidden_size | int | `4096` | |
228+| Seq_len | int | `8192` | |
229+| head_dim | int | `128` / `512` | |
230+| block_size | int | `128` | |
231+| cmp_ratio | int | `4` / `128` | |
232+| coff | int | `1` / `2` | |
233+| start_p | int | `0` / `8192` | |
234+| cache_mode | int | `1` / `2` | 1: continuous buffer, 2: circular buffer |
235+| layout_x | str | `TH` / `BSH` | |
236+| data_type | str | `BF16` / `FP16` | |
237+| cu_seqlens | None/str | `None` or `"[0,8192]"` | Required for the TH layout |
238+| seqused | None/str | `None` or `"[8192]"` | |
239+| start_pos | None/str | `None` or `"[0]"` | |
240+| x_datarange | str | `"[-10,10]"` | |
241+| wkv_datarange | str | `"[-10,10]"` | |
242+| wgate_datarange | str | `"[-10,10]"` | |
243+| ape_datarange | str | `"[-10,10]"` | |
244+| kv_state_datarange | str | `"[-10,10]"` | |
245+| score_state_datarange | str | `"[-10,10]"` | |
246+ 
247+**Precautions**:
248+ 
249+- When `layout_x=TH`, `cu_seqlens` is required (a list with a length of batch_size+1)
250+ 
251+## Output Files
252+ 
253+| File | Description |
254+|---|---|
255+| `result.xlsx` | Test results (accuracy, parameters, and so on) |
256+| `result_perf.xlsx` | Test results + performance data (msprof mode only) |
257+| `batch_summary.log` | Detailed log of batch isolated execution |
258+| `batch_fail_list.log` | List of failed test cases |
259+| `PROF_*/` | Directory of raw msprof performance data |
Rattention/flash_attention_score/docs/FA算子设计介绍.md→attention/flash_attention_score/docs/FA_design.md+0-0
Rattention/flash_attention_score_grad/docs/FAG算子设计介绍.md→attention/flash_attention_score_grad/docs/FAG_design.md+0-0
Rattention/fused_floyd_attention/docs/FFA算子设计介绍.md→attention/fused_floyd_attention/docs/FFA_design.md+0-0
Rattention/incre_flash_attention/docs/IFA算子设计介绍.md→attention/incre_flash_attention/docs/IFA_design.md+0-0
Aattention/incre_flash_attention/docs/IFA_design_en.md+245-0File content is under review, please try again later
Rattention/kda_input_proj/docs/KdaInputProj算子设计说明.md→attention/kda_input_proj/docs/KdaInputProj_design.md+0-0
Rattention/mla_prolog/docs/MlaProlog算子设计介绍.md→attention/mla_prolog/docs/MlaProlog_design.md+0-0
Rattention/mla_prolog_v2/docs/MlaPrologV2算子设计介绍.md→attention/mla_prolog_v2/docs/MlaPrologV2_design.md+0-0
Rattention/mla_prolog_v3/docs/MlaPrologV3算子设计介绍.md→attention/mla_prolog_v3/docs/MlaPrologV3_design.md+0-0
Aattention/pool_key_indexer/README_en.md+94-0File content is under review, please try again later
Rattention/prompt_flash_attention/docs/PFA算子设计介绍.md→attention/prompt_flash_attention/docs/PFA_design.md+0-0
Rexamples/mc2/all_gather_add/docs/AllGatherAdd算子设计实现介绍.md→examples/mc2/all_gather_add/docs/AllGatherAdd_design.md+596-596
Amc2/all_gather_matmul_v2/tests/bi/README_en.md+74-0File content is under review, please try again later
Amc2/common/op_kernel/apace/README_en.md+106-0File content is under review, please try again later
Rmc2/moe_distribute_dispatch_v2/docs/MoeDistributeDispatch-Combine算子设计介绍.md→mc2/moe_distribute_dispatch_v2/docs/MoeDistributeDispatch-Combine_design.md+0-0
Amc2/moe_update_expert/README_en.md+135-0File content is under review, please try again later
Amc2/quant_reduce_scatter/README_en.md+5-0File content is under review, please try again later
Amhc/mhc_post_backward/README_en.md+122-0
Amoe/moe_init_routing_v2/docs/moe_init_routing_v2_design_en.md+86-0File content is under review, please try again later
Aposembedding/norm_rope_concat/README_en.md+362-0File content is under review, please try again later
Aposembedding/norm_rope_concat_grad/README_en.md+354-0File content is under review, please try again later
Aposembedding/qkv_rms_norm_rope_cache_with_k_scale/README_en.md+174-0File content is under review, please try again later
Aposembedding/rotary_position_embedding_grad/README_en.md+254-0File content is under review, please try again later