Star1662
3343
Code
Code
Issues458
Pull Requests686
Actions
Discussion
Wiki
Analysis
Star1662
3343
Merged
docs: add English translations for markdown files #15306
jenny_xiangAdd on 4 days ago
docs: add English translations for markdown files #15306
Merged
jenny_xiangAdd on 4 days ago
Merged fromjenny_xiang/ops-transformerdocs/add-english-translationsintocann/ops-transformermaster
Total: 42 files changed+10942-8
ACONTRIBUTING_en.md+232-0
@@ -0,0 +1,232 @@
1+# Contribution Guide
2+ 
3+This project welcomes all developers to try it out and contribute. Before participating in community contributions, see [cann-community](https://atomgit.com/cann/community) to learn about the code of conduct, sign the CLA agreement, and familiarize yourself with the contribution process of the source code repository.
4+ 
5+## Preparation
6+ 
7+When preparing local code and submitting PRs, developers need to pay special attention to the following items:
8+ 
9+1. When submitting a PR, carefully fill in the PR template with information such as the business background, purpose, and solution of the PR.
10+2. If your changes involve new features, new interfaces, new configuration parameters, or modifications to the code flow (that is, changes beyond a simple bug fix), be sure to discuss the solution through an Issue first to avoid having your code rejected. If you are not sure whether a change counts as a "simple bug fix", it is recommended to submit an Issue for discussion.
11+ 
12+## Contribution Scenarios
13+ 
14+Developer contribution scenarios mainly include:
15+ 
16+- [Contributing a New Operator](#1-contributing-a-new-operator)
17+- [Operator Bug Fixes](#2-operator-bug-fixes)
18+- [Operator Optimization](#3-operator-optimization)
19+- [Documentation Corrections](#4-documentation-corrections)
20+- [Helping Resolve Issues Raised by Others](#5-helping-resolve-issues-raised-by-others)
21+ 
22+---
23+ 
24+## 1. Contributing a New Operator
25+ 
26+The operator development contribution process is as follows:
27+ 
28+![Operator development contribution process](./docs/en/figures/op_dev_and_contribute_flow.png "Operator development contribution flowchart")
29+ 
30+The ecosystem operators in this repository are custom-developed for specific networks and scenarios and support quick verification and use of customized algorithms. If you have a brand-new operator that you want to design and implement based on the NPU, you are welcome to present your ideas and design proposals in an Issue and contribute the operator to the corresponding operator classification directory under `experimental`. The complete contribution process is as follows:
31+ 
32+### 1. Create a Requirement Issue
33+ 
34+Create a `Requirement|需求建议` type Issue and describe the design proposal for the new operator. An Issue generally needs to include the following content:
35+ 
36+- **Background information**
37+- **Value/benefits**
38+- **Design proposal**
39+ 
40+For detailed operation steps, refer to: [Issue Operation Guide](https://atomgit.com/cann/community/blob/master/contributor/issue-operation.md).
41+ 
42+### 2. Requirement Review
43+ 
44+After creating a requirement Issue, apply for a SIG topic and attend the Ops-transformer SIG regular meeting for review at the specified time.
45+ 
46+- [Meeting topics and minutes](https://etherpad-cann.meeting.osinfra.cn/p/sig-ops-transformer)
47+- [Meeting link](https://meeting.osinfra.cn/cann/)
48+ 
49+**Handling urgent requirements:**
50+ 
51+If a requirement is urgent, you can request an ad-hoc SIG meeting:
52+ 
53+1. Fill in the topic application: [Topic application](https://etherpad-cann.meeting.osinfra.cn/p/sig-ops-transformer)
54+2. Send an email to the [maintainer](https://atomgit.com/cann/community/blob/master/CANN/sigs/ops-transformer/README.md). It is recommended to use an email subject that states the ad-hoc SIG topic application and the topic content
55+ 
56+**Requirement acceptance:**
57+ 
58+If the requirement is accepted, [SIG members](https://atomgit.com/cann/community/blob/master/CANN/sigs/ops-transformer/README.md) will assign you a suitable operator classification path (for example, `experimental/attention`). Submit the contributed operator to the corresponding operator classification directory under `experimental`.
59+ 
60+### 3. PR Submission
61+ 
62+#### Deliverables for an Ecosystem Minimal Operator
63+ 
64+```text
65+experimental/${op_class} # Operator classification (under the experimental directory)
66+├── ${op_name} # Operator name
67+│ ├── ${op_name}.cpp # Operator kernel implementation file
68+│ ├── examples # Operator example directory (guards the basic scenario; must be compilable and executable)
69+│ ├── tests # Test directory (containing customized scenario cases)
70+│ │ └── test_${op_name}.py # Operator test file
71+│ ├── CMakeLists.txt # Operator compilation configuration file
72+│ └── README.md # Operator README document
73+```
74+ 
75+#### Notes on PR Submission
76+ 
77+**Preparation before submission:**
78+ 
79+- It is recommended to submit the PR only after the requirement review is approved.
80+- Before submitting a PR, complete the development environment setup and carefully learn the project-specific development specifications and copyright statement requirements (if the code references open-source code snippets, refer to the snippet reference guide). Ensure that your contribution complies with the project standards and sign the CLA. For detailed operation steps, refer to: [PR Operation Guide](https://atomgit.com/cann/community/blob/master/contributor/pull_request_operation.md).
81+ 
82+**Deliverable requirements:**
83+ 
84+| Type | Requirement | Reference |
85+|------|------|---------|
86+| Code deliverables | Provide the operator kernel implementation and the operator test file | [fast_kernel_launch_example](examples/fast_kernel_launch_example/README_en.md) |
87+| Documentation deliverables | The operator README document is required; other documents are optional depending on the situation | [Documentation Contribution Guide](docs/CONTRIBUTING_DOCS.md) |
88+| Precision requirement | Newly contributed operators must meet the precision standard | [Open-source precision standard for ecosystem operators](https://atomgit.com/cann/opbase/blob/master/docs/zh/ops_precision_standard/mixed_tolerance_standard.md) |
89+| Coverage requirement | The operator test file must cover the core functions and main branches of the operator and meet the UT coverage requirements. Locally, you can run `bash build.sh --experimental -u --cov` to generate an lcov coverage report for self-check. On the CI side, coverage is calculated according to the configuration in [tests/test_config.yaml](tests/test_config.yaml) | [Build and test instructions](docs/zh/install/build.md) |
90+| Example guarding requirement | Provide an operator example to guard the basic scenario of the operator. The example must be compilable and executable (locally, you can run `bash build.sh --experimental --run_example ${op_name} eager` to compile and run it for verification; for running modes and optional parameters, see `bash build.sh --help`). CI will trigger example verification according to the configuration in [tests/test_config.yaml](tests/test_config.yaml) | [fast_kernel_launch_example](examples/fast_kernel_launch_example/README_en.md) |
91+| Case submission requirement | For an ecosystem operator customized for specific networks and scenarios, its customized scenario cases (function and precision verification cases under typical network shapes and typical configurations) must be committed to the operator `tests` directory together with the code to ensure sustainable regression verification of the operator. For the case generation rules, refer to the open-source precision standard for ecosystem operators | [Open-source precision standard for ecosystem operators](https://atomgit.com/cann/opbase/blob/master/docs/zh/ops_precision_standard/mixed_tolerance_standard.md) |
92+ 
93+**Compliance checks:**
94+ 
95+- [ ] Does the code comply with the [C++ Coding Standards](https://atomgit.com/cann/community/blob/master/contributor/coding-standards/C++%20Coding%20standards.md)?
96+- [ ] Does the code compile successfully?
97+- [ ] Does the Markdown document syntax comply with the specifications?
98+- [ ] Before committing code with git, you can refer to the [pre-commit tool usage guide](docs/zh/develop/pre-commit_guide.md) to make your code submission more compliant and efficient
99+ 
100+**Submission specifications:**
101+ 
102+- **Contribution directory**: Submit to the specified directory `experimental/${op_class}` as advised by SIG members. You can refer to the file placement rules of existing operators.
103+- **PR submission**: Use `git` commands to submit a PR to the target branch. Check whether the PR title is clear, whether the PR description is standard (indicating what is changed and why, and whether the corresponding Issue is linked), and whether the CLA is signed.
104+ 
105+> **Description**: If you want to directly contribute a project standard operator, see [PR Submission Requirements for Project Standard Operators](#pr-submission-requirements-for-project-standard-operators) for its deliverables and submission specifications. For the procedure of converting an ecosystem operator into a project standard operator, see [Converting an Ecosystem Operator into a Project Standard Operator](#converting-an-ecosystem-operator-into-a-project-standard-operator).
106+ 
107+### 4. CI Gate
108+ 
109+Trigger the open-source repository gate by commenting with the `compile` command, and make modifications based on the CI check results. Currently, the CI gate includes the following check items:
110+ 
111+- Code compilation
112+- Static check (if a codecheck false positive is involved, submit it to SIG members for blocking)
113+- UT tests (including coverage statistics)
114+- Smoke tests
115+ 
116+After the gate is passed, @ the Committer in the associated Issue.
117+ 
118+### 5. Committer Review
119+ 
120+The Committer will provide review comments after the review. Modify the code according to the comments, and @ the Committer after completion.
121+ 
122+### 6. Maintainer Merge
123+ 
124+After the Committer review is approved, the `/lgtm` label will be added. The Maintainer performs the final review and, after confirming that there are no issues, adds the `/approve` label to merge the PR.
125+ 
126+---
127+ 
128+## 2. Operator Bug Fixes
129+ 
130+If you find bugs in some operators in this project and want to fix them, you are welcome to create an Issue for feedback and tracking.
131+ 
132+You can follow the [Issue submission/handling guide](https://atomgit.com/cann/community#提交Issue处理Issue任务) to create a `Bug-Report|缺陷反馈` type Issue that describes the bug.
133+ 
134+---
135+ 
136+## 3. Operator Optimization
137+ 
138+If you have ideas for generalization enhancements or performance optimizations for some operator implementations in this project and want to implement these optimization points, you are welcome to contribute operator optimizations.
139+ 
140+You can follow the [Issue submission/handling guide](https://atomgit.com/cann/community#提交Issue处理Issue任务) to create a `Requirement|需求建议` type Issue that describes the optimization points and provides your design proposal.
141+ 
142+---
143+ 
144+## 4. Documentation Corrections
145+ 
146+If you find description errors in some operator documents in this project, you are welcome to create an Issue for feedback and fixes. For documentation specifications, refer to the [Documentation Contribution Guide](docs/CONTRIBUTING_DOCS.md).
147+ 
148+You can follow the [Issue submission/handling guide](https://atomgit.com/cann/community#提交Issue处理Issue任务) to create a `Documentation|文档反馈` type Issue that points out the problems in the corresponding document.
149+ 
150+---
151+ 
152+## 5. Helping Resolve Issues Raised by Others
153+ 
154+If you have a suitable solution to a problem encountered by others in the community, you are welcome to post comments in the Issue to share ideas, help others solve problems and pain points, and jointly improve usability.
155+ 
156+If the corresponding Issue requires code changes, you can enter `/assign` or `/assign @yourself` in the Issue comment box to assign the Issue to yourself, and then track and help resolve the problem.
157+ 
158+---
159+ 
160+## Appendix
161+ 
162+### Introduction to Project Standard Operators
163+ 
164+Operators in this repository fall into two categories: the ecosystem operators described in the main text, and project standard operators (that is, non-experimental operators). The two categories are compared as follows:
165+ 
166+| Comparison dimension | Project standard operator | Ecosystem operator |
167+|---------|--------------|---------------------|
168+| Target scenario | Mainstream model scenarios, pursuing high performance, high precision, and high generalization | Custom development for specific networks and scenarios |
169+| Usage | Released with official versions and supported for use as project standard operators | Supports quick verification and use of customized algorithms |
170+| Contribution directory | `${op_class}/${op_name}` | `experimental/${op_class}/${op_name}` |
171+| Deliverable requirements | Complete deliverables (op_host, op_kernel, UT, README, etc.) | Minimal deliverables (kernel implementation, example, test files, and README), plus the coverage, example guarding, and customized scenario case submission requirements |
172+| Review focus | Comprehensive review of performance, precision, and generalization | Functional correctness, customized scenario adaptation, and completeness of cases and coverage |
173+ 
174+### Deliverables for a Project Standard Operator
175+ 
176+```text
177+${op_class} # Operator classification
178+├── ${op_name} # Operator name
179+│ ├── op_host # Operator definition and Tiling-related implementation
180+│ │ ├── ${op_name}_def.cpp # Operator definition file
181+│ │ ├── ${op_name}_tiling.cpp # Operator Tiling implementation file
182+│ │ ├── ${op_name}_tiling_${sub_case}.cpp # Optional, sub-scenario Tiling implementation. ${sub_case} indicates the sub-scenario (such as arch35)
183+│ │ ├── ${op_name}_tiling_${sub_case}.h # Optional, header file of the sub-scenario Tiling implementation
184+│ │ └── CMakeLists.txt
185+│ ├── op_kernel # Operator kernel directory
186+│ │ ├── ${op_name}.cpp # Kernel entry file, which contains the main function and scheduling logic
187+│ │ ├── ${op_name}.h # Kernel implementation file, which defines the kernel header file and contains function descriptions, structure definitions, and logic implementation
188+│ │ ├── ${op_name}_tiling_data.h # TilingData file, which stores configuration information related to the Tiling strategy
189+│ │ └── ${op_name}_tiling_key.h # TilingKey file, which defines the key of the Tiling strategy to identify different partitioning methods
190+│ ├── CMakeLists.txt # Operator compilation configuration file. Keep the original file.
191+│ ├── README.md # Operator description document
192+│ └── tests # Operator test files
193+│ └── ut # Operator UT test files
194+```
195+ 
196+For the complete operator directory structure (including all optional deliverables), see [Project Directory Structure](docs/zh/install/dir_structure.md#项目目录).
197+ 
198+> **Description**: For Tiling implementation files in the op_host directory that participate in compilation, the file name must contain the `_tiling` identifier (such as `${op_name}_tiling.cpp` or `${op_name}_tiling_${sub_case}.cpp`). Otherwise, the compilation system will not recognize them. When splitting the Tiling implementation for sub-scenarios (such as a specific architecture like arch35), follow this naming rule.
199+ 
200+### PR Submission Requirements for Project Standard Operators
201+ 
202+**Deliverable requirements:**
203+ 
204+| Type | Requirement | Reference |
205+|------|------|---------|
206+| Code deliverables | Provide the operator Tiling implementation in op_host, the operator kernel implementation in op_kernel, and the operator UT test files | [Operator Development Guide](docs/zh/develop/aicore_develop_guide.md) |
207+| Documentation deliverables | The operator README document is required; other documents are optional depending on the situation | [Documentation Contribution Guide](docs/CONTRIBUTING_DOCS.md) |
208+ 
209+**Compliance checks:**
210+ 
211+- [ ] Does the code comply with the [C++ Coding Standards](https://atomgit.com/cann/community/blob/master/contributor/coding-standards/C++%20Coding%20standards.md)?
212+- [ ] Does the code comply with the basic programming specifications for standard operators?
213+- [ ] Does the code compile successfully?
214+- [ ] Does the Markdown document syntax comply with the specifications?
215+ 
216+**Submission specifications:**
217+ 
218+- **Contribution directory**: Submit to the specified operator classification directory (such as `attention/${op_name}`) as advised by SIG members. You can refer to the file placement rules of existing operators.
219+- **PR submission**: Use `git` commands to submit a PR to the target branch. Check whether the PR title is clear, whether the PR description is standard (indicating what is changed and why, and whether the corresponding Issue is linked), and whether the CLA is signed.
220+ 
221+### Converting an Ecosystem Operator into a Project Standard Operator
222+ 
223+Ecosystem operators are custom-developed for specific networks and scenarios. After an operator becomes stable and mature, you can apply to convert it into a project standard operator. The conversion must meet the following conditions:
224+ 
225+- **Project standard operator requirements**: The operator has been stably verified in the target network and scenarios, and its performance and precision meet the usage requirements of project standard operators.
226+- **Generalization requirement**: The operator has good generalization capability and can support the common usage requirements of mainstream model scenarios.
227+ 
228+The conversion procedure is as follows:
229+ 
230+1. Create a `Requirement|需求建议` type Issue, describe the actual scenario verification status and generalization performance of the operator, and apply to migrate the operator from `experimental/${op_class}/${op_name}` to the corresponding operator module directory `${op_class}/${op_name}`.
231+2. Apply for a SIG topic and attend the review: fill in the conversion topic on the [topic application page](https://etherpad-cann.meeting.osinfra.cn/p/sig-ops-transformer) to complete the application, and attend the Ops-transformer SIG regular meeting review as scheduled by the topic ([meeting link](https://meeting.osinfra.cn/cann/)). After the application is accepted by the SIG review, complete the deliverables according to the "Deliverables for a Project Standard Operator" standard (such as the op_host Tiling implementation).
232+3. The conversion is completed after the CI gate and Committer review are passed and the Maintainer merges the changes. The converted operator will be released with official versions later and supported for use as a project standard operator.
AREADME_en.md+92-0
@@ -0,0 +1,92 @@
1+<div align="center">
2+ 
3+# ops-transformer
4+ 
5+English | [简体中文](./README.md)
6+ 
7+[![CANNcommunity](https://img.shields.io/badge/docs-Community_Operator_List-brightgreen.svg?style=flat)](https://www.hiascend.com/document/redirect/CannCommunityOplist)
8+[![repo](https://img.shields.io/badge/docs-Project_Documentation-green.svg?style=flat)](docs)
9+[![license](https://img.shields.io/badge/license-CANN_Open_2.0-blue.svg)](LICENSE)
10+[![contributing](https://img.shields.io/badge/Contribution_Guide-teal)](CONTRIBUTING_en.md)
11+[![SIG](https://img.shields.io/badge/SIG-ops_transformer-yellow)](https://atomgit.com/cann/community/blob/master/CANN/sigs/ops-transformer/README.md)
12+[![Zread](https://img.shields.io/badge/Zread-Ask_AI-_.svg?style=flat&color=0052D9&labelColor=000000&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZD0iTTQuOTYxNTYgMS42MDAxSDIuMjQxNTZDMS44ODgxIDEuNjAwMSAxLjYwMTU2IDEuODg2NjQgMS42MDE1NiAyLjI0MDFWNC45NjAxQzEuNjAxNTYgNS4zMTM1NiAxLjg4ODEgNS42MDAxIDIuMjQxNTYgNS42MDAxSDQuOTYxNTZDNS4zMTUwMiA1LjYwMDEgNS42MDE1NiA1LjMxMzU2IDUuNjAxNTYgNC45NjAxVjIuMjQwMUM1LjYwMTU2IDEuODg2NjQgNS4zMTUwMiAxLjYwMDEgNC45NjE1NiAxLjYwMDFaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00Ljk2MTU2IDEwLjM5OTlIMi4yNDE1NkMxLjg4ODEgMTAuMzk5OSAxLjYwMTU2IDEwLjY4NjQgMS42MDE1NiAxMS4wMzk5VjEzLjc1OTlDMS42MDE1NiAxNC4xMTM0IDEuODg4MSAxNC4zOTk5IDIuMjQxNTYgMTQuMzk5OUg0Ljk2MTU2QzUuMzE1MDIgMTQuMzk5OSA1LjYwMTU2IDE0LjExMzQgNS42MDE1NiAxMy43NTk5VjExLjAzOTlDNS42MDE1NiAxMC42ODY0IDUuMzE1MDIgMTAuMzk5OSA0Ljk2MTU2IDEwLjM5OTlaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik0xMy43NTg0IDEuNjAwMUgxMS4wMzg0QzEwLjY4NSAxLjYwMDEgMTAuMzk4NCAxLjg4NjY0IDEwLjM5ODQgMi4yNDAxVjQuOTYwMUMxMC4zOTg0IDUuMzEzNTYgMTAuNjg1IDUuNjAwMSAxMS4wMzg0IDUuNjAwMUgxMy43NTg0QzE0LjExMTkgNS42MDAxIDE0LjM5ODQgNS4zMTM1NiAxNC4zOTg0IDQuOTYwMVYyLjI0MDFDMTQuMzk4NCAxLjg4NjY0IDE0LjExMTkgMS42MDAxIDEzLjc1ODQgMS42MDAxWiIgZmlsbD0iI2ZmZiIvPgo8cGF0aCBkPSJNNCAxMkwxMiA0TDQgMTJaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00IDEyTDEyIDQiIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSIxLjUiIHN0cm9rZS1saW5lY2FwPSJyb3VuZCIvPgo8L3N2Zz4K&logoColor=ffffff)](https://zread.ai/hicann/ops-transformer)
13+ 
14+</div>
15+ 
16+---
17+ 
18+## 🔥Latest News
19+ 
20+- [2026/08] The repository integrates the [pre-commit code check framework](docs/zh/develop/pre-commit_guide.md). After configuration, `git commit` automatically checks and fixes format issues, intercepting specification issues at the commit stage and improving code merge efficiency. For details, see the [pre-commit configuration guide](docs/zh/develop/pre-commit_guide.md).
21+- [2026/07] The [quant_flash_attn](attention/quant_flash_attn) operator is newly supported on A5, dedicated to the fully quantized Dense Attention scenario. The experimental directory adds an [operator with the same name](experimental/attention/quant_flash_attn) that supports MxFP4 full quantization on A5. The [chunk_gated_delta_rule](attention/chunk_gated_delta_rule) operator supports A5. A5 adds the DSV4 dense training-stage TopK filtering backward operators [dense_lightning_indexer_grad_kl_loss](attention/dense_lightning_indexer_grad_kl_loss) and [dense_lightning_indexer_softmax_lse](attention/dense_lightning_indexer_softmax_lse). mc2 newly supports the [engram_fetch_grad](mc2/engram_fetch_grad) operator.
22+- [2026/06] [sparse_flash_mla](attention/sparse_flash_mla), [mixed_quant_sparse_flash_mla](attention/mixed_quant_sparse_flash_mla), and the corresponding metadata operators are newly supported on A5 for sparse Attention computation in the DSV4 scenario. [lightning_indexer_v2](attention/lightning_indexer_v2), [quant_lightning_indexer_v2](attention/quant_lightning_indexer_v2), and the corresponding metadata operators are added for TopK filtering in the DSV4 scenario. A5 adds the DSV4 sparse training-stage Attention backward operators [sparse_flash_mla_grad](attention/sparse_flash_mla_grad) and [sparse_lightning_indexer_kl_loss_grad](attention/sparse_lightning_indexer_kl_loss_grad). mc2 newly supports the [engram_fetch](mc2/engram_fetch) and [engram_fetch_wait](mc2/engram_fetch_wait) operators.
23+- [2026/05] The [fused_infer_attention_score](attention/fused_infer_attention_score) operator supports MxFP8 full quantization on A5. The [flash_attn](attention/flash_attn) operator is newly supported on A5, dedicated to the non-quantized Dense Attention scenario.
24+- [2026/04] mc2 newly supports the [mega_moe](mc2/mega_moe) operator.
25+- [2026/03] The [recurrent_gated_delta_rule](attention/recurrent_gated_delta_rule) operator supports A5.
26+- [2026/02] Newly supported operators: [mhc_post](experimental/mhc/mhc_post), [mhc_pre](experimental/mhc/mhc_pre), and [mhc_res](experimental/mhc/mhc_res).
27+- [2026/01] Added the [grouped_matmul<<<>>> invocation example](examples/fast_kernel_launch_example/csrc/grouped_matmul) for convenient customized use.
28+- [2026/01] Newly supported operators: [fused_floyd_attention](attention/fused_floyd_attention), [fused_floyd_attention_grad](attention/fused_floyd_attention_grad), and [matmul_allto_all](mc2/matmul_allto_all).
29+- [2025/12] Added [QuickStart](docs/QUICKSTART.md), which guides beginners with zero experience through operator project deployment (Docker environment supported), operator development, and the contribution process.
30+- [2025/12] Optimized the guide documents, focusing on the [Operator Development Guide](docs/zh/develop/aicore_develop_guide.md), clarifying the minimal deliverables and key sample code, and providing guidance on migrating operators from the [Ascend/samples](https://gitee.com/ascend/samples/tree/master) repository to this project.
31+- [2025/12] Supports transformer-type ONNX operator plugins, including [NPUMultiHeadAttention](common/src/framework), [NPUMoeComputeExpertTokens](moe/moe_compute_expert_tokens/framework), and so on.
32+- [2025/12] Newly supported operators: [kv_rms_norm_rope_cache](posembedding/kv_rms_norm_rope_cache), [attention_update](attention/attention_update), [attention_worker_scheduler](attention/attention_worker_scheduler), [gather_pa_kv_cache](attention/gather_pa_kv_cache), [kv_quant_sparse_flash_attention](attention/kv_quant_sparse_flash_attention), [lightning_indexer_grad](attention/lightning_indexer_grad), [mla_preprocess](attention/mla_preprocess), [mla_preprocess_v2](attention/mla_preprocess_v2), [grouped_matmul_swiglu_quant_v2](gmm/grouped_matmul_swiglu_quant_v2), [attention_to_ffn](mc2/attention_to_ffn), and [ffn_to_attention](mc2/ffn_to_attention).
33+- [2025/12] Open-source operators support development and debugging with the [NPU Simulator](docs/zh/debug/npu_sim.md) simulation tool.
34+- [2025/12] Open-source operators support the [Ascend 950PR&950DT products](docs/zh/ascend950_op_list.md)/KirinX90 products.
35+- [2025/11] Newly supported operators: [kv_quant_sparse_flash_attention](attention/kv_quant_sparse_flash_attention), [lightning_indexer](attention/lightning_indexer), [quant_lightning_indexer](attention/quant_lightning_indexer), and [sparse_flash_attention](attention/sparse_flash_attention).
36+- [2025/11] Added the sample operators [rope_matrix](experimental/posembedding/rope_matrix) and [all_gather_add](examples/mc2/all_gather_add).
37+- [2025/11] Added the operator development project template [NpuOpsTransformerExt](experimental/npu_ops_transformer_ext), which seamlessly integrates PyTorch tensor operations and supports automatic differentiation and a unified GPU/NPU interface.
38+- [2025/10] Added the [experimental](experimental) directory and improved the [Contribution Guide](CONTRIBUTING_en.md), allowing developers to debug and contribute custom operators.
39+- [2025/09] The ops-transformer project goes online for the first time, with open-source operators supporting Atlas A2/A3 products.
40+ 
41+## 🚀Overview
42+ 
43+ops-transformer is an advanced operator library in the [CANN](https://hiascend.com/software/cann) (Compute Architecture for Neural Networks) operator library that provides computation for transformer-type large models. It includes the attention, moe, and mc2 operator categories, covering various scenarios such as attention computation, MoE computation, and communication-computation fusion. The position of the operator library in the architecture diagram is as follows.
44+ 
45+<img src="docs/en/figures/architecture.png" alt="Architecture diagram" width="700px" height="320px">
46+ 
47+## 📌Version Compatibility
48+ 
49+The source code of this project is released along with CANN software versions. For the mapping between CANN software versions and tags of this project, refer to the corresponding version descriptions in the [release repository](https://atomgit.com/cann/release-management).
50+Note that to ensure smooth customized development of your source code, select a compatible CANN version and the source code of the corresponding Gitcode tag. Using the master branch may cause version mismatch risks.
51+ 
52+## 🛠️Environment Preparation
53+ 
54+[Environment deployment](docs/zh/install/quick_install.md) is the prerequisite for experiencing the capabilities of this project. Complete the NPU driver and CANN package installation first to ensure that the environment works properly.
55+ 
56+## ⬇️Source Code Download
57+ 
58+After the environment is ready, download the branch source code that matches the CANN version. The general command is as follows, where \$\{tag\_version\} is replaced with the branch tag name. The following uses downloading the 9.0.0 branch source code as an example:
59+ 
60+```bash
61+# General command: git clone -b ${tag_version} https://atomgit.com/cann/ops-transformer.git
62+git clone -b 9.0.0 https://atomgit.com/cann/ops-transformer.git
63+```
64+ 
65+> Description: If the compatible branch source code already exists in the environment, **you can skip this step**. For example, CANNLab provides the source code corresponding to the latest CANN version by default.
66+ 
67+## 📖Learning Tutorials
68+ 
69+- [Quick Start](docs/QUICKSTART.md): Quickly experience the core basic capabilities of the project from scratch, covering source code compilation, operator invocation, development, debugging, and other operations.
70+- [Advanced Tutorials](docs/README.md): If you need to learn more about project compilation and deployment, operator invocation, development, debugging and tuning, and other capabilities, see the documentation center for detailed guidance.
71+ 
72+## 💬Related Information
73+ 
74+- [Directory Structure](docs/zh/install/dir_structure.md)
75+- [Operator Index](docs/zh/op_list.md)
76+- [aclnn API Index](docs/en/op_api_list.md)
77+- [torch API Index](docs/en/torch_api_list.md)
78+- [Contribution Guide](CONTRIBUTING_en.md)
79+- [Security Statement](SECURITY_en.md)
80+- [License](LICENSE)
81+- [SIG](https://atomgit.com/cann/community/tree/master/CANN/sigs/ops-transformer)
82+- [Contact Us] WeChat: Clark_Dan
83+ 
84+-----
85+PS: The features and documentation of this project are continuously being updated and improved. Stay tuned for the latest version.
86+ 
87+- **Issue feedback**: Submit issues through GitCode [Issues](https://atomgit.com/cann/ops-transformer/issues).
88+- **Community interaction**: Join discussions through GitCode [Discussions](https://atomgit.com/cann/ops-transformer/discussions).
89+- **Technical column**: Get technical articles (such as serialized tutorials and best practices) through GitCode [Wiki](https://atomgit.com/cann/ops-transformer/wiki).
90+ 
91+ 
92+​
ASECURITY_en.md+64-0
@@ -0,0 +1,64 @@
1+# Security Statement
2+ 
3+## Recommendations for the Running User
4+ 
5+For security reasons, do not use administrator accounts such as root to execute any command. Follow the principle of least privilege.
6+ 
7+## File Permission Control
8+ 
9+- It is recommended that you set the umask value of the running system to 0027 or higher on the host (including the host machine) and in containers. This ensures that the default maximum permission is 750 for new folders and 640 for new files.
10+- It is recommended that you take security measures such as permission control for sensitive content, including personal privacy data, commercial assets, source files, and various files saved during operator development. For example, this involves permission control over the installation directory of this project and over input public data files. For recommended permission settings, refer to [A-Recommended Maximum File (Folder) Permissions for Different Scenarios](#a-recommended-maximum-file-folder-permissions-for-different-scenarios).
11+- During operator running, operator compilation files may be cached and stored in the `kernel_meta_*` folder in the running directory to speed up subsequent operator invocations. You can perform permission control on the generated files as needed.
12+- You need to perform proper permission control during installation and use. For the settings, it is recommended that you refer to the file permission reference in [A-Recommended Maximum File (Folder) Permissions for Different Scenarios](#a-recommended-maximum-file-folder-permissions-for-different-scenarios).
13+ 
14+## Build Security Statement
15+ 
16+When compiling and installing this project from source code, you need to compile it yourself. The compilation process generates some intermediate files. After the compilation is complete, it is recommended that you perform proper permission control on the intermediate files to ensure file security.
17+ 
18+## Running Security Statement
19+ 
20+- It is recommended that you write operator invocation scripts based on the resource status of the running environment. If an operator invocation script does not match the resource status, for example, the space used for generating input data or benchmark computation results exceeds the memory capacity limit, or the data saved locally by the script exceeds the disk space, errors may occur and the process may exit unexpectedly.
21+- When an operator runs abnormally, it exits the process and prints error information. It is recommended that you locate the specific error cause based on the error message, using methods such as setting the operator to run synchronously and viewing log files.
22+- When operators are invoked through [PyTorch](https://gitee.com/ascend/pytorch), version mismatches may cause running errors. For details, refer to the [PyTorch security statement](https://gitee.com/ascend/pytorch#%E5%AE%89%E5%85%A8%E5%A3%B0%E6%98%8E).
23+ 
24+## Public Network Address Statement
25+ 
26+The public network addresses contained in the code of this project are declared as follows:
27+ 
28+| Type | Open-source code address | File name | Public IP address/Public URL/Domain name/Email address/Archive file address | Purpose description |
29+| :------------: |:------------------------------------------------------------------------------------------:|:----------------------------------------------------------| :---------------------------------------------------------- |:-----------------------------------------|
30+| Dependency | N/A | cmake/third_party/makeself-fetch.cmake | https://cann-3rd.obs.cn-north-4.myhuaweicloud.com/makeself/makeself-release-2.5.0.tar.gz | Downloads the makeself source code from obs as a compilation dependency |
31+| Dependency | N/A | cmake/third_party/protobuf.cmake | https://cann-3rd.obs.cn-north-4.myhuaweicloud.com/protobuf/protobuf-25.1.tar.gz | Downloads the protobuf source code from obs as a compilation dependency |
32+| Dependency | N/A | cmake/third_party/json.cmake | https://cann-3rd.obs.cn-north-4.myhuaweicloud.com/json/json-3.12.0.tar.gz | Downloads the JSON source code from obs as a compilation dependency |
33+| Dependency | N/A | cmake/third_party/gtest.cmake | https://cann-3rd.obs.cn-north-4.myhuaweicloud.com/googletest/googletest-1.14.0.tar.gz | Downloads the googletest source code from obs as a compilation dependency |
34+| Dependency | N/A | cmake/third_party/secure_c.cmake | https://gitee.com/openeuler/libboundscheck/repository/archive/v1.1.10.tar.gz | Downloads the libboundscheck source code from gitee as a compilation dependency |
35+ 
36+---
37+ 
38+## Vulnerability Mechanism Description
39+ 
40+[Vulnerability management](https://atomgit.com/cann/community/blob/master/security/security.md)
41+ 
42+## Appendix
43+ 
44+### A-Recommended Maximum File (Folder) Permissions for Different Scenarios
45+ 
46+| Type | Reference maximum Linux permission |
47+| -------------- | --------------- |
48+| User home directory | 750 (rwxr-x---) |
49+| Program files (including script files, library files, and so on) | 550 (r-xr-x---) |
50+| Program file directory | 550 (r-xr-x---) |
51+| Configuration file | 640 (rw-r-----) |
52+| Configuration file directory | 750 (rwxr-x---) |
53+| Log file (recording completed or archived) | 440 (r--r-----) |
54+| Log file (being recorded) | 640 (rw-r-----) |
55+| Log file directory | 750 (rwxr-x---) |
56+| Debug file | 640 (rw-r-----) |
57+| Debug file directory | 750 (rwxr-x---) |
58+| Temporary file directory | 750 (rwxr-x---) |
59+| Maintenance and upgrade file directory | 770 (rwxrwx---) |
60+| Service data file | 640 (rw-r-----) |
61+| Service data file directory | 750 (rwxr-x---) |
62+| Key component, private key, certificate, and ciphertext file directory | 700 (rwx------) |
63+| Key components, private keys, certificates, and encrypted ciphertext | 600 (rw-------) |
64+| Encryption/decryption interfaces and encryption/decryption scripts | 500 (r-x------) |
Aattention/chunk_kda_fwd/docs/ChunkKdaFwd_design_en.md+300-0
@@ -0,0 +1,300 @@
1+# ChunkKdaFwd Operator Design Introduction
2+ 
3+## 1 Design Objectives
4+ 
5+1. The top-level interface aligns with the non-CP FLA `chunk_kda_fwd`.
6+2. No public operator prototype is added; the A5 specialization reuses the existing
7+ `ChunkKdaFwd` L0 calling interface and the `chunk_kda_fwd` device kernel launch entry.
8+3. A2/A3/A5 use the same mathematical definition; A5 retains the regbase dual-issue specialization.
9+4. The input layout and the output layout are decoupled.
10+5. FwdH is maintained separately by KDA under arch22/arch35, and supports key-wise gate and `state_v_first`.
11+ 
12+## 2 Computation Process
13+ 
14+### 2.1 Interface Layering
15+ 
16+```text
17+raw g -> ChunkKdaFwd[
18+ gate cumsum -> Prepare/Post-WU -> FwdH -> Finalize
19+] -> attn_out
20+```
21+ 
22+The interface call chain is divided by layer:
23+ 
24+```text
25+cann_ops_transformer.ops.chunk_kda_fwd
26+ -> PyTorch custom op registration and wrapper
27+ -> aclnnChunkKdaFwd
28+ -> ChunkKdaFwd L0 registration and launcher
29+ -> chunk_kda_fwd device kernel launch entry
30+ -> Gate / Prepare / Post-WU / FwdH / Finalize internal stages
31+```
32+ 
33+`aclnnChunkKdaFwd` is responsible for contiguizing the public layouts, performing necessary view
34+transformations, and calling the registered `ChunkKdaFwd` L0 interface once. The L0 launcher continues to
35+use the `chunk_kda_fwd` device kernel launch entry. Gate, Prepare, Post-WU, FwdH, and Finalize are only
36+internal computation stages of this kernel; no `ChunkKdaFwdPrepare`, `ChunkKdaFwdPostWu`, or
37+`ChunkKdaFwdFinalize` L0 operators are registered. All information required for template selection is
38+derived from existing inputs, shapes, and attributes; no public attributes or interface fields are added.
39+ 
40+### 2.2 Computation Stages
41+ 
42+#### 2.2.1 Gate Cumsum
43+ 
44+Converts the raw/activated gate into FP32 chunk-local log2 cumulative values:
45+ 
46+```text
47+gamma = g use_gate_in_kernel=false
48+gamma = -exp(A_log) * softplus(g + dt_bias) use_gate_in_kernel=true, safe_gate=false
49+gamma = lower_bound * sigmoid(exp(A_log) * (g + dt_bias)) use_gate_in_kernel=true, safe_gate=true
50+gk = chunk_local_cumsum(gamma) / ln(2)
51+```
52+ 
53+#### 2.2.2 Prepare
54+ 
55+Reads only `q/k/v/gk/beta` and variable-length metadata, and produces:
56+ 
57+```text
58+Aqk, Akk, qg, w_seed, u_seed
59+```
60+ 
61+```text
62+qg[i] = q[i] * exp2(gk[i])
63+w_seed[i] = beta[i] * k[i] * exp2(gk[i])
64+u_seed[i] = beta[i] * v[i]
65+```
66+ 
67+The score dot product takes BF16/FP16 as input and accumulates in FP32. The `L` of `Akk=(I+L)^{-1}`, the
68+identity matrix, diagonal block inversion, block merging, Cube inputs/accumulation, and workspace are all
69+FP32, with HF32 explicitly disabled. The public `Akk` is written back as the q dtype only after the
70+complete FP32 solve finishes.
71+ 
72+#### 2.2.3 Post-WU
73+ 
74+Reads only `k/gk/w_seed/Akk/u_seed`, and produces:
75+ 
76+```text
77+w, u, kg
78+w = Akk @ w_seed
79+u = Akk @ u_seed
80+kg[i] = k[i] * exp2(gk[last] - gk[i])
81+```
82+ 
83+When `H_v>H_k`, the scenario is GVA (Grouped Value Attention): each group of `H_v/H_k` consecutive Value
84+heads shares one q/k head. When head numbering starts from 0, the q/k head corresponding to Value head
85+`h_v` is numbered `floor(h_v/(H_v/H_k))`. The head loop of `Akk` runs over `H_v`, and the GVA mapping is
86+converted only when q/k heads are read, avoiding duplication or omission based on `H_k`.
87+Whether Post-WU is executed as a separate internal stage is decided by tiling. When the A5 dense-aligned
88+fusion conditions are met, Prepare directly calls the Post-WU computation components; in other scenarios,
89+the internal Post-WU stage is executed after Prepare finishes and completes synchronization. Both modes
90+reuse the same `ChunkKdaFwd` L0 calling interface and device kernel launch entry.
91+ 
92+#### 2.2.4 FwdH
93+ 
94+Reads `kg/w/u/gk` and the optional `initial_state`, and computes the cross-chunk state update:
95+ 
96+```text
97+v_new = u - w @ h_prev
98+h_next = exp2(gk_last) * h_prev + kg^T @ v_new
99+h[chunk] = h_prev
100+final_state = h_after_last_chunk
101+```
102+ 
103+Both arch22 and arch35 use the KDA private FwdH implementation inside the `ChunkKdaFwd` directory and do
104+not depend on kernel source code of other operators. The cross-chunk `h_prev/h_next` is always stored in
105+the FP32 workspace or the FP32 `final_state`; the public `h` is written back as a snapshot for backward
106+only, in the q dtype, and cannot participate in the state update of the next chunk.
107+ 
108+#### 2.2.5 Finalize
109+ 
110+Reads only `qg/Aqk/v_new/h` and computes with FP32 accumulation:
111+ 
112+```text
113+attn_out = scale * (qg @ h) + Aqk @ v_new
114+```
115+ 
116+`scale` is applied after the first term completes the FP32 dot product; it is not pre-quantized into the
117+BF16 `qg`.
118+ 
119+`attn_out` is written out directly in BSND/TND inside the kernel. Intermediate quantities for backward
120+remain in BNSD/NTD.
121+ 
122+## 3 Numerical Design
123+ 
124+- The `q`, `k`, `v`, and intermediate quantities stored in the input dtype in the formula serve as matmul
125+ inputs, and matmuls use FP32 accumulation.
126+- The $L$ of `Akk=(I+L)^{-1}`, the identity matrix, diagonal block inversion, block merging, Cube inputs,
127+ and accumulation remain FP32, and HF32 is not enabled.
128+- The computation of the cross-chunk state $h_c$ and `final_state` remains FP32; the low-precision snapshot
129+ of the public `h` is not used to continue state propagation to the next chunk.
130+- The two matmuls and the addition of `attn_out` use FP32 accumulation, and the result is converted to the
131+ output dtype during the final write-back.
132+ 
133+## 4 Output Layout and Recomputation
134+ 
135+### 4.1 State Layout
136+ 
137+Internal states uniformly use `[...,K,V]`. When `state_v_first=true`, L2 transposes the initial state
138+before entering FwdH. The internal `hCompute` always remains head-major for Finalize to consume; the
139+public `hOut` is converted to sequence-major at the L2 export boundary, with the order of the last two
140+dimensions decided by `state_v_first`. `final_state` is arranged by sequence, consistent with the FLA
141+top-level output.
142+ 
143+### 4.2 Recomputation Policy
144+ 
145+L2 does not understand autograd recomputation policies. `final_state/gk/w/u/qg/kg/v_new/h` are mutually
146+independent `OPTIONAL_OUTPUT` items; a non-null pointer indicates export, and a null pointer indicates that
147+the result is not exposed. The arch35 fast path passes fixed ABI placeholders for hidden outputs, and
148+tiling receives the actual intermediate results in the kernel workspace; the generic path uses the same
149+rule. When a public output exists, it is used directly as the internal target.
150+ 
151+The PyTorch API wrapper layer aligns with the fla-org `chunk_kda_fwd` commit
152+`0f0f0c97af39343855b43bbbaddcedfda5cb9d77`:
153+ 
154+- `attn_out` is always returned.
155+- `final_state` is returned only when `output_final_state=true`.
156+- `gk` is returned when `use_gate_in_kernel=false` or `disable_recompute=true`.
157+- `Aqk` is always returned.
158+- `Akk` is always returned.
159+- `w` is returned when `disable_recompute=true`.
160+- `u` is returned when `disable_recompute=true`.
161+- `qg` is returned when `disable_recompute=true`.
162+- `kg` is returned when `disable_recompute=true`.
163+- `v_new` is returned when `disable_recompute=true`.
164+- `h` is returned when `disable_recompute=true` or `return_intermediate_states=true`.
165+- `initial_state` is passed through as the original object by the Python layer; it is not an aclnn output.
166+ 
167+The internal `hCompute` and the public `hOut` have two lifecycles: `hCompute` is the required head-major
168+stage result from FwdH to Finalize, and the kernel workspace holds it when `hOut` is empty; when `hOut` is
169+not empty, L2 provides a head-major temporary output and converts it to sequence-major at the export
170+boundary. This rule aligns with the non-CP low-level 12-return-value interface; the 12th item
171+`initial_state` is passed through as the original object by the Python layer.
172+ 
173+## 5 Tiling Design
174+ 
175+A "template" in this document refers to a computation path selected by tiling or kernel conditions that
176+has an independent data split or pipeline organization. Ordinary C++ `template` helpers used only for
177+dtype, constant propagation, or code reuse are not listed separately as scenario templates.
178+ 
179+### 5.1 Interface and Implementation Directory
180+ 
181+`ChunkKdaFwd` retains one set of L0 registration, launcher, and the `op_kernel/chunk_kda_fwd.cpp` device
182+kernel launch entry. The A2/A3 generic implementations are located in `op_kernel/*.h`, the A5
183+specializations are located in `op_kernel/arch35/*.h`, and the A5 host derived options are located in
184+`op_host/arch35/chunk_kda_fwd_tiling_impl.h`. The `.h` files of Prepare, Post-WU, and Finalize are internal
185+computation components of the same kernel, not independent L0 operator interfaces.
186+ 
187+### 5.2 Compile-Time Shape Templates
188+ 
189+| tiling key | Compile-Time Constants | A2/A3 Implementation | A5 Implementation | Covered Scenarios |
190+| --- | --- | --- | --- | --- |
191+| `key=1` generic shape template | `COMPILE_BT/K/V=0/0/0` | Generic implementation in the root directory | Generic implementation in the root directory | Scenarios that do not satisfy `chunk=64,K=128,V=128`, including dense, tail, varlen, and GVA |
192+| `key=2` chunk64/K128/V128 template | `COMPILE_BT/K/V=64/128/128` | Generic implementation in the root directory, using compile-time constants | `arch35` specialized implementation | dense, tail, varlen, and different head counts with `chunk=64,K=V=128` |
193+ 
194+`SetTilingKey` checks only chunk, K, and V, and does not check the SoC, layout, whether the input is
195+varlen, or whether a tail block exists. Therefore, the tiling key represents a family of shape templates,
196+not a platform, interface, or operator version. Platform differences are selected by the architecture
197+branches compiled under the same key.
198+ 
199+### 5.3 A5 key2 Execution Combination Templates
200+ 
201+`ConfigureChunkKdaFwdArch35` derives four internal options from the instantiated states of existing inputs
202+and outputs. The options are not public attributes and do not change the L0, aclnn L2, or Python
203+interfaces.
204+ 
205+| Execution Combination | Selection Condition | Internal Stage Order |
206+| --- | --- | --- |
207+| standalone Gate | `computeGateInPrepare=false` | Gate executes fully and synchronizes, then Prepare starts |
208+| Prepare inline Gate | BF16 q/k/v, raw g in FP32, `A_log` present, `use_gate_in_kernel=true`, and `safe_gate=true` | The Prepare AIV computes chunk-local gk in the gate-product pipeline |
209+| A5 dense FwdH | BF16 and non-varlen, and `T % 64 == 0` | Uses the `ChunkKdaFwdFwdH` arch35 dense template |
210+| Intra-stage pipeline fusion of Prepare/Post-WU | A5 key2, BF16, `safe_gate=true`, dense-aligned, and `H_v % 2 == 0`, and the next-row template is not selected | After the Prepare AIC produces data by head pair, it directly calls the Post-WU components to consume data in batches |
211+| Intra-stage pipeline fusion of Post-WU/FwdH | The conditions of the previous row hold, Gate has already been computed inside Prepare, and `qg/v_new/h` are not exported | The dense FwdH directly completes the computation required for W/U and state propagation |
212+| Independent internal Post-WU stage | Neither of the two Post-WU fusions is selected | `RunChunkKdaPostWu` is called after Prepare finishes synchronization |
213+| Generic FwdH backend | FP16, varlen, non-aligned tail blocks, or conditions of dense FwdH not satisfied | Selects the KDA private `KdaFwdHTileShapes128/256` templates |
214+ 
215+Therefore, "Prepare/Post-WU fusion" describes only the execution combinations that satisfy the conditions
216+in the table; it does not mean that all key2 scenarios remove the internal Post-WU stage. In particular,
217+varlen and non-aligned tail blocks use the independent internal Post-WU stage.
218+ 
219+### 5.4 Gate and Prepare Internal Templates
220+ 
221+| Template | Selection Condition | Main Implementation |
222+| --- | --- | --- |
223+| Gate generic chunk-cumsum | standalone Gate | Reuses `KdaGateCumsum::DispatchKdaGateCumsum`, supporting raw/activated gates and safe true/false |
224+| Gate/Prepare arch35 regbase | A5 key2, BF16, raw FP32, `A_log`, and safe gate | 16-row gate tiles and two sets of UB staging, with MTE2 and VEC overlapped; gk, qg, w seed, and kg factors are generated in batches |
225+| Prepare generic score/solve | key1, and A2/A3 key2 | Generic Aqk/Akk MMAD, causal mask, triangular solve, and seed write-back |
226+| Prepare arch35 key2 generic boundary | A5 key2 scenarios that miss batched fusion, including odd heads, safe false, tail, and varlen | Retains the same mathematical definition, using a bounded score queue, AIC/AIV handshaking, and a padded tail |
227+| Prepare arch35 safe-gate head-pair | A5 key2, safe gate, and `H_v % 2 == 0` | Two head lanes sharing two levels of score queues; 32-row score blocks; Akk solve overlapped with subsequent score movement |
228+| Prepare arch35 full-chunk fused-score | The previous row with `curT=64`, hitting the fused Post-WU | The two QK/Akk MMADs use a manual L1/L0 pipeline, and the Aqk/Akk masks and beta seeds are completed before write-back |
229+ 
230+### 5.5 Post-WU Internal Templates
231+ 
232+| Template | Selection Condition | Main Implementation |
233+| --- | --- | --- |
234+| Post-WU generic template | key1, and A2/A3 key2 | Generic Cube/Vector computes `w=Akk@w_seed` and `u=Akk@u_seed`, and generates kg |
235+| Post-WU arch35 full-chunk head-pair | A5 key2, `curT=64`, and fusion combination | At most 4 chunk tasks per batch and two head lanes per task; consumed directly after Prepare produces |
236+| Post-WU arch35 full-chunk standalone | A5 key2, `curT=64`, and not fused | W/U uses two sets of MTE2 staging with the MMAD/Fixpipe pipeline; Kg uses 16-row tiles and 32-row double buffering |
237+| Post-WU arch35 tail 16-63 | A5 key2, `16<=curT<64` | Padded Cube and write-back by valid rows, retaining the generic tail block synchronization protocol |
238+| Post-WU arch35 BF16 sub-16 | A5 key2, BF16, and `1<=curT<16` | `ComputeTailWuRegbaseArch35` fully stages the W/U seeds and Akk first; a single AIV accumulates in FP32 regbase and writes back, avoiding cross-AIV read/write contention when the W seed and W output reuse memory in place |
239+| Post-WU arch35 sub-16 fallback | A5 key2, non-BF16, and `1<=curT<16` | Uses the generic row-by-row Vector path; key1 and A2/A3 are covered by the Post-WU generic template |
240+ 
241+The worst-case staging of the Post-WU sub-16 template is
242+`2*15*128*sizeof(BF16)+15*64*sizeof(BF16)=9,600 B`, and no workspace or public output is added.
243+ 
244+### 5.6 FwdH Internal Templates
245+ 
246+| Template | Selection Condition | Main Implementation |
247+| --- | --- | --- |
248+| `KdaFwdHTileShapes128` | Generic backend and `V<=128` | KDA private arch22/arch35 FwdH; on arch22, complete chunks and chunks of at least 16 rows compute `kg^T@v_new` with Cube MMAD using BF16 inputs and FP32 accumulation, and only sub-16 tail blocks use the bounded Vector fallback; cross-chunk states remain FP32; supports dense/varlen/tail |
249+| `KdaFwdHTileShapes256` | Generic backend and `128<V<=256` | Chunked with V=256; other semantics are the same as the 128 template |
250+| `ChunkKdaFwdFwdH` arch35 dense | A5 key2, BF16, non-varlen, and `T%64=0` | 16-token sub-chunks and 4-slot L1 staging, with AIC MMAD/Fixpipe cooperating with AIV gate/state updates |
251+| arch35 dense fused Post-WU/FwdH | `fusePostWuIntoFwdH=true` | FwdH directly forms W/U and v_new from the Prepare seeds and does not export the unrequested qg/v_new/h |
252+ 
253+The state, W, and v_new L1 slots of the arch35 dense FwdH use ready/free credits to protect reuse. After
254+the last task completes, the AIV explicitly consumes the terminal free credit returned by the AIC before
255+exiting the kernel, preventing unconsumed credits from entering subsequent executions.
256+ 
257+### 5.7 Finalize Internal Templates
258+ 
259+| Template | Selection Condition | Main Implementation |
260+| --- | --- | --- |
261+| Finalize generic template | key1, and A2/A3 key2 | Generic Cube computes the two matmuls, and the AIV completes the FP32 addition, type conversion, and sequence-major write-back |
262+| Finalize arch35 dense-aligned pipeline | A5 key2, non-varlen, and `T%64=0` | `ProcessOutAicPipelinedArch35` uses dual L1 slots so that MTE2 of the next output tile overlaps with the current MTE1/MMAD/Fixpipe |
263+| Finalize arch35 full-chunk staged | A5 key2, `curT=64`, and dense-aligned pipeline not hit | Computes `scale*(qg@h)` and `Aqk@v_new` in stages, keeping FP32 intermediate results |
264+| Finalize arch35 tail 16-63 | A5 key2, `16<=curT<64` | Padded Cube computes the two results, and the AIV adds them by valid rows and writes back |
265+| Finalize arch35 BF16 sub-16 | A5 key2, BF16, and `1<=curT<16` | `ComputeTailOutputRegbaseRows` stages h, v_new, qg, and Aqk at one time, and after a single MTE2-to-V handoff directly accumulates the two terms in FP32 regbase, bypassing the tail block scalar EventID chain |
266+| Finalize arch35 sub-16 fallback | A5 key2, non-BF16, and `1<=curT<16` | Generic FP32 Vector row-by-row reduction; key1 and A2/A3 are covered by the Finalize generic template |
267+ 
268+The worst-case staging of the Finalize sub-16 template is
269+`(128*128+2*15*128+15*16)*sizeof(BF16)=40,928 B`, which is smaller than the existing `40,960 B` of
270+`gateWritebackBuf_` under key2 BF16. Therefore, the threshold is fixed at 16 and does not depend on
271+runtime probing.
272+ 
273+All templates keep the same formulas, public output order, optional output semantics, and layout contract.
274+Template selection must not change the interfaces of `ChunkKdaFwd` L0, `aclnnChunkKdaFwd`, or
275+`cann_ops_transformer.ops.chunk_kda_fwd`.
276+ 
277+## 6 Pipeline Design
278+ 
279+- The right matrix of Prepare resides in L1, avoiding repeated movement and transposition of K/K^T.
280+- The AIC organizes MTE2, MTE1, Cube, and Fixpipe with L1/L0 double buffering.
281+- The AIV uses input staging ping-pong so that MTE2 of the next tile overlaps with VEC of the current tile.
282+- The A5 VEC path uses regbase dual-issue; the main numerical computation remains FP32.
283+- Inter-sub-chunk merging uses an independent workspace region to avoid blocking the main tile pipeline.
284+ 
285+Performance conclusions use only `msopprof`. Target regression cases are maintained by the test directory
286+of this operator. For the A5 H96/T16K, K=V=128, chunk=64, and BF16 key scenario, the target value is
287+10 ms; the gating check reads all valid samples of the `ChunkKdaFwd` main kernel and of the `beta` layout
288+conversion for BSND inputs, and conservatively accumulates the maximum of each of the two components. The
289+check fails if the total exceeds 11 ms or if no valid sample is collected for either component. The same
290+key scenario is executed 50 consecutive times with non-zero inputs, and `attn_out/final_state/Aqk/Akk`
291+must be binary identical.
292+ 
293+## 7 Verification Matrix
294+ 
295+- Platforms: A2/A3/A5.
296+- dtypes: FP16/BF16.
297+- layouts: BSND/BNSD/TND/NTD.
298+- gates: raw/activated, safe true/false.
299+- Shapes: K=128, V=128/256, chunk=64/128, dense/varlen/tail/GVA.
300+- Attributes: final state, recomputation policy, and `state_v_first`.
Aattention/chunk_kda_fwd/examples/README_en.md+65-0
@@ -0,0 +1,65 @@
1+# ChunkKdaFwd Invocation Sample
2+ 
3+English | [简体中文](./README.md)
4+ 
5+`test_aclnn_chunk_kda_fwd.cpp` demonstrates the complete process of calling `ChunkKdaFwd` through the two-stage aclnn API, including:
6+ 
7+1. Initializing AscendCL, the device, the context, and the stream.
8+2. Constructing valid inputs, allocating device memory, and creating contiguous ND tensors.
9+3. Calling `aclnnChunkKdaFwdGetWorkspaceSize` to obtain the workspace size and the executor.
10+4. Allocating the workspace and calling `aclnnChunkKdaFwd` to perform the computation.
11+5. Copying `attn_out`, `Aqk`, and `Akk` to the host side, checking that they are finite values, and printing the first 5 elements.
12+6. Releasing the tensors, device memory, and AscendCL resources.
13+ 
14+## Sample Configuration
15+ 
16+The sample uses the most basic fixed-length BSND layout:
17+ 
18+| Parameter | Value or Description |
19+| :--- | :--- |
20+| Product | Ascend 950PR&950DT products, Atlas A3 products, Atlas A2 products |
21+| Data type | `q`, `k`, `v`, `attn_out`, `Aqk`, and `Akk` are BFLOAT16; `g`, `beta`, and `a_log` are FLOAT32 |
22+| Scale | B=1, S=64, H=HV=2, K=V=128, C=64 |
23+| `layout` | `BSND` |
24+| gate configuration | `use_gate_in_kernel=true`, `safe_gate=true`, `lower_bound=-5.0` |
25+| Optional inputs | `dt_bias`, `initial_state`, `cu_seqlens`, and `chunk_indices` are not passed |
26+| Optional outputs | Only the fixed `attn_out`, `Aqk`, and `Akk` are exported |
27+ 
28+The input and output shapes are as follows:
29+ 
30+| Tensor | Shape | Description |
31+| :--- | :--- | :--- |
32+| `q`, `k` | `[1, 64, 2, 128]` | Query and Key |
33+| `v` | `[1, 64, 2, 128]` | Value |
34+| `g` | `[1, 64, 2, 128]` | Raw gate when `use_gate_in_kernel=true` |
35+| `beta` | `[1, 64, 2]` | Delta coefficient |
36+| `a_log` | `[2]` | Gate parameter of each Value head |
37+| `attn_out` | `[1, 64, 2, 128]` | Attention output |
38+| `Aqk`, `Akk` | `[1, 2, 64, 64]` | Intra-chunk Query-Key coefficient matrices and triangular inversion results |
39+ 
40+For the API, formulas, and complete parameter constraints, see [aclnnChunkKdaFwd](../docs/aclnnChunkKdaFwd_en.md).
41+ 
42+## Compilation and Running
43+ 
44+Run the commands in the repository root directory of ops-transformer. For Ascend 950, use `--soc=ascend950`; for Atlas A2/A3, replace it with the SoC corresponding to the build environment.
45+ 
46+```bash
47+bash build.sh --pkg --soc=ascend950 --ops=chunk_kda_fwd
48+bash build_out/cann-ops-transformer-*.run --quiet --install-path=<install_path>
49+source <install_path>/vendors/custom_transformer/bin/set_env.bash
50+bash build.sh --run_example chunk_kda_fwd eager cust --soc=ascend950 --vendor_name=custom
51+```
52+ 
53+After successful execution, `build.sh` outputs `execute samples success`, and the sample itself outputs the following content:
54+ 
55+```text
56+ChunkKdaFwd execution succeeded.
57+ attnOut[0] = ...
58+ ...
59+ Aqk[0] = ...
60+ ...
61+ Akk[0] = ...
62+ ...
63+```
64+ 
65+The sample verifies that all printed results are finite values. The numerical results can be used to confirm that the invocation chain and output layout are correct. To perform a numerical comparison with the reference implementation, use the [pytest test framework](../tests/pytest/README_en.md).
Aattention/chunk_kda_fwd/tests/pytest/README_en.md+119-0
@@ -0,0 +1,119 @@
1+# ChunkKdaFwd Operator Test Framework
2+ 
3+English | [简体中文](./README.md)
4+ 
5+## Function Description
6+ 
7+Based on pytest and msopprof, this directory implements function, determinism, and performance verification for the ChunkKdaFwd operator:
8+ 
9+- **Function verification**: Calls the operator through `cann_ops_transformer.ops.chunk_kda_fwd`, and verifies the four input layouts, optional output matrices, BFLOAT16 gate parameter promotion, variable-length sequence tail blocks, and `state_v_first`.
10+- **Consistency verification**: Compares the results of BNSD, TND, and NTD inputs with the results of BSND inputs across layouts.
11+- **Determinism verification**: Repeats execution with the same input, and verifies that the specified outputs are binary identical.
12+- **Performance verification**: Uses msopprof to collect the device-side time consumption of the key scenario, and does not use Python wall time as the performance conclusion.
13+ 
14+## Current Implementation Scope
15+ 
16+### Case Coverage
17+ 
18+| Test Item | Covered Content |
19+| --- | --- |
20+| Input layouts | BSND, BNSD, TND, NTD |
21+| Data types | BFLOAT16 main inputs, FLOAT32 gate, and BFLOAT16 and FLOAT32 `A_log` and `dt_bias` |
22+| Optional outputs | Default return, `return_intermediate_states=True`, and `disable_recompute=True` |
23+| Variable-length sequences | NTD, `cu_seqlens=[0,64,65]`, and a sequence tail block of length 1 |
24+| State layout | `state_v_first=False` and `state_v_first=True` |
25+| Performance scenario | BFLOAT16, BSND, B=1, S=16384, H=HV=96, K=V=128, `chunk_size=64`, `safe_gate=True`, and `use_gate_in_kernel=True` |
26+ 
27+### Environment Configuration
28+ 
29+#### Prerequisites
30+ 
31+1. Complete the installation and environment variable configuration of CANN, `torch_npu`, and ops-transformer. For details, refer to [ops-transformer](../../../../README_en.md).
32+2. The custom operator package that contains the `ChunkKdaFwd` operator and the `torch_extension` wheel that provides `cann_ops_transformer.ops.chunk_kda_fwd` are installed.
33+3. Before performance verification, confirm that the `msopprof` command is available.
34+ 
35+## File Structure
36+ 
37+```text
38+pytest/
39+├── README.md # Test framework description
40+├── common.py # Input construction, layout conversion, and common assertions
41+├── test_chunk_kda_fwd.py # Function, consistency, and determinism cases
42+├── profile_a5_h96_t16k.py # Ascend 950 key scenario performance and 50-run determinism entry
43+└── check_msopprof.py # msopprof result parsing and performance gating
44+```
45+ 
46+## Usage
47+ 
48+Run the commands in the repository root directory of ops-transformer.
49+ 
50+### Running Function Cases
51+ 
52+```bash
53+python3 -m pytest -rA -s \
54+ attention/chunk_kda_fwd/tests/pytest/test_chunk_kda_fwd.py \
55+ -v
56+```
57+ 
58+### Running 50-Round Binary Determinism Verification
59+ 
60+```bash
61+python3 attention/chunk_kda_fwd/tests/pytest/profile_a5_h96_t16k.py \
62+ --check-determinism \
63+ --determinism-repeats 50
64+```
65+ 
66+The determinism verification performs binary consistency comparison on `attn_out`, `final_state`, `Aqk`, and `Akk`.
67+ 
68+### Running Ascend 950 Key Scenario Performance Verification
69+ 
70+Create the performance data directory:
71+ 
72+```bash
73+mkdir -m 700 -p output/chunk_kda_fwd_profiles
74+```
75+ 
76+Collect the `ChunkKdaFwd` main kernel:
77+ 
78+```bash
79+msopprof \
80+ --application="python3 attention/chunk_kda_fwd/tests/pytest/profile_a5_h96_t16k.py --repeats 8" \
81+ --output=output/chunk_kda_fwd_profiles/a5_h96_t16k \
82+ --aic-metrics=BasicInfo \
83+ --kernel-name=ChunkKdaFwd \
84+ --launch-count=5 \
85+ --warm-up=3 \
86+ --kill=off
87+```
88+ 
89+Collect the BSND `beta` layout conversion kernel:
90+ 
91+```bash
92+msopprof \
93+ --application="python3 attention/chunk_kda_fwd/tests/pytest/profile_a5_h96_t16k.py --repeats 8" \
94+ --output=output/chunk_kda_fwd_profiles/a5_h96_t16k_beta \
95+ --aic-metrics=BasicInfo \
96+ --kernel-name=Transpose_float16_int64_high_performance_10001 \
97+ --launch-count=5 \
98+ --warm-up=3 \
99+ --kill=off
100+```
101+ 
102+Verify the end-to-end device-side time consumption:
103+ 
104+```bash
105+python3 attention/chunk_kda_fwd/tests/pytest/check_msopprof.py \
106+ output/chunk_kda_fwd_profiles/a5_h96_t16k \
107+ --component Transpose_float16_int64_high_performance_10001 \
108+ output/chunk_kda_fwd_profiles/a5_h96_t16k_beta \
109+ --max-ms 11
110+```
111+ 
112+`check_msopprof.py` takes the maximum of the valid samples of the main kernel and of the `beta` layout conversion kernel respectively, and then sums them. The performance gating check fails when the total time consumption exceeds 11 ms or when no valid sample is collected for either component.
113+ 
114+## Output Description
115+ 
116+- The pytest function cases directly output the PASS/FAIL result of each case.
117+- `profile_a5_h96_t16k.py` outputs `binary_determinism=PASS` when determinism passes.
118+- `check_msopprof.py` outputs the sample count, minimum/average/maximum time consumption, and end-to-end maximum time consumption of each component.
119+- `output/chunk_kda_fwd_profiles/` contains local profiling artifacts and must not be committed to the code repository.
Aattention/compressor/README_en.md+283-0
@@ -0,0 +1,283 @@
1+# Compressor
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: Compressor is a preprocessing operator for SMLA and QLI in inference scenarios. It compresses the KV cache of every 4 or 128 tokens into one, and then each token performs DSA computation with these compressed KV caches. In long-sequence cases, Compressor effectively reduces the computation overhead. The main computation process is:
19+ 1. Performs matmul on the input $X$ and $W^{KV}$ to obtain $kv\_state$, and performs matmul on the input $X$ and $W^{Gate}$ followed by an Add operation with $Ape$ to obtain $score\_state$. $kv\_state$ and $score\_state$ are updated based on the input start_pos and cu_seqlens.
20+ 2. When coff is 2, rearranges the data of $kv\_state$ and $score\_state$.
21+ 3. Performs softmax on $score\_state$, multiplies the softmax result with $kv\_state$ through a Mul operation, and then performs a ReduceSum operation.
22+ 
23+- Computation formula:
24+ 
25+ 1. Computes the matmul:
26+ 
27+ $$
28+ C4A:\left[kv\_state^a, score\_state^a\right] = X @ \left[W^{aKV}, W^{aGate}\right], \left[kv\_state^b, score\_state^b\right] = X @ \left[W^{bKV}, W^{bGate}\right];
29+ $$
30+ 
31+ $$
32+ C128A:\left[kv\_state, score\_state\right] = X @ \left[W^{KV}, W^{Gate}\right]
33+ $$
34+ 
35+ 2. Computes the grouped addition:
36+ 
37+ $$
38+ C4A:score\_state_i^\prime = \left[score\_state_{\left[4(i-1)+1:4i,:\right]}^a; score\_state_{\left[4i+1:4(i+1),:\right]}^b\right] + Ape,~i=1,2,\cdots, \frac{s}{4};
39+ $$
40+ 
41+ $$
42+ C128A:score\_state_i^\prime = score\_state_{\left[128(i-1)+1:128i,:\right]} + Ape,~i=1,2,\cdots, \frac{s}{128};
43+ $$
44+ 
45+ 3. Computes the grouped softmax:
46+ 
47+ $$
48+ C4A:S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{4};
49+ $$
50+ 
51+ $$
52+ C128A:S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{128};
53+ $$
54+ 
55+ 4. Computes the Hadamard product:
56+ 
57+ $$
58+ C4A:(S_H)_i = S_i^\prime \odot \left[kv\_state^a_{\left[4(i-1)+1:4i,:\right]} ;kv\_state^b_{\left[4i+1:4(i+1),:\right]}\right],~i=1,2,\cdots, \frac{s}{4};
59+ $$
60+ 
61+ $$
62+ C128A:S_H = S_i^\prime \odot kv\_state;
63+ $$
64+ 
65+ 5. Computes the grouped sum along the compression axis:
66+ 
67+ $$
68+ C4A:C_{i}^{\text{Comp}} = \left[1\right]_{1\times8} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{4};
69+ $$
70+ 
71+ $$
72+ C128A:C_{i}^{\text{Comp}} = \left[1\right]_{1\times128} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{128};
73+ $$
74+ 
75+## Parameter Description
76+ 
77+<table style="undefined;table-layout: fixed; width: 1100px"><colgroup>
78+ <col style="width: 200px">
79+ <col style="width: 150px">
80+ <col style="width: 330px">
81+ <col style="width: 250px">
82+ <col style="width: 170px">
83+</colgroup>
84+<thead>
85+ <tr>
86+ <th>Parameter Name</th>
87+ <th>Input/Output/Attribute</th>
88+ <th>Description</th>
89+ <th>Data Type</th>
90+ <th>Data Format</th>
91+ </tr>
92+</thead>
93+<tbody>
94+ <tr>
95+ <td>x</td>
96+ <td>Input</td>
97+ <td><em>X</em> in the formula, indicating the original data without compression.</td>
98+ <td>FLOAT16, BFLOAT16</td>
99+ <td>ND</td>
100+ </tr>
101+ <tr>
102+ <td>wkv</td>
103+ <td>Input</td>
104+ <td><em>W<sup>KV</sup></em> in the formula, indicating the kv compression weight.</td>
105+ <td>FLOAT16, BFLOAT16</td>
106+ <td>ND</td>
107+ </tr>
108+ <tr>
109+ <td>wgate</td>
110+ <td>Input</td>
111+ <td><em>W<sup>Gate</sup></em> in the formula, indicating the gate compression weight.</td>
112+ <td>FLOAT16, BFLOAT16</td>
113+ <td>ND</td>
114+ </tr>
115+ <tr>
116+ <td>state_cache</td>
117+ <td>Input/Output</td>
118+ <td><em>[kv_state, score_state]</em> in the formula, indicating the historical data of kv_state and score_state.</td>
119+ <td>FLOAT</td>
120+ <td>ND</td>
121+ </tr>
122+ <tr>
123+ <td>ape</td>
124+ <td>Input</td>
125+ <td><em>Ape</em> in the formula, indicating positional biases.</td>
126+ <td>FLOAT</td>
127+ <td>ND</td>
128+ </tr>
129+ <tr>
130+ <td>cmp_ratio</td>
131+ <td>Attribute</td>
132+ <td>Used for sparse computation. It indicates the data compression ratio, and the value is an integer in the range [2, 128].</td>
133+ <td>INT32</td>
134+ <td>-</td>
135+ </tr>
136+ <tr>
137+ <td>state_block_table</td>
138+ <td>Optional input</td>
139+ <td>Indicates the block mapping table used for state_cache storage.</td>
140+ <td>INT32</td>
141+ <td>ND</td>
142+ </tr>
143+ <tr>
144+ <td>cu_seqlens</td>
145+ <td>Optional input</td>
146+ <td>Indicates the number of valid tokens in different batches. This parameter is a prefix-sum array. The value of each later element must be greater than or equal to the value of the previous element, and the first element must be 0.</td>
147+ <td>INT32</td>
148+ <td>ND</td>
149+ </tr>
150+ <tr>
151+ <td>seqused</td>
152+ <td>Optional input</td>
153+ <td>Indicates the number of tokens that actually participate in compression in different batches. When set to None, it is the same as the Sequence Length of each batch.</td>
154+ <td>INT32</td>
155+ <td>ND</td>
156+ </tr>
157+ <tr>
158+ <td>start_pos</td>
159+ <td>Optional input</td>
160+ <td>Indicates the computation start position. When set to None, computation starts from 0.</td>
161+ <td>INT32</td>
162+ <td>ND</td>
163+ </tr>
164+ <tr>
165+ <td>coff</td>
166+ <td>Optional attribute</td>
167+ <td>Indicates whether overlap data rearrangement is performed. coff=1: no overlap data rearrangement is required; coff=2: overlap data rearrangement is required. The default value is 1.</td>
168+ <td>INT32</td>
169+ <td>-</td>
170+ </tr>
171+ <tr>
172+ <td>cache_mode</td>
173+ <td>Optional attribute</td>
174+ <td>Indicates the storage mode of state_cache. cache_mode=1: continuous buffer; cache_mode=2: circular buffer. The default value is 1.</td>
175+ <td>INT32</td>
176+ <td>-</td>
177+ </tr>
178+ <tr>
179+ <td>state_cache_stride_dim0</td>
180+ <td>Optional attribute</td>
181+ <td>Indicates the stride of axis 0 of state_cache. The default value is 0.</td>
182+ <td>INT32</td>
183+ <td>-</td>
184+ </tr>
185+ <tr>
186+ <td>grad_enabled</td>
187+ <td>Optional attribute</td>
188+ <td>Whether to export the softmax_score/kv intermediate results (for the backward CompressorGrad). They are exported when true and not exported when false. The default value is false.</td>
189+ <td>BOOL</td>
190+ <td>-</td>
191+ </tr>
192+ <tr>
193+ <td>cmp_kv</td>
194+ <td>Output</td>
195+ <td>Indicates the compressed data.</td>
196+ <td>FLOAT16, BFLOAT16</td>
197+ <td>ND</td>
198+ </tr>
199+ <tr>
200+ <td>softmax_score</td>
201+ <td>Output</td>
202+ <td>Intermediate result of the softmax weight of compressed blocks. The output is valid only when grad_enabled is true.</td>
203+ <td>FLOAT</td>
204+ <td>ND</td>
205+ </tr>
206+ <tr>
207+ <td>kv</td>
208+ <td>Output</td>
209+ <td>Intermediate result of the raw KV of compressed blocks. The output is valid only when grad_enabled is true.</td>
210+ <td>FLOAT</td>
211+ <td>ND</td>
212+ </tr>
213+</tbody>
214+</table>
215+ 
216+- <term>Atlas A3 products</term> and <term>Atlas A2 products</term>:
217+ - cmp_ratio supports only 2/4/8/16/32/64/128.
218+ 
219+## Constraints
220+ 
221+- Dimension meanings of x: B (Batch Size) indicates the batch size of input samples, S (Sequence Length) indicates the sequence length of input samples, H (Head Size) indicates the size of the hidden layer, D (Head Dim) indicates the size of the minimum unit of the hidden layer, and T indicates the accumulated sequence length of input samples across all batches.
222+- Input shape restrictions:
223+ - state_cache supports the input shape [block_num, block_size, 2\*coff\*D] and requires block_num>0. When cache_mode=2, block_size >= coff * cmp_ratio + S - 1 must be satisfied.
224+ - If the BS dimensions of x are merged, that is, the input shape of x is [T,H]:
225+ - The input shape of cu_seqlens must be [B+1,]. The value of each element in this parameter indicates the total number of tokens of the current batch and all preceding batches, that is, the prefix sum. Therefore, the value of each later element must be greater than or equal to the value of the previous element, the first element must be 0, and the last element must be T.
226+ - seqused supports the input shape [B,] and requires that the number of valid tokens of each batch be less than or equal to the corresponding Sequence Length, that is, seqused[n] <= cu_seqlens[n+1] - cu_seqlens[n], and not less than 0.
227+ - When cache_mode=1, state_block_table supports the input shape [B, ceil(Smax/block_size)]. Smax is the maximum Sequence Length across batches, that is, Smax=max(start_pos)+max(cu_seqlens[n+1] - cu_seqlens[n]). When cache_mode=2, state_block_table supports the input shape [B].
228+ - cmp_kv has the output shape [min(T, T/cmp_ratio+B), D]: compressed_tokens + compressed_tokens + ... + compressed_tokens + pad.
229+ - If the BS dimensions of x are not merged, that is, the input shape of x is [B,S,H]:
230+ - cu_seqlens must be empty.
231+ - seqused supports the input shape [B,] and requires that the number of valid tokens of each batch be less than or equal to the corresponding Sequence Length, that is, seqused[n] <= S, and not less than 0.
232+ - When cache_mode=1, state_block_table supports the input shape [B, ceil(Smax/block_size)]. Smax is the maximum Sequence Length across batches, that is, Smax=max(start_pos)+S. When cache_mode=2, state_block_table supports the input shape [B].
233+ - cmp_kv has the output shape [B, ceil(S/cmp_ratio), D]: (compressed_tokens+pad0) + (compressed_tokens+pad1) + ... + (compressed_tokens+padN).
234+- Input value range restrictions:
235+ - This API supports generalization over B and S, with the following scenario restrictions:
236+ - In some long-sequence scenarios, an excessive computation amount may cause an error that the memory usage exceeds NPU memory. Note: the computation amount is affected by the input shape of x; the larger the values, the greater the computation amount. Typical long-sequence scenarios (that is, a large product of B and S, or a large T) include the following cases:
237+ 
238+ <table style="undefined;table-layout: fixed; width: 400px"><colgroup>
239+ <col style="width: 100px">
240+ <col style="width: 100px">
241+ <col style="width: 100px">
242+ </colgroup>
243+ <thead>
244+ <tr>
245+ <th>B</th>
246+ <th>S</th>
247+ <th>H</th>
248+ </tr>
249+ </thead>
250+ <tbody>
251+ <tr><td>100</td><td>65525</td><td>4096</td></tr>
252+ <tr><td>25</td><td>261120</td><td>4096</td></tr>
253+ <tr><td>100</td><td>131072</td><td>4096</td></tr>
254+ <tr><td>100</td><td>261120</td><td>4096</td></tr>
255+ </tbody>
256+ </table>
257+ - This API supports B, S, and T being 0. That is, input parameters whose shapes depend on the values of B, S, and T allow empty tensors, while other input parameters do not support empty tensors. In this scenario, state_cache is not updated, and the output cmp_kv is an empty tensor.
258+ - The value range of state_block_table elements is [0, block_num), where block_num is the size of dimension 0 of state_cache. Element values are directly used as block indices of state_cache, and out-of-range values cause invalid memory access. When an element value is 0: in cache_mode=1 (continuous buffer), the write state_cache operation skips this position; in cache_mode=2 (circular buffer), neither read nor write operations skip it. The operator does not perform duplicate validation; the caller must ensure the uniqueness of element values: in cache_mode=1, the element value 0 is the "unallocated" sentinel and may appear repeatedly, while non-zero values must be globally unique; in cache_mode=2, 0 is a valid physical block number, and all element values must be globally unique. Duplicates cause multiple logical blocks/batches to write the same physical block region, resulting in state_cache data being overwritten.
259+- Input attribute restrictions:
260+ - D of 128/512 is supported.
261+ - H of 1K to 10K, aligned to 512, is supported.
262+ - block_size of 1 to 1024 is supported.
263+ - The following three typical combination scenarios are supported:
264+ - C4A: D=512, coff=2, and cmp_ratio=4.
265+ - C4Li: D=128, coff=2, and cmp_ratio=4.
266+ - C128A: D=512, coff=1, and cmp_ratio=128.
267+- Deterministic computation and batch consistency:
268+ - A deterministic implementation is used by default, and multiple calls with the same input produce consistent results.
269+<!-- npu="950" id1 -->
270+ - <term>Ascend 950PR&950DT products</term>: batch consistency: set ACL_OPT_DETERMINISTIC to 3 through aclrtSetSysParamOpt() to enable batch consistency. After it is enabled, computation results are independent of the batch size and position.
271+<!-- end id1 -->
272+<!-- npu="A3,910b" id2 -->
273+ - <term>Atlas A3 products</term> and <term>Atlas A2 products</term>: batch consistency: set ACL_OPT_DETERMINISTIC to 3 through aclrtSetSysParamOpt() to enable batch consistency. After it is enabled, computation results are independent of the batch size and position.
274+<!-- end id2 -->
275+ 
276+## Invocation Description
277+ 
278+- <term>Ascend 950PR&950DT products</term>, <term>Atlas A3 products</term>, and <term>Atlas A2 products</term>:
279+ 
280+ | Invocation Mode | Sample Code | Description |
281+ | ---------- | -------- | --------------------------------------------- |
282+ | aclnn API | [test_aclnn_compressor](./examples/arch35/test_aclnn_compressor.cpp) | Calls the Compressor operator through the [aclnnCompressor](docs/aclnnCompressor_en.md) API. |
283+ | PyTorch API | - | Calls the Compressor operator through the [compressor](docs/torchapi_compressor_en.md) API. |
Aattention/flash_attention_score/docs/FA_design_en.md+281-0
@@ -0,0 +1,281 @@
1+## 1 Computation Process
2+ 
3+Figure 1 Training computation flowchart
4+ 
5+![Flowchart](./fig/FlashAttentionScore.png)
6+ 
7+It is implemented according to the forward computation flow of FlashAttention. The overall computation flow is as follows:
8+ 
9+1. After query and the transposed key undergo matmul computation, the initial attention_score is obtained. Then it is added to the positional encoding pse and multiplied by the scale factor scale_value. The result is processed by the select operation with atten_mask, masking the positions where atten_mask is true, to obtain masked_attention_score. That is, after select, the positions where atten_mask is true have a tiny negative value, which becomes 0 after the softmax computation, achieving the masking effect.
10+ 
11+2. To implement the FlashAttention acceleration algorithm, the FlashSoftmax operation is used to compute masked_attention_score, replacing the softmax operation in the original formula, and then the result is multiplied with value by matmul. Since the FlashSoftmax operation splits masked_attention_score along the Skv direction (the sequence length of the input key and value), a refresh process exists in the implementation, as follows:
12+ 1. Each FlashSoftmax computation operates on only one SkvSplit after splitting (SkvSplit is short for the sequence length after splitting along the Skv axis), and exp is recorded starting from the second loop iteration, where i is the loop variable after Skv splitting. For exp, i starts from 1. The formula for exp is as follows:
13+ $$
14+ exp[i] = e^{max_{i - 1} - max_{i}}
15+ $$
16+ 2. When i = 0, the computed MM[PV] result is directly saved to the ub of ub_attention_out[0].
17+ 3. Starting from i = 1, Mul and Add operations need to be added. That is, the previous MM[PV] result is multiplied by the current exp, and the product is added to the current MM[PV] result. The sum is saved to the ub of ub_attention_out[1]. This process repeats until the traversal of Skv is complete.
18+ 4. Since the division by sum in the FlashSoftmax computation is postponed to before the output of attention_out, the ub_attention_out in the ub needs to be divided by softmax_sum row by row at the end, and the final complete result is saved to the output memory attention_out(Final).
19+ 
20+## 2 Entry Point of Each Computation Stage
21+ 
22+Taking the S1 template as an example
23+ 
24+1. After obtaining the number of task blocks processed by each core, call the function: `Process`
25+2. Cube1 stage (Cube1 in the preceding figure), computation entry: `IterateBmm1`
26+3. Vector1 stage (Vector1 in the preceding figure), computation entry: `ProcessVec1`
27+4. Cube2 stage (Cube2 in the preceding figure), computation entry: `IterateBmm2`
28+5. Vector2 stage (Vector2 in the preceding figure), computation entry: `ProcessVec2`
29+ 
30+## 3 Tiling Design
31+ 
32+1. Atlas A2 training products
33+ The purpose of tiling is to find a more efficient NPU execution method. The original data volume is generally very large, and all computation cannot be completed with a single instruction call. Therefore, the data volume needs to be distributed across multiple cores for parallel computation. Each core also needs to consider how to achieve the best performance in cyclic computation. Different inputs may have different optimal execution methods, so the tiling policy needs to decide how to allocate data to the cores for computation.
34+ 
35+ Based on the hardware architecture, the AI Core is divided into two independent cores: AIC and AIV. The AIC and AIV cores have their own independent Scalar computing units, can independently load their own code segments, and execute separately. The separated AIC/AIV architecture allows AIC and AIV to execute in parallel. The paths for data exchange between AIC and AIV are L2 and GM (Global Memory, high-bandwidth memory), and the number of interactions between them greatly affects performance. In addition, due to the computing power difference between AIC and AIV, the two cores need to use different basic block sizes. Guided by the principle of minimizing the number of communications between AIC and AIV and maximizing the computing power, the CVtiling separation strategy was developed. This strategy effectively reduces the number of CV communications and selects different basic blocks for computation based on the buffer characteristics of different units, thereby improving operator performance.
36+ 
37+ For the FA operator, Vector computation involves multiple inputs, outputs, intermediate computation results, and double-buffer designs, so the buffer needs to be divided into multiple parts. In the optimal allocation scheme, the largest part is 32KB. Since the data type used by Vector computation is float32, the basic tiling block of Vector is 8 *1024. To fully leverage the computing power of the Cube, the data volume of one round of computation is allocated at a ratio of 1:16 between C and V. In addition, since the input data type on the Cube side is float16 and the output is float32, the basic block of the Cube is 128* 128. Therefore, a data volume of 128 * 1024 is allocated through nRatio=8. The pseudocode is as follows:
38+ 
39+ ```c++
40+ // C-Tiling: (S1_c_i,D)x(D,S2_c_i) => (S1_c_i, S2_c_i):(128,1024)
41+ // V-Tiling: (S1_v_i, S2_v_i) => (8,1024)
42+ 
43+ // Matmul computation on the C side
44+ Bmm((S1_c_i,D)x(D,S2_c_i)) => 128*1024 // Output result 128*1024, stored on the workspace
45+ // Vector computation on the V side
46+ for S1_c_i/S1_v_i=128/8:
47+ copy_gm_to_ub(S1_v_i*S2_v_i) // Copy the bmm result data from the bmm workspace
48+ Vector(S1_v_i,S2_v_i) // Perform Vector computation
49+ copy_ub_to_gm(S1_v_i*S2_v_i) // Vector computation completes and the final output data is obtained; copy it to GM
50+ 
51+ // Since the data computed on the Cube side is larger than that on the Vector side, Vector tiling needs to be performed again within the UB, which produces the S1-dimension ratio: S1_c_i/S1_v_i
52+ ```
53+ 
54+ In the preceding example, the ratio is enabled only in the S1 dimension, and the C/V computation lengths in the S2 dimension are the same. Of course, the ratio can also be enabled in both the S1 and S2 dimensions. The benefit is that the Cube can launch a large block of data at a time, avoiding the communication overhead caused by continuously launching small data blocks, and the buffer of the Cube unit can be used to the greatest extent.
55+ 
56+2. Ascend 950PR&950DT products
57+ Ascend 950PR&950DT products support both the separated AIC/AIV architecture and the fused AIC/AIV architecture. The interaction paths between AIC and AIV include L2, GM, and UB (Unified Buffer). UB is a more efficient interaction path between AIC and AIV than GM. AIC can directly output to the UB, and AIV uses the data on the UB to perform vec computation, which greatly reduces the data movement time. Since the UB capacity is limited, when D_v>128, the output of mm2 still needs to be saved to GM first; the output of mm1 is always saved in the UB.
58+ 
59+## 4 Pipeline Design
60+ 
61+To achieve ultimate performance, hardware resources must be fully utilized, and pipeline designs for different pipelines are usually required. The purpose of pipeline design is to make a certain pipeline reach the bound state as much as possible, so that a certain hardware unit keeps working and the performance upper limit is reached.
62+ 
63+### 4.1 V-Side Pipeline
64+ 
65+The V-side pipeline design needs to consider the data movement and computation processes of Vector. The main optimization method is double buffer.
66+Take the following pipeline task diagram as an example. Vec uses ping-pong to represent two data processing tasks, and each task needs to perform DataCopy movement and Clc computation in sequence. Data dependencies exist between tasks; for example, Clc can be processed only after DataCopy is completed.
67+After double buffer is applied between the ping-pong pipelines of the pipeline tasks in the figure, the run diagram shows that for the same piece of data, DataCopy movement and Clc computation are processed serially, while for different data slices, multiple tasks can be processed in parallel at the same time point, thereby achieving task parallelism and improving performance.
68+ 
69+![Design Diagram 2](./fig/fa_design_diagram_2.png)
70+ 
71+The memory resources occupied by the ping and pong computation data are independent of each other.
72+​The V-side computation process of FA-class fused operators involves many steps and is complex. A simple double buffer usually cannot cover all cases, so different computation pipeline arrangements appear. Different computation pipelines are applicable to different types of shape features, so as to achieve the best pipeline design for each type of feature.
73+ 
74+### 4.2 CV Pipeline
75+ 
76+1. Atlas A2 training products
77+ Fused operators usually include both Vector computation and Cube computation. For the FA operator, the V-side computation depends on the C-side computation result. If only the V-side pipeline is considered and the C side is ignored, the C side and the V side are likely to run as serial pipelines, which cannot achieve parallel computation or optimal fused operator performance. Therefore, the CV pipeline design was introduced. In addition, the CV pipeline behaves differently for different operators. The Cube dual-launch mechanism of FA (also known as inter-CV preload pipeline) enables pipeline optimization:
78+ 
79+ Under the pipeline characteristics of this scenario, there are few Vector computation nodes and the computation speed is fast. In the case of C:V=1:2 on <term>Atlas A2 training products</term>, the Cube data movement duration is sufficient to cover the Vector computation duration, so only the MTE2 time consumption of the Cube needs to be considered, finally achieving the MTE2 bound. Under the Cube dual-launch mechanism, two Cube computations are launched in advance, and Cube1 and Cube2 computations can be connected, maximizing Cube utilization and achieving the Cube bound.
80+ 
81+ ![FA Pipeline.jpg](../../../docs/zh/figures/fa_pipeline.png)
82+ 
83+2. Ascend 950PR&950DT products
84+ The CV pipeline design idea of Ascend 950PR&950DT products is basically the same as that of A2. The difference is that the number of Cube preload operations is 3: the mm2 computation starts only after three mm1 computations are completed. The purpose is to optimize the CV pipeline in the startup phase and make it tighter, so as to achieve the best overall performance.
85+ 
86+## 5 Multi-Template Design
87+ 
88+To allow different inputs to reuse the same tiling and pipeline, fused operators are implemented using templates. However, when all different inputs use the same set of templates, optimal performance and functional generalization cannot be achieved. Therefore, different templates need to be implemented based on the characteristics of the input shape.
89+The multi-template design idea of the FA (FlashAttentionScore, FA for short) fused operator is mainly as follows:
90+ 
91+- **Template splitting based on intra-core and inter-core splitting**
92+ 
93+ Since the hardware buffer size is limited and the data volume to be computed is huge, the computation cannot be completed at once, so tiling splitting is required. Different shapes lead to different splitting axes of the operator, and the splitting axes of the operator affect the functions and performance of the template. Simple elewise operators often fuse all axes into one axis for splitting, with simple logic, so their templates are relatively simple. However, fused operators combine multiple types of scenarios such as elewise, broadcast, reduce, and matmul, and their functions are complex. To meet high performance requirements, templates usually need to be split based on the splitting axes. To achieve optimal performance when splitting templates, the following points need to be considered:
94+ 
95+ a. Use all cores to prevent some cores from being idle.
96+ 
97+ b. The computation volume assigned to each core is relatively even, avoiding a situation where some cores compute excessively large amounts of data while the remaining cores stay idle.
98+ 
99+ c. The data volume processed between AIC and AIV must match their corresponding computing power, preventing long idle periods of AIC or AIV.
100+ 
101+ The FA operator contains five axes in total: B, N2 (the N of key and value), G (query_N/kv_N), S1 (the S of query), and S2 (the S of key and value). The splitting order is intra-core first and then inter-core. Intra-core splitting selects the splitting axes based on the basic block size, and inter-core splitting merges the remaining axes after intra-core splitting and then splits them based on the number of AI Cores. The splitting axes change with different shape sizes.
102+ 
103+ - Atlas A2 training products: From the Vector perspective, the FA operator is divided into the following template types. Templates are prioritized by sequence number; the smaller the number, the higher the priority, and the earlier the template is matched.
104+ 
105+ <table style="undefined;table-layout: fixed; width: 1576px">
106+ <colgroup>
107+ <col style="width: 170px">
108+ <col style="width: 170px">
109+ <col style="width: 310px">
110+ <col style="width: 212px">
111+ </colgroup>
112+ <thead>
113+ <tr>
114+ <th>Template</th>
115+ <th>Splitting Axis</th>
116+ <th>Entry Condition</th>
117+ <th>Applicable Scope</th>
118+ </tr>
119+ </thead>
120+ <tbody>
121+ <tr>
122+ <td>TNDSameAB template</td>
123+ <td>Split S1 in UB</td>
124+ <td>(B >= 4 and accumS1 >= 8192 and accumS2 >= 8192 and maxS1 >= 512 and maxS2 >= 512) or
125+ (maxS2 >= 5120 and maxS1 >= 5120). </td>
126+ <td>TND scenario</td>
127+ </tr>
128+ <tr>
129+ <td>TND template</td>
130+ <td>Split S1S2 in UB</td>
131+ <td>TND scenario, but the TNDSameAB template conditions are not met. </td>
132+ <td>TND scenario</td>
133+ </tr>
134+ <tr>
135+ <td>SameAB template</td>
136+ <td>Split S1 in UB</td>
137+ <td>Non-FP32, S2 >= 512 and D % 16 != 0 or D = 96, or S2 > 1024 and 128 < D < 196. </td>
138+ <td>General scenario</td>
139+ </tr>
140+ <tr>
141+ <td>S1S2 template</td>
142+ <td>Split S1S2 in UB</td>
143+ <td>S2 > 1024 and the preceding conditions are not met. </td>
144+ <td>General scenario</td>
145+ </tr>
146+ <tr>
147+ <td>S1 template</td>
148+ <td>Split S1D in UB</td>
149+ <td>(N2 * G * ((alignedS1 + alignedS2) * alignedD + alignedS2) * dtypeSize) >= 256 * 1024 or (N2 * G * (alignedS1 + alignedD) * alignedS2 * dtypeSize) >= 256 * 1024 or N2 * G * alignedS1 * alignedS2 * dtypeSize > 65536 * 2.</td>
150+ <td>General scenario</td>
151+ </tr>
152+ <tr>
153+ <td>B template</td>
154+ <td>Split B across cores</td>
155+ <td>The preceding conditions are not met.</td>
156+ <td>General scenario</td>
157+ </tr>
158+ </tbody>
159+ </table>
160+ 
161+ Each type of template has its own unique UB and Block splitting axes and can handle scenarios with inputs of a specific shape feature. The template is designed for that type of shape feature.
162+ - Ascend 950PR&950DT products
163+ 
164+ Since the basic block for intra-core splitting on Ascend 950PR&950DT products is 128 * 128, optimal performance can be achieved for all shapes. Therefore, only one set of templates is designed, supporting all shapes. This set of templates also does not distinguish whether the layout is TND.
165+ 
166+- **Template specialization based on special scenarios and specific optimizations**
167+ 
168+ Each fused operator has a basic template, supplemented by multiple specialized templates:
169+ 
170+ - The basic template covers functions and the performance of most inputs with shape features of this type.
171+ - Specialized templates cover the ultimate performance of specific scenarios. They are mainly optimized using special methods applicable to specific scenarios and are not suitable for generalization.
172+ 
173+ For example, the empty_input template specialized for the special empty tensor scenario, and the deterministic computation template designed based on the role-based Cube core management (RCM) optimization solution.
174+- **Template specialization based on different computation pipelines**
175+ 
176+ To fully leverage hardware advantages, fused operators usually require pipeline design to improve fused operator performance. Different pipeline designs greatly affect the code architecture. To improve the maintainability, testability, and readability of the code, templates need to be specialized based on different computation pipelines to achieve the ultimate performance in specific scenarios.
177+ 
178+### 5.1 Detailed Multi-Template Design
179+ 
180+- **Basic concepts:**
181+ 
182+ **CV basic block:** Indicates the data volume computed by the Cube or Vector in a single computation. It describes the data volume of a complete interaction between the Cube and Vector, and is usually equivalent to the **inter-core basic block**. On the chip, communication between the Cube and Vector has a certain overhead, so the CV basic block is set to a large value. A suitable data volume is generally 512*1024 (unit: Bytes). Since the intra-core buffers of the Cube and Vector are limited, the inter-core basic block may need to be completed through **multiple intra-core computations**.
183+ 
184+ **Intra-core basic block:**
185+ 
186+ If the CV basic block is too large, the Cube and Vector cores further split the CV basic block into basic blocks that fit the intra-core L0A, L0B, L0C, UB, and other buffer sizes. This is called the **intra-core basic block**. On the Cube side, the intra-core basic block is generally 32KB (unit: Bytes), which allows the DoubleBuffer capability of L0A and L0B to be fully utilized, and the computing power and bandwidth can also be used as fully as possible. On the Vector side, the basic block size is generally 32KB (unit: Bytes).
187+ 
188+ **xxx.i:** Indicates the size of an axis in the CV basic block after splitting. A basic block is generally 2-dimensional, and xxx.i indicates the size of one of its dimensions. xxx can be B, N2, G, S1, or S2;
189+ For example, if S1 = 512, S2 = 1024, S1.i = 64, and S2.i = 128, [S1, S2] is split into basic blocks of size [64, 128]. The basic block size of the S1 axis is 64, and the basic block size of the S2 axis is 128.
190+ 
191+ **xxx.o**: Indicates the number of segments into which an axis is divided after basic block splitting. xxx can be B, N2, G, S1, or S2;
192+ 
193+ For example, if S1 = 512, S2 = 1024, S1.i = 64, and S2.i = 128, [S1, S2] is split into basic blocks of size [64, 128], S1.o = 512 / 64 = 8, and S2.o = 1024 / 128 = 8, so a total of 8 * 8 basic blocks are split.
194+ - Atlas A2 training products: The FA operator is divided into the following template types based on the splitting axes. For the following FA templates, the S2 axis is currently not split across multiple cores:
195+ > 1. Template that splits the B, N2, G, and S1 axes across cores and splits the S1 and S2 axes within a core. This template is the most general template and supports all inputs (except TND):
196+ >
197+ > tiling code file: ops-transformer-dev/attention/flash_attention_score/op_host/flash_attention_score_tiling_general.cpp
198+ >
199+ > tiling code class: FlashAttentionScoreTilingS1s2Bn2gs1
200+ >
201+ > kernel code: ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_s1s2_bn2gs1.h
202+ >
203+ > Condition: S2 > 1024;
204+ > Rationale: Since the shape of the FlashSoftmax computation on the Vector side is [S1, S2], and the FlashSoftmax operation requires the S2 split size to be as large as possible, the basic block of the S2 axis is usually set to 1024 in FlashAttention. In addition, since FlashSoftmax has a refresh process that continuously updates the Softmax result when the S2 axis is split, more Softmax update processes appear in the computation flow when S2 is greater than 1024, which differs from other templates.
205+ > CV basic block selection:
206+ > S1.i: The default value is 64. When splitting by 64, if B *N2* G * S1.o exceeds the number of Vector cores, S1.i is set to 128. The purpose is to fully use all cores first when S1 is small, since full use of multiple cores delivers higher performance.
207+ > S2.i: 1024
208+ >
209+ >
210+ >
211+ > 2. Template that splits the B, N2, G, and S1 axes across cores and splits the S1 axis within a core
212+ >
213+ > tiling code file: ops-transformer-dev/attention/flash_attention_score/op_host/flash_attention_score_tiling_general.cpp
214+ >
215+ > tiling code class: FlashAttentionScoreTilingS1Bn2gs1
216+ >
217+ > kernel code: ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_s1_bn2gs1.h
218+ >
219+ > Condition: (128 < S2 <= 1024) || (N2 *G* (align(S1,16) + align(S2,16)) *align(D,16)* sizeof(INPUT_T)>= 512KB) || (N2 *G* (align(S1,16) + align(D,16)) *align(S2,16)* sizeof(INPUT_T))
220+ >
221+ > Rationale: When S2 < 1024, S2 is not split, so the FlashSoftmax result does not need to be updated, making the flow simpler. The generalization template described in point 1 is not used. Since the B-splitting template does not bring performance improvement when S2 > 128, this template is used by default. In addition, when S2 <= 128, if the input of matmul1 or matmul2 without the Batch axis has filled the entire L1, the B-splitting template is also not used. The B-splitting template means splitting the batch axis, and the inter-core basic block size is B.i *N2* G *S1* D (the D of query), or B.i *N2* G *S2* D (the D of key/value), or B.i *N2* G *S1* S2 (the Softmax result, that is, P). If B splitting meets the splitting requirements, B.i is at least 2, so the product of these axes inside B.i must be smaller than the L1 size. The condition check above is based on this: if the input of either matmul in Q *K or P* V is greater than the L1 size, using the B-splitting template brings no benefit.
222+ >
223+ > CV basic block selection:
224+ >
225+ > S1.i: It is dynamically adjusted based on the size of S2 to make the data volume of S1 * S2 as large as possible, in order to reduce the number of communications and the communication overhead.
226+ >
227+ > S2 is not split, and the basic block size of S2 = S2
228+ >
229+ >
230+ >
231+ > 3. Template that splits the B axis across cores. On the intra-core Cube side, S1 and S2 are not split; B.i, N2, and G are used as loop axes to open loops for processing batch matmul. The Vector core comprehensively splits B.i *N2* G * S1 intra-core to find the most suitable intra-core basic block.
232+ >
233+ > tiling code file: ops-transformer-dev/attention/flash_attention_score/op_host/flash_attention_score_tiling_general.cpp
234+ >
235+ > tiling code class: FlashAttentionScoreTilingB
236+ >
237+ > kernel code: ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_bn2gs1s2_b.h
238+ >
239+ > Condition: Other shapes that cannot enter the preceding two templates
240+ > Rationale: When S1, S2, and D are all small, the CV basic block is small. B.i, N2, and G are also included in the CV basic block, so the data volume of one CV interaction is larger, improving execution performance.
241+ 
242+ - Ascend 950PR&950DT products
243+ > 1. Template that splits the B, N2, G, and S1 axes across cores and splits the S1 and S2 axes within a core. This template is the most general template and supports all inputs (except TND):
244+ >
245+ > tiling code file: attention/flash_attention_score/op_host/arch35/flash_attention_score_tiling_regbase.cpp
246+ >
247+ > tiling code class: FlashAttentionScoreTilingS1S2Const, FlashAttentionScoreTilingVarLenConst
248+ >
249+ > kernel code: attention/flash_attention_score/op_kernel/arch35/flash_attention_score_kernel_train.h
250+ >
251+ > CV basic block selection:
252+ > S1.i: 128, or 64 in a few scenarios
253+ > S2.i: 128, or 256 in a few scenarios
254+ 
255+## 6 Programming Perspective
256+ 
257+### 6.1 AscendC High-Level APIs
258+ 
259+Currently, AscendC high-level APIs provide two programming modes. In the first mode, Vector is the main core and Cube is the secondary core. Vector0 and Vector1 independently launch Matmul tasks, and the two are independent of each other.
260+ 
261+In the second mode, Cube is the main core and Vector is the secondary core. V0 uniformly launches the Matmul task, and the result of this task is processed by V0 and V1 together; generally, V0 and V1 each process half. Currently, if a template name ends with "_sab", the template is a template with Cube as the main core. For example:
262+ 
263+```c++
264+ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_s1s2_bn2gs1_sab.h
265+ 
266+ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_sab.h
267+```
268+ 
269+For FlashAttention, using Cube as the main core allows the Matmul tasks of V0 and V1 to reuse the left matrix, and part of the output results can be accumulated in L0C, reducing the bandwidth dependency. Performance is better in most scenarios.
270+ 
271+### 6.2 AscendC Low-Level APIs
272+ 
273+Currently, some templates still do not use AscendC high-level APIs. For example:
274+ 
275+```c++
276+ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_basic.h
277+```
278+ 
279+This template uses Cube as the main core and Vector as the secondary core more thoroughly. All Matmul tasks are fully launched from the Cube side, and the Vector side is notified through synchronization.
280+ 
281+On Ascend 950PR&950DT products, FA is implemented using AscendC low-level APIs.
Aattention/flash_attention_score_grad/docs/FAG_design_en.md+405-0
@@ -0,0 +1,405 @@
1+## 1 Computation Process
2+ 
3+Figure 1 Training computation flow diagram
4+ 
5+![Flow diagram](./fig/FlashAttentionScoreGrad.png)
6+ 
7+This operator is implemented following the FlashAttention backward computation flow. The overall computation flow is as follows:
8+ 
9+1. Recompute p: $p = SimpledSoftmax(Mask(Matmul(query, key^T) + pse\_shift) * scale)$. This step recomputes the softmax result p from the FA forward process, and the computation result is saved in the UB.
10+2. Compute dp: $dp = Dropout(Matmul(dy, value^T))$. This computation includes a matmul computation and a dropout computation. In the matmul computation, the left matrix is dy and the right matrix is the transposed value.
11+3. Compute ds: $ds = p * Sub(dp, FlashSoftmaxGrad(dy, attention\_in))$. In this computation, the inputs of the FlashSoftmaxGrad computation are dy and the forward output attention\_in. The result of FlashSoftmaxGrad is subtracted from dp, and the final result is multiplied by p to obtain ds.
12+4. Compute dq: $dq = Matmul(ds, key) * scale$. This computation performs matmul on ds and key, and multiplies the result by scale to obtain dq.
13+5. Compute dk: $dk = Matmul(ds^T, query) * scale$. This computation performs matmul on the transposed ds and query, and multiplies the result by scale to obtain dk.
14+6. Compute dv: $dv = Matmul(DropOut(p)^T, dy)$. In this computation, the result p is processed by dropout, transposed, and then multiplied by dy through matmul.
15+ 
16+## 2 Entry Point of Each Computation Stage
17+ 
18+This section uses the sameAB template as an example.
19+ 
20+1. After obtaining the number of task blocks processed by each core, call the function: `Process`
21+2. Cube1 stage, step 1 in the flow, computation entry: `ComputeMM1`
22+3. Vector1 stage, steps 2-3 in the flow, computation entry: `ComputeVec`
23+4. Cube2 stage, steps 4-6 in the flow, computation entry: `ComputeMMDqkv`
24+ 
25+## 3 Tiling Design
26+ 
27+Tiling aims to find a more efficient way to execute on the NPU. The original data volume is generally very large, so all computations cannot be completed in a single instruction call. Therefore, the data needs to be distributed to multiple cores for parallel computation. In addition, each core needs to consider how to loop over the computations for optimal performance. Different inputs may have different optimal execution approaches, so a tiling strategy is required to decide how to distribute data to the cores for computation. Based on hardware architecture characteristics, an AI Core is divided into two independent cores: AIC and AIV. The AIC and AIV cores each have their own independent scalar computation unit and can load and execute their own code segments independently.
28+ 
29+### 3.1 <term>Atlas A2 training products</term>
30+ 
31+The separated AIC/AIV architecture of the <term>Atlas A2 training products</term> chip allows the AIC and AIV to execute in parallel. Data is exchanged between the AIC and AIV through L2 and GM (Global Memory, the high-bandwidth memory). The number of interactions between the two greatly affects performance. In addition, due to the difference in computing power between the AIC and AIV, the two need to use different basic block sizes. Guided by the principles of minimizing the number of AIC-AIV communications and maximizing computing power, the CV tiling separation strategy was developed. This strategy effectively reduces the number of CV communications and, based on the buffer characteristics of different units, selects different basic blocks for computation, thereby improving operator performance.
32+ 
33+ For the FAG operator, the Vector computation involves multiple inputs, outputs, intermediate computation results, and the double-buffer design. Therefore, the buffer needs to be split into multiple portions, and the largest portion in the optimal allocation scheme is 32KB. Because the Vector computation uses the float32 data type, the basic tiling block of the Vector is 8 *1024. To fully utilize the computing power of the Cube, the data volume of one round of computation is allocated at a ratio of 1:16 between C and V. In addition, because the input data type on the Cube side is float16 and the output is float32, the basic block of the Cube is 128* 128. Therefore, a data volume of 128 * 1024 is derived through the ratio nRatio=8. The pseudocode is as follows:
34+ 
35+```c++
36+// C-Tiling: (S1_c_i,D)x(D,S2_c_i) => (S1_c_i, S2_c_i):(128,1024)
37+// V-Tiling: (S1_v_i, S2_v_i) => (8,1024)
38+ 
39+// matmul computation on the C side
40+Bmm((S1_c_i,D)x(D,S2_c_i)) => 128*1024 // Output result 128*1024, placed on the workspace
41+// Vector computation on the V side
42+for S1_c_i/S1_v_i=128/8:
43+ copy_gm_to_ub(S1_v_i*S2_v_i) // Copy in the bmm result data from the bmm workspace
44+ Vector(S1_v_i,S2_v_i) // Perform the Vector computation
45+ copy_ub_to_gm(S1_v_i*S2_v_i) // Vector computation ends, final output data is obtained, copy to GM
46+ 
47+// Because the data volume computed on the Cube side is larger than that on the Vector side, Vector Tiling needs to be performed again inside the UB, which produces the S1-dimension ratio: S1_c_i/S1_v_i
48+```
49+ 
50+In the preceding example, the ratio is enabled only in the S1 direction, and the C/V computation lengths are the same in the S2 direction. The ratio can also be enabled in both the S1 and S2 directions. The advantage is that the Cube can dispatch a large block of data at a time, avoiding the communication overhead caused by continuously dispatching small blocks, and the buffer of the Cube unit can be used to the maximum extent.
51+ 
52+### 3.2 **<term>Ascend 950PR&950DT products</term>**
53+ 
54+The Ascend 950PR&950DT products also adopt the separated AIC/AIV architecture and retain the AIC-AIV parallel execution feature. In addition, high-speed data exchange paths between the AIC and AIV, L0C->UB and UB->L1, are added. These paths reduce the cost of CV interactions and increase the pipelining parallelism. Compared with the complex tiling block strategy of the <term>Atlas A2 training products</term>, good performance can be achieved on the 950 with only one basic block.
55+ 
56+ For the FAG operator, the Vector computation involves multiple inputs, outputs, intermediate computation results, and the double-buffer design. Therefore, the buffer needs to be split into multiple portions, and the largest portion in the optimal allocation scheme is 32KB. Because the Vector computation uses the float32 data type, the basic tiling block of the Vector is 64 *128. Because the number of Cube cores to Vector cores is 1:2, to fully utilize the computing power of the Cube cores, the Cube side uses a basic block of 128* 128. That is, after each Cube core computes a 128 * 128 data block, the block is evenly split between the two Vector cores for processing. The pseudocode is as follows:
57+ 
58+```c++
59+// C-Tiling: (S1_i,D)x(D,S2_i) => (S1_i, S2_i):(128,128)
60+// V-Tiling: (S1_i / 2, S2_i) => (64,128)
61+ 
62+// matmul computation on the C side
63+Bmm((S1_i,D)x(D,S2_i)) => 128*128 // value@dy,q@k output result 128*128, placed on L0C
64+fixp_l0c_to_ub(S1_i / 2,S2_i); => v0: 64*128, v1: 64*128 // Through the CV path L0C->UB, evenly written to the UBs of the two V cores
65+// Vector computation on the V side
66+Vector(S1_i / 2,S2_i) // Perform the Vector computation
67+copy_ub_to_l1(S1_i/2*S2_i) // Vector computation ends, output data is obtained, and through the CV path UB->L1, data from the two V cores is aggregated and copied to L1
68+Bmm((S1_i,S2_i)*(S2_i,D)); // dq, final result output to the workspace
69+Bmm((S2_i,S1_i)*(S1_i,D)); // dkv, final result output to the workspace
70+```
71+ 
72+In the preceding example, the two copy instructions fixp_l0c_to_ub and copy_ub_to_l1 copy the data computed on the Cube side to the Vector side and the data computed on the Vector side to the Cube side, respectively. Both copies use intra-core high-speed channels.
73+ 
74+## 4 Pipeline Design
75+ 
76+To pursue ultimate performance, hardware resources must be fully utilized, which usually requires pipeline design for different pipelines. The purpose of pipeline design is to make one pipeline reach a bound state, keeping a hardware unit constantly working to reach the performance limit.
77+ 
78+### 4.1 V-Side Pipeline
79+ 
80+The V-side pipeline design needs to consider the data movement and computation processes of the Vector. The main optimization method is double buffer.
81+Taking the following pipeline task diagram as an example, the Vec function is split into two pipeline tasks: subA and subB, each focusing on a single function. The data to be processed is split into two slices, and ping-pong denotes the two data processing tasks. Each task needs to sequentially perform the DataCopy movement and Clc computation operations (Clc denotes the Vector computation). Arrows between tasks indicate data dependencies. For example, subB can process Clc only after subA completes DataCopy.
82+As shown in the figure, without pipeline design, the movement and computation tasks are executed serially, causing a stream interruption. That is, the movement pipeline stays idle after the first DataCopy completes, and it continues to work and starts the second DataCopy only after the first batch of moved-in data is computed and moved out (the Vector computation and move-out pipelines have the same problem). In this case, the performance is usually extremely poor.
83+ 
84+![Design diagram 1](./fig/fag_design_diagram_1.png)
85+ 
86+After double buffer is applied between the ping-pong pipelines of the pipeline tasks in the preceding figure, the pipeline tasks run as shown in the following diagram. As shown in the run diagram, for the same data slice, the DataCopy movement and Clc computation have a dependency and need to be processed serially. For different data slices, multiple tasks can be processed in parallel at the same time point, thereby achieving task parallelism and improving performance.
87+ 
88+![Design diagram 2](./fig/fag_design_diagram_2.png)
89+ 
90+The memory resources occupied by the two computation data blocks, ping and pong, are independent of each other.
91+The V-side computation of the FAG fused operator involves many steps and complex situations. A simple double buffer usually cannot cover all cases, so different computation pipeline arrangements exist. Different computation pipelines are applicable to different types of shape features, so as to achieve the best pipeline design for each type of shape feature.
92+ 
93+### 4.2 CV Pipeline
94+ 
95+Fused operators usually include both Vector computation and Cube computation. For the FA operator, the V-side computation depends on the C-side computation results. If only the V-side pipeline is considered and the C side is ignored, the C side and V side are likely to run as serial pipelines. This cannot achieve parallel computation and cannot make the fused operator perform at its best. Therefore, the CV pipeline design is introduced. In addition, the CV pipeline behaves differently for different operators. The Cube dual-dispatch mechanism of FA/FAG, also known as the inter-CV preload pipeline, enables pipeline optimization in two scenarios:
96+ 
97+- Total C-side time > Total V-side time
98+ 
99+ In this scenario, the Vector computation has few nodes and runs fast. On the <term>Atlas A2 training products</term> with C:V=1:2, the Cube movement duration is long enough to cover the Vector computation duration, so only the MTE2 time of the Cube needs attention, achieving MTE2 bound. Under the Cube dual-dispatch mechanism, two Cube computations are dispatched in advance, so the Cube1 and Cube2 computations can be chained, maximizing Cube utilization and achieving Cube bound.
100+ 
101+![Design diagram 3](./fig/fag_design_diagram_3.png)
102+ 
103+- Total C-side time < Total V-side time
104+ 
105+ In this scenario, the Vector computation has many nodes and is the bottleneck. The C-side movement is not enough to cover the V-side pipeline, so CV pipeline arrangement is needed to achieve CV parallelism as much as possible. The most common optimization method is early dispatch on the C side.
106+ 
107+![Design diagram 4](./fig/fag_design_diagram_4.png)
108+ 
109+The C side dispatches two Cube computations consecutively, so when the V side finishes computing the previous round, it can immediately start the computation of the current round without waiting for the data of Cube1. This keeps the V side constantly working and achieves Vector bound.
110+ 
111+## 5 Multi-Template Design
112+ 
113+Templates are used to implement fused operators so that different inputs can reuse the same tiling and pipeline. However, using the same template for all different inputs cannot achieve optimal performance and functional generalization. Therefore, different templates need to be distinguished based on the characteristics of the input shape.
114+The multi-template design of the FAG (FlashAttentionScoreGrad, abbreviated as FAG) fused operator mainly includes:
115+ 
116+- **Splitting templates based on intra-core and inter-core partitioning**
117+ 
118+ Because the hardware buffer size is limited while the data volume to be computed is huge, the computation cannot be completed at one time, so tiling partitioning is required. Different shapes lead to different partitioning axes of the operator, and the partitioning axes affect the function and performance of the template. Simple elewise operators usually fuse all axes into one axis for partitioning, which has simple logic and therefore relatively simple templates. In contrast, fused operators combine elewise, broadcast, reduce, matmul, and other scenarios with complex functions. To meet high performance requirements, templates often need to be split based on partitioning axes. To achieve optimal performance when splitting templates, consider the following points:
119+ 
120+ - a. Use all cores to prevent some cores from being idle
121+ 
122+ - b. Distribute the computation volume relatively evenly across cores, preventing some cores from computing too much data while other cores stand by.
123+ 
124+ - c. Match the data volume processed between the AIC and AIV with their respective computing power, preventing long idle periods of the AIC or AIV.
125+ 
126+ The FAG operator contains five axes in total: B, N2 (the N of key and value), G (query_N/kv_N), S1 (the S of query), and S2 (the S of key and value). The partitioning order is intra-core first and then inter-core. For intra-core partitioning, the partitioning axes are selected based on the basic block size. For inter-core partitioning, the remaining axes after intra-core partitioning are merged and then partitioned based on the number of AI Cores. Because shapes differ in size, the partitioning axes change. From the Vector perspective, the FAG operator is divided into the following template types. Templates are prioritized by sequence number: the smaller the sequence number, the higher the priority and the earlier the match.
127+ 
128+- **Depending on the applicable scope, the FAG operator is divided into the following template types:**
129+ 
130+**<term>Atlas A2 training products</term>**
131+ 
132+<table style="undefined;table-layout: fixed; width: 1576px">
133+<colgroup>
134+ <col style="width: 170px">
135+ <col style="width: 170px">
136+ <col style="width: 310px">
137+ <col style="width: 212px">
138+</colgroup>
139+<thead>
140+ <tr>
141+ <th>Template</th>
142+ <th>Partitioning Axis</th>
143+ <th>Entry Condition</th>
144+ <th>Applicable Scope</th>
145+ </tr>
146+</thead>
147+<tbody>
148+ <tr>
149+ <td>B template</td>
150+ <td>Multi-core partitioning on B</td>
151+ <td>Non-FP32, bestBasicBlockNum = S1 >= 4 ? 64 * 128 / 4 * 3 : 64 * 128;
152+N1 * G * alignedS1 * alignedS2 <= bestBasicBlockNum. </td>
153+ <td>Normal scenario</td>
154+ </tr>
155+ <tr>
156+ <td>N2 template</td>
157+ <td>Multi-core partitioning on N2</td>
158+ <td>Non-FP32, S1 * S2 < 64 * 128. </td>
159+ <td>Normal scenario</td>
160+ </tr>
161+ <tr>
162+ <td>SameAB template</td>
163+ <td>UB partitioning on S1S2, multi-core partitioning on S2</td>
164+ <td>In the deterministic computation scenario, if non-FP32 or S1 >= 1024 or S2 >= 1024, the SameAB template is used; otherwise, template 1.2 is used.<br>
165+ In the TND scenario, if average S1 >= 1024 and average S2 >= 1024, the SameAB template is used; otherwise, template 1.1 is used.<br>
166+ In the normal scenario, S1 >= 1024 or S2 >= 1024.
167+ </td>
168+ <td>Deterministic computation scenario<br>TND scenario<br>Normal scenario</td>
169+ </tr>
170+ <tr>
171+ <td>S1S2 template</td>
172+ <td>UB partitioning on S1S2, multi-core partitioning on N2</td>
173+ <td>In the deterministic computation scenario, the SameAB template conditions are not met.<br>
174+ In the normal scenario, the conditions of other templates are not met.
175+ </td>
176+ <td>Deterministic computation scenario<br>Normal scenario</td>
177+ </tr>
178+ <tr>
179+ <td>TND template</td>
180+ <td>UB partitioning on S1S2, multi-core partitioning on S2</td>
181+ <td>The SameAB template conditions are not met.</td>
182+ <td>TND scenario</td>
183+ </tr>
184+</tbody>
185+</table>
186+ 
187+**Ascend 950PR&950DT products**
188+ 
189+<table style="undefined;table-layout: fixed; width: 1576px">
190+<colgroup>
191+ <col style="width: 170px">
192+ <col style="width: 170px">
193+ <col style="width: 310px">
194+ <col style="width: 212px">
195+</colgroup>
196+<thead>
197+ <tr>
198+ <th>Template</th>
199+ <th>Partitioning Axis</th>
200+ <th>Entry Condition</th>
201+ <th>Applicable Scope</th>
202+ </tr>
203+</thead>
204+<tbody>
205+ <tr>
206+ <td>BN2</td>
207+ <td>Multi-core partitioning on BN2</td>
208+ <td>S1<=128 and S2 <= 128 and G=1, non-FP32.</td>
209+ <td>Normal scenario</td>
210+ </tr>
211+ <tr>
212+ <td>BN2S2 template</td>
213+ <td>Multi-core partitioning on BN2S2</td>
214+ <td>Non-FP32, TND, B * N2 * S2 > cubeCoreNum and G = 1. </td>
215+ <td>Normal scenario</td>
216+ </tr>
217+ <tr>
218+ <td>Deterministic computation template</td>
219+ <td>Multi-core partitioning on BN2GS1S2</td>
220+ <td>Deterministic computation is enabled.</td>
221+ <td>Deterministic computation scenario</td>
222+ </tr>
223+ <tr>
224+ <td>BN2GS1S2 template</td>
225+ <td>Multi-core partitioning on BN2GS1S2</td>
226+ <td>None of the preceding scenarios is met.</td>
227+ <td>Normal scenario</td>
228+ </tr>
229+</tbody>
230+</table>
231+ 
232+Each template type has its own UB and Block partitioning axes and can process scenarios with inputs of specific shape features. The template is designed for that type of shape feature.
233+ 
234+- **Specializing templates based on special scenarios and specific optimizations**
235+ 
236+ Each fused operator has a base template supplemented by multiple specialized templates:
237+ 
238+ - The base template covers the functionality and the performance for most inputs of this type of shape feature.
239+ - The specialized templates deliver ultimate performance for specific scenarios. They mainly apply optimizations based on special methods that suit specific scenarios and are not suitable for generalization.
240+ 
241+ For example, the empty_input template is specialized for the special scenario of empty tensors, and the deterministic computation template is designed based on the role-based Cube core management (RCM) optimization solution.
242+ 
243+- **Specializing templates based on different computation pipelines**
244+ 
245+ To fully exploit hardware advantages, fused operators usually require pipeline design to improve performance. Different pipeline designs have a great impact on the code architecture. To improve the maintainability, testability, and readability of the code, templates need to be specialized based on different computation pipelines to achieve ultimate performance in specific scenarios.
246+ 
247+### 5.1 Detailed Multi-Template Design
248+ 
249+- **Basic concepts:**
250+ 
251+ **CV basic block:** Represents the data volume computed by the Cube or Vector in a single computation. It describes the data volume of one complete Cube-Vector interaction and is usually equivalent to the **inter-core basic block**. On the chip, communication between the Cube and Vector has a certain overhead, so the CV basic block is set relatively large. A suitable data volume is generally 512*1024 (unit: Bytes). In addition, because the intra-core buffers of the Cube and Vector are limited, the inter-core basic block may need to be completed through **multiple intra-core computations**.
252+ 
253+ **Intra-core basic block:**
254+ 
255+ If the CV basic block is too large, the Cube and Vector further split the CV basic block inside the cores into basic blocks suitable for the sizes of intra-core L0A, L0B, L0C, UB, and so on. This is called the **intra-core basic block**. On the Cube side, the intra-core basic block is generally 32KB (unit: Bytes), which fully unfolds the double buffer capability of L0A and L0B and makes full use of the computing power and bandwidth. On the Vector side, the basic block size is generally 32KB (unit: Bytes).
256+ 
257+ **xxx.i:** Represents the size of an axis in the CV basic block after partitioning. A basic block is generally two-dimensional, and xxx.i represents the size of one dimension. xxx can be B, N2, G, S1, or S2.
258+ For example, S1 = 512, S2 = 1024, S1.i = 64, S2.i = 128 means that [S1, S2] is partitioned into basic blocks of size [64, 128]; the basic block size of the S1 axis is 64, and the basic block size of the S2 axis is 128.
259+ 
260+ **xxx.o**: Represents the number of parts of an axis after basic block partitioning. xxx can be B, N2, G, S1, or S2;
261+ 
262+ For example, S1 = 512, S2 = 1024, S1.i = 64, S2.i = 128 means that [S1, S2] is partitioned into basic blocks of size [64, 128]; S1.o = 512 / 64 = 8, S2.o = 1024 / 128 = 8, and a total of 8 * 8 basic blocks are partitioned.
263+ 
264+ - **<term>Atlas A2 training products</term>**
265+ The templates of the FAG operator are divided as follows. For the following templates, the larger the sequence number, the higher the priority of the template. Template 1 is the generalized template (supports all shapes):
266+ 
267+ > 1. Template with inter-core partitioning on the B, N2, G, and S1 axes and intra-core partitioning on the S1 and S2 axes:
268+ >
269+ > tiling code file: ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_s1s2_bn2gs1s2.cpp
270+ >
271+ > tiling code class: FlashAttentionScoreGradTilingS1s2Bn2gs1s2
272+ >
273+ > kernel code: ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2.h
274+ >
275+ > Condition: supports all shapes. If other templates do not support a shape, this template is used.
276+ > Rationale: Following the most general approach, this template supports all shapes. However, because the intra-core computation can process only S1 * S2 at a time, frequent CV interaction overhead occurs and the performance is poor when both S1 and S2 are small. When both S1 and S2 are small, the flow routes to the following templates.
277+ >
278+ >
279+ >
280+ > 2. Template with inter-core partitioning on the B and N2 axes and intra-core partitioning on the G, S1, and S2 axes. This template optimizes performance for scenarios where both S1 and S2 are small and (B * N2 is relatively large or G = 1) through pure core partitioning changes.
281+ >
282+ > tiling code file: ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_s1s2_bn2.cpp
283+ >
284+ > tiling code class: FlashAttentionScoreGradTilingS1s2Bn2
285+ >
286+ > kernel code: ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2.h
287+ >
288+ > Condition: (S1 < 1024) && (S2 < 1024) && (B *N2* 2 > CoreNum || G == 1)
289+ >
290+ > Rationale: When both S1 and S2 are small and both B and N2 are large, B and N2 are used for core partitioning. N2.i is not partitioned intra-core, and the Cube and Vector computations are performed N2.i times in a loop.
291+ >
292+ >
293+ >
294+ > 3. Template with inter-core partitioning on the B and N2.o axes and intra-core partitioning on the N2.i, G, S1, and S2 axes. This template optimizes performance for scenarios where G *S1* S2 are all small by partitioning the N2 axis intra-core and also partitioning the N2.i axis intra-core to accelerate the Vector computation. Compared with template 2, template 3 is more complex. Template 3 also partitions the N2.i axis in intra-core computations, making the computation volume of each round larger.
295+ >
296+ > tiling code file: ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_ngs1s2_bn.cpp
297+ >
298+ > tiling code class: FlashAttentionScoreGradUngs1s2BbnTiling
299+ >
300+ > kernel code: FlashAttentionScoreGradUngs1s2Bbn
301+ >
302+ > Condition: G *S1* S2 <= 32KB && G = 1 && S2 < 1536
303+ >
304+ > Rationale: When G *S1* S2 is smaller than 32KB, part of the N2 axis can be partitioned intra-core to enlarge the CV basic block. Meanwhile, within the Vector core, N2.i is also partitioned so that the computation volume of each Vector round is larger, improving Vector utilization.
305+ >
306+ >
307+ >
308+ > 4. Template with inter-core partitioning on the B axis and intra-core computation on the B.i, N2, G, S1, and S2 axes. This template optimizes performance when N2 *G* S1 * S2 is relatively small.
309+ >
310+ > tiling code file: ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_bngs1s2_b.cpp
311+ >
312+ > tiling code class: FlashAttentionScoreGradUbngs1s2BbTiling
313+ >
314+ > kernel code: FlashAttentionScoreGradUngs1s2Bbn
315+ >
316+ > Condition: N2 *G* S1 *S2 <= 64* 128
317+ >
318+ > Rationale: If B.i alone is to be placed into the CV basic block, the inner axes N2 *G* S1 * S2 need to be small enough. Generally, when this value is smaller than 64KB, the data volume of Bmm1 and Bmm2 usually does not exceed half of L1, making partitioning on the B axis meaningful. Otherwise, a single Matmul fills L1, and the data move-in between multiple Matmuls cannot be parallelized with the computation.
319+ 
320+ - **<term>Ascend 950PR&950DT products</term>**
321+ The templates of the FAG operator are divided as follows. For the following templates, the larger the sequence number, the higher the priority of the template. Template 1 is the generalized template (supports all shapes). Although multiple templates exist, only two template files exist in the implementation: one for deterministic computation and the other for non-deterministic computation. The non-deterministic computation template includes three partitioning templates: BN2, BN2S2, and BN2GS1S2. In the code, their implementation logic is isolated through template parameters:
322+ 
323+ > 1. Template with inter-core partitioning on the B, N2, G, S1, and S2 axes:
324+ >
325+ > tiling code file: ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpp
326+ >
327+ > tiling code class: FlashAttentionScoreGradTilingUs1s2Bs2Regbase
328+ >
329+ > kernel code: ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h
330+ >
331+ > Condition: supports all shapes. If other templates do not support a shape, this template is used.
332+ > Rationale: Following the most general approach, this template supports all shapes. However, the performance is not optimal in some scenarios, in which case the flow routes to the following specialized templates based on the scenario.
333+ >
334+ > 2. Template with inter-core partitioning on the B and N2 axes:
335+ >
336+ > tiling code file: ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpp
337+ >
338+ > tiling code class: FlashAttentionScoreGradTilingUs1s2Bs2Regbase
339+ >
340+ > kernel code: ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h
341+ >
342+ > Condition: S1<=128 and S2 <= 128 and G=1, non-FP32.
343+ > Rationale: When S1<=128 and S2 <= 128 and G=1, the S1S2 axes are not partitioned across cores, meaning no multi-core accumulation exists. This eliminates the preceding GM zeroing and the following cast and musl computations, resulting in better performance.
344+ >
345+ > 3. Template with inter-core partitioning on the B, N2, and S2 axes:
346+ >
347+ > tiling code file: ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpp
348+ >
349+ > tiling code class: FlashAttentionScoreGradTilingUs1s2Bs2Regbase
350+ >
351+ > kernel code: ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h
352+ >
353+ > Condition: Non-FP32, TND, B *N2* S2 > cubeCoreNum and G = 1.
354+ > Rationale: In the scenario of B *N2* S2 > cubeCoreNum and G = 1, the dk and dv computations do not require multi-core accumulation. This eliminates the preceding GM zeroing for dk and dv and the following cast and musl computations, resulting in better performance.
355+ >
356+ > 4. Deterministic computation template
357+ >
358+ > tiling code file: ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpp
359+ >
360+ > tiling code class: FlashAttentionScoreGradTilingUs1s2Bs2Regbase
361+ >
362+ > kernel code: ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel_deter.h
363+ >
364+ > Condition: Deterministic computation is enabled.
365+ > Rationale: The deterministic computation scenario has strict requirements on core partitioning. A specific core partitioning method avoids multi-core accumulation at the same address, achieving deterministic computation.
366+ 
367+##
368+ 
369+## 6 Programming Perspective
370+ 
371+### 6.1 AscendC High-Level API
372+ 
373+**<term>Atlas A2 training products</term>**
374+ 
375+The current AscendC high-level API provides two programming modes. In the first mode, the Vector is the master core and the Cube is the slave core. Vector0 and Vector1 independently dispatch Matmul tasks, and the two are independent of each other.
376+ 
377+In the second mode, the Cube is the master core and the Vector is the slave core. V0 uniformly dispatches the Matmul task, and the result of this task is processed jointly by V0 and V1, usually with each processing half. Currently, if a template name ends with "_sab", the template uses the Cube as the master core. For example:
378+ 
379+```c++
380+ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_s1s2_bn2gs1_sab.h
381+ 
382+ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_sab.h
383+```
384+ 
385+For FlashAttention, using the Cube as the master core allows the Matmul tasks of V0 and V1 to reuse the left matrix, and part of the output results can be accumulated in L0C. This reduces the bandwidth dependency, so the performance is better in most scenarios.
386+ 
387+**<term>Ascend 950PR&950DT products</term>**
388+ 
389+To achieve ultimate performance, all scenarios except some deterministic computation scenarios are implemented using the AscendC low-level API. Both the high-level and low-level implementations use the Cube as the master core.
390+ 
391+### 6.2 AscendC Low-Level API
392+ 
393+**<term>Atlas A2 training products</term>**
394+ 
395+Some templates do not use the AscendC high-level API. For example:
396+ 
397+```c++
398+ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_basic.h
399+```
400+ 
401+This template more thoroughly uses the Cube as the master core and the Vector as the slave core. All Matmul tasks are fully dispatched from the Cube side, which notifies the Vector side through synchronization.
402+ 
403+**<term>Ascend 950PR&950DT products</term>**
404+ 
405+Except for some deterministic computation scenarios, all other scenarios are implemented using the low-level API, achieving ultimate memory reuse.
Aattention/flash_attn/tests/pytests/readme_en.md+478-0
@@ -0,0 +1,478 @@
1+# FlashAttention Accuracy Test Framework
2+ 
3+English | [简体中文](./readme.md)
4+ 
5+## 1. Architecture Overview
6+ 
7+```
8+test_flash_attn.py (entry point)
9+ │
10+ ├── case_loader.load_case_modules() ← loads test_cases/*.py + Cartesian product expansion
11+ ├── case_loader.resolve_case_ids() ← --case_id filtering
12+ │
13+ ▼ loop over cases
14+ ├── case_loader.normalize_params() ← completes default values, TND/PA special handling
15+ │
16+ ▼
17+ ├── orchestrator.run_case()
18+ │ │
19+ │ ├── data.flash_attn_inputs.generate() ← declarative tensor generation (Q/K/V)
20+ │ ├── data.build_flash_attn_params() ← constructs metadata/kernel kwargs
21+ │ │
22+ │ ├── primary.compute(inputs, params) ← NPU/GPU backend
23+ │ ├── golden.compute(cpu_inputs, params) ← CPU reference (tfoward)
24+ │ │
25+ │ └── compare.check_result() ← accuracy comparison
26+ │
27+ └── reporter.record() → terminal table + CSV
28+```
29+ 
30+**Core modules** (`core/`):
31+ 
32+| Module | Responsibility |
33+|------|------|
34+| `backends/base.py` | Backend abstract base class (`is_available`, `compute`, `device`) |
35+| `backends/cpu.py` | CPU Golden — calls `backends/cpu_impl.tforward` to perform the reference computation |
36+| `backends/gpu.py` | GPU — wraps the `flash_attn` library |
37+| `backends/npu.py` | NPU — calls `flash_attn_metadata` + `flash_attn`, including `print_metadata` |
38+| `data.py` | `InputSpec` declarative tensor generation + `build_flash_attn_params` parameter group construction |
39+| `case_loader.py` | Case loading (`load_case_modules`), parameter normalization (`normalize_params`), ID resolution (`resolve_case_ids`) |
40+| `orchestrator.py` | Orchestrator — chains data generation → backend computation → accuracy comparison |
41+| `reporter.py` | `Reporter` — terminal table + incremental CSV output |
42+ 
43+---
44+ 
45+## 2. Directory Structure
46+ 
47+```
48+pytests/
49+├── core/ # Core modules
50+│ ├── backends/
51+│ │ ├── base.py # Backend abstract base class
52+│ │ ├── cpu.py # CPU Golden (calls tforward)
53+│ │ ├── gpu.py # GPU (flash_attn library)
54+│ │ └── npu.py # NPU (flash_attn + print_metadata)
55+│ ├── data.py # InputSpec tensor generation + build_flash_attn_params
56+│ ├── case_loader.py # Case loading + parameter normalization
57+│ ├── orchestrator.py # Execution orchestration
58+│ └── reporter.py # Result reporting (terminal + CSV)
59+├── backends/ # Low-level implementations (called by core/backends)
60+│ ├── cpu_impl.py # tforward — CPU reference implementation
61+│ ├── gpu_impl.py # flash_attn_gpu — GPU implementation
62+│ └── npu_impl.py # NPU helper functions
63+├── test_cases/ # Test case definitions
64+│ ├── functional_stc.py # Functional STC cases
65+│ ├── performance_redline_train.py
66+│ ├── performance_redline_infer.py
67+│ ├── functional_redline_train.py
68+│ ├── functional_redline_train_tnd.py
69+│ ├── functional_redline_infer.py
70+│ ├── base.py # Base cases
71+│ └── fia_stc.py # FIA STC cases
72+├── utils/ # Utility functions
73+│ ├── compare.py # Accuracy comparison + failure distribution analysis
74+│ ├── data.py # Data generation, mask, layout conversion
75+│ ├── io.py # Tensor file read/write
76+│ ├── perf_parser.py # Performance data parsing
77+│ ├── perf_runner.py # Performance test runner
78+│ └── precision_visual.py # Accuracy heatmap visualization
79+├── tools/
80+│ └── xlsx_to_testcase.py # Excel → test case conversion
81+├── test_flash_attn.py # Main entry point
82+├── test_npu_metadata.py # Standalone metadata operator test
83+└── readme.md
84+```
85+ 
86+---
87+ 
88+## 3. Environment Dependencies
89+ 
90+### General
91+ 
92+| Dependency | Description |
93+|------|------|
94+| Python 3.8+ | — |
95+| PyTorch 2.2+ | Foundation for CPU/GPU computation |
96+| `einops` | Layout conversion (`pip install einops`) |
97+| `numpy` | Numerical computation |
98+ 
99+### NPU Mode
100+ 
101+| Dependency | Description |
102+|------|------|
103+| `TorchNPU` | PyTorch NPU extension |
104+| `cann_ops_transformer` | Provides `flash_attn` and `flash_attn_metadata` |
105+ 
106+### GPU Mode (Optional)
107+ 
108+| Dependency | Description |
109+|------|------|
110+| `flash-attn` | FlashAttention library (`pip install flash-attn --no-build-isolation`) |
111+| CUDA 12.0+ | GPU environment |
112+ 
113+### Environment Check
114+ 
115+```bash
116+# NPU
117+python -c "import torch_npu, cann_ops_transformer; print('NPU OK')"
118+ 
119+# GPU
120+python -c "import torch, flash_attn, einops; print('GPU OK')"
121+ 
122+# CPU-only
123+python -c "import torch, einops, numpy; print('CPU OK')"
124+```
125+ 
126+---
127+ 
128+## 4. Running Methods
129+ 
130+### 4.1 Accuracy Test (Default)
131+ 
132+```bash
133+# Run all cases (NPU mode)
134+python test_flash_attn.py
135+ 
136+# Specify cases
137+python test_flash_attn.py --case_id TND_05,TND_06
138+ 
139+# GPU mode
140+python test_flash_attn.py --case_id BASE_01 --use_gpu
141+ 
142+# Specify case files
143+python test_flash_attn.py --case_files functional_stc --case_id TND_05
144+```
145+ 
146+### 4.2 Metadata Debugging
147+ 
148+This mode calls only the AICPU metadata operator, prints the core partitioning result, and skips kernel execution:
149+ 
150+```bash
151+python test_flash_attn.py --case_id TND_05 --meta_only
152+```
153+ 
154+The output contains:
155+ 
156+- **Header** — `sectionNum`, `isFD`, `mBaseSize`, `s2BaseSize`
157+- **FA Metadata** — core partitioning information of the 36 AIC cores for each section (start/end of BN2/M/S2)
158+- **FD Metadata** — reduction tasks of active AIV cores (only when M_NUM > 0)
159+ 
160+```
161+==============================================================================
162+ [Metadata Header]
163+ sectionNum = 1
164+ isFD = 1
165+ mBaseSize = 128
166+ s2BaseSize = 64
167+------------------------------------------------------------------------------
168+ [Section 0] FA Metadata — AIC cores (36 cores × 16 slots, 7 fields used)
169+------------------------------------------------------------------------------
170+ Core BN2_START M_START S2_START BN2_END M_END S2_END FIRST_FD_WS_IDX
171+ AIC00 0 0 0 7 127 63 0
172+ ...
173+```
174+ 
175+### 4.3 Performance Test
176+ 
177+```bash
178+# Collect performance data in batch
179+python test_flash_attn.py --case_files performance_redline_train --perf_mode --perf_output ./perf_out
180+ 
181+# Execute cases one by one
182+python test_flash_attn.py --perf_mode --one_by_one
183+```
184+ 
185+### 4.4 Three-Way Comparison
186+ 
187+```bash
188+# Real-time CPU vs GPU vs NPU comparison
189+python test_flash_attn.py --case_id BASE_01 --compare_mode --use_gpu
190+ 
191+# Offline comparison (from dump files)
192+python test_flash_attn.py --case_id BASE_01 --compare_mode \
193+ --load_gpu_dump ./dump/BASE_01/gpu_out.txt \
194+ --load_npu_dump ./dump/BASE_01/npu_out.txt
195+```
196+ 
197+---
198+ 
199+## 5. Command-Line Parameters
200+ 
201+### Basic Parameters
202+ 
203+| Parameter | Type | Default Value | Description |
204+|------|------|--------|------|
205+| `--case_files` | str | Automatic scan of `test_cases/` | Case files, separated by commas |
206+| `--case_id` | str | `all` | Case names, separated by commas, prefix matching supported |
207+| `--device_id` | int | `0` | NPU device ID |
208+| `--use_gpu` | flag | — | Use the GPU backend |
209+| `--gpu_device` | int | `0` | GPU device ID |
210+ 
211+### Golden Persistence Parameters
212+ 
213+| Parameter | Type | Default Value | Description |
214+|------|------|--------|------|
215+| `--save_golden` | str | — | Runs the normal comparison process and saves inputs/golden to the specified directory (in .pt format) |
216+| `--load_golden` | str | — | Loads golden from the specified directory and runs only the device comparison (skips the CPU golden computation) |
217+| `--case_timeout` | int | `300` | Timeout for a single case in subprocess isolation mode (seconds) |
218+ 
219+### Debugging Parameters
220+ 
221+| Parameter | Type | Default Value | Description |
222+|------|------|--------|------|
223+| `--meta_only` | flag | — | Calls only the metadata operator and prints the core partitioning result |
224+| `--verbose_diff` | flag | — | Prints all elements that exceed the threshold when accuracy does not meet the requirement (prints only the first 10 by default) |
225+| `--visualize` | flag | — | Generates a heatmap PNG when accuracy does not meet the requirement |
226+| `--viz_dir` | str | `./viz_output` | Directory for saving heatmaps/analysis reports |
227+| `--fail_analysis` | flag | — | Outputs a multi-dimensional failure distribution analysis report when accuracy does not meet the requirement |
228+| `--compare_mode` | flag | — | Multi-backend comparison |
229+| `--graph_mode` | flag | — | Executes in graph mode |
230+ 
231+### Output Parameters
232+ 
233+| Parameter | Type | Default Value | Description |
234+|------|------|--------|------|
235+| `--result_csv` | str | `result.csv` | Path of the accuracy result CSV |
236+| `--report_interval` | int | `20` | Interval for printing intermediate statistics |
237+| `--skip_accuracy` | flag | — | Skips the accuracy comparison |
238+ 
239+### Performance Parameters
240+ 
241+| Parameter | Type | Default Value | Description |
242+|------|------|--------|------|
243+| `--perf_mode` | flag | — | Performance test mode |
244+| `--perf_runs` | int | `5` | Number of performance collection rounds |
245+| `--perf_cold_thr` | int | `16` | Cold/hot case boundary (S1) |
246+| `--perf_output` | str | `./perf_output` | Performance output directory |
247+| `--one_by_one` | flag | — | Executes cases one by one |
248+ 
249+### Offline Comparison Parameters
250+ 
251+| Parameter | Type | Default Value | Description |
252+|------|------|--------|------|
253+| `--load_gpu_dump` | str | None | Path of the GPU dump file |
254+| `--load_npu_dump` | str | None | Path of the NPU dump file |
255+ 
256+---
257+ 
258+## 6. Test Case Format
259+ 
260+Test cases are defined in the `TestCases` dictionary in `test_cases/*.py`. Each field is a list, and the framework automatically performs Cartesian product expansion.
261+ 
262+### 6.1 Field Description
263+ 
264+```python
265+TestCases = {
266+ "MY_CASE": {
267+ # ── Required ──
268+ "B": [2], # batch size (can be omitted for TND)
269+ "N1": [8], # number of query heads
270+ "D": [128], # head dim
271+ "Dtype": ["fp16"], # fp16 / bf16
272+ 
273+ # ── Layout ──
274+ "layout_q": ["BNSD"], # BNSD / BSND / TND
275+ "layout_kv": ["BNSD"], # BNSD / BSND / TND / PA_BBND / PA_BNBD
276+ "layout_out": ["BNSD"], # BNSD / BSND / TND
277+ 
278+ # ── Sequence Length ──
279+ "S1": [64], # query seq len (max_seqlen_q for TND)
280+ "S2": [1024], # KV seq len (max_seqlen_kv for TND)
281+ 
282+ # ── Optional ──
283+ "N2": [1], # number of KV heads, default = N1
284+ "DV": [128], # value head dim, default = D; when DV != D, metadata passes head_dim_v=DV through (the (192,128) combination is supported)
285+ "scale": [None], # scale factor, default 1/sqrt(D)
286+ "mask_mode":[0], # 0=no mask, 3=causal, 4=band+causal
287+ "win_left": [-1],
288+ "win_right":[-1],
289+ "return_softmax_lse": [False],
290+ 
291+ # ── Q/K/V Value Range ──
292+ "q_range": [(-5.0, 5.0)],
293+ "k_range": [(-5.0, 5.0)],
294+ "v_range": [(-5.0, 5.0)],
295+ 
296+ # ── TND-Specific ──
297+ "cu_seqlens_q": [[None]], # cumulative sequence lengths (B+1 elements)
298+ "cu_seqlens_kv": [[None]],
299+ "seqused_q": [[None]], # actual sequence length per batch
300+ "seqused_kv": [[None]],
301+ },
302+}
303+```
304+ 
305+### 6.2 TND Cases
306+ 
307+```python
308+"TND_05": {**BASE,
309+ "layout_q": ["TND"], "layout_kv": ["TND"], "layout_out": ["TND"],
310+ "cu_seqlens_q": [[0, 64, 128, 192, 256, 320, 384, 448, 512]],
311+ "cu_seqlens_kv": [[0, 1024, 2048, 3072, 4096, 5120, 6144, 7168, 8192]]},
312+```
313+ 
314+`normalize_params` automatically derives `seqused` (difference) and `B` (len-1) from `cu_seqlens`.
315+ 
316+### 6.3 Paged Attention Cases
317+ 
318+```python
319+"PA_07": {**BASE,
320+ "layout_kv": ["PA_BBND", "PA_BNBD"],
321+ "block_size": [128]},
322+```
323+ 
324+`normalize_params` automatically generates `block_table` from `seqused_kv` + `block_size`.
325+ 
326+---
327+ 
328+## 7. Execution Flow Details
329+ 
330+### 7.1 Case Loading
331+ 
332+```
333+test_cases/functional_stc.py::TestCases
334+ │
335+ ▼ load_case_modules()
336+ Cartesian product expansion: each parameter combination → an independent case
337+ Naming: "test_cases.functional_stc/TND_05"
338+ Multiple combinations: "test_cases.functional_stc/GS1_14_0", "..._1", ...
339+ │
340+ ▼ resolve_case_ids()
341+ --case_id TND_05 → matches all cases containing "/TND_05"
342+```
343+ 
344+### 7.2 Parameter Normalization
345+ 
346+Auto-completion logic of `normalize_params()`:
347+ 
348+| Condition | Auto-Completion |
349+|------|----------|
350+| `N2` is missing | `N2 = N1` |
351+| `S2` is missing | `S2 = S1` |
352+| `DV` is missing | `DV = D` |
353+| TND + `seqused_q` missing | Derived from the difference of `cu_seqlens_q` |
354+| TND + `cu_seqlens_kv` missing | Copied from `cu_seqlens_q` |
355+| PA + `block_table` missing | Automatically generated from `seqused_kv` + `block_size` |
356+ 
357+### 7.3 Tensor Generation
358+ 
359+`InputSpec` declaratively defines the shape, dtype, and generation method of Q/K/V:
360+ 
361+```python
362+flash_attn_inputs = (
363+ InputSpec("flash_attn")
364+ .tensor("q", layout={"BNSD": "B,N1,S1,D", "BSND": "B,S1,N1,D", "TND": "total_s1,N1,D"})
365+ .tensor("k", layout={...}, layout_override="layout_kv")
366+ .tensor("v", layout={...}, layout_override="layout_kv")
367+)
368+```
369+ 
370+Special tokens: `total_s1` = `cu_seqlens_q[-1]`, `total_s2` = `cu_seqlens_kv[-1]`.
371+ 
372+### 7.4 Parameter Group Construction
373+ 
374+`build_flash_attn_params()` returns `(meta_kwargs, kernel_kwargs, out_layout)`:
375+ 
376+- **meta_kwargs** → passed to `flash_attn_metadata()`
377+ - `cu_seqlens_q/kv`, `seqused_q/kv`, `max_seqlen_q/kv`, `num_heads_q/kv`, `head_dim`, `batch_size`, `mask_mode`, `layout_*`
378+- **kernel_kwargs** → passed to `flash_attn()`
379+ - `softmax_scale`, `cu_seqlens_q/kv`, `seqused_q/kv`, `max_seqlen_q/kv` (actual values are passed for TND, otherwise -1), `mask_mode`, `layout_*`
380+ 
381+### 7.5 Accuracy Comparison
382+ 
383+Decision criteria of `check_result()`:
384+ 
385+```
386+Per-element: diff ≤ max(|golden| × 0.5%, 0.000025)
387+Overall: proportion of elements exceeding the threshold ≤ 0.5%
388+```
389+ 
390+Output: `PASS` / `FAIL` + MaxAbsErr + MeanAbsErr + MaxRelErr + FailRatio.
391+ 
392+---
393+ 
394+## 8. Layout Support Matrix
395+ 
396+| Layout | CPU | GPU | NPU | Description |
397+|--------|-----|-----|-----|------|
398+| **BNSD** | ✓ | ✓ | ✓ | (B, N, S, D) default |
399+| **BSND** | ✓ | ✓ | ✓ | (B, S, N, D) |
400+| **TND** | ✓ | ✓ | ✓ | (total_tokens, N, D) variable-length sequences |
401+| **PA_BBND** | ✗ | ✓ | ✓ | (num_blocks, block_size, N, D) Paged KV |
402+| **PA_BNBD** | ✗ | ✓ | ✓ | (num_blocks, N, block_size, D) Paged KV |
403+ 
404+---
405+ 
406+## 9. mask_mode
407+ 
408+| Value | Name | Description |
409+|----|------|------|
410+| `0` | NO_MASK | Full attention |
411+| `3` | RIGHT_DOWN_CAUSAL | Right-down aligned causal mask |
412+| `4` | BAND_CAUSAL | BAND + CAUSAL combination, `win_left`/`win_right` take effect |
413+ 
414+---
415+ 
416+## 10. File Responsibility Quick Reference
417+ 
418+| File | Change Frequency | Typical Changes |
419+|------|----------|----------|
420+| `test_cases/*.py` | **Frequent** | Add/modify test cases |
421+| `core/data.py` | Occasional | Parameter group construction, tensor generation rules |
422+| `core/backends/npu.py` | Occasional | NPU invocation, `print_metadata` |
423+| `test_flash_attn.py` | Occasional | CLI parameters, feature entry points |
424+| `core/case_loader.py` | Rare | Parameter normalization logic |
425+| `core/orchestrator.py` | Rare | Execution flow orchestration |
426+| `core/reporter.py` | Rare | Report format |
427+| `utils/compare.py` | Rare | Accuracy comparison, failure analysis |
428+| `utils/data.py` | Rare | Mask generation, layout conversion |
429+| `backends/cpu_impl.py` | Normally unchanged | CPU golden (tfoward) |
430+| `backends/gpu_impl.py` | Rare | GPU implementation |
431+| `utils/perf_*.py` | Rare | Performance collection and parsing |
432+ 
433+---
434+ 
435+## 11. Frequently Asked Questions
436+ 
437+### Q1: `ModuleNotFoundError: No module named 'cann_ops_transformer'`
438+ 
439+Install `cann_ops_transformer` for NPU mode, or use GPU mode instead:
440+ 
441+```bash
442+python test_flash_attn.py --case_id BASE_01 --use_gpu
443+```
444+ 
445+### Q2: Abnormal Accuracy for TND Cases
446+ 
447+Check the format of `cu_seqlens_q` (it must include a leading 0, and the length = batch + 1):
448+ 
449+```python
450+"cu_seqlens_q": [[0, 64, 128, 192]], # 3 batches, 64 tokens each
451+```
452+ 
453+### Q3: TND Case Kernel Hang
454+ 
455+Confirm that `max_seqlen_q`/`max_seqlen_kv` are passed in correctly. The kernel requires the actual max values under the TND layout:
456+ 
457+```bash
458+python test_flash_attn.py --case_id TND_05 --meta_only # check metadata first
459+```
460+ 
461+### Q4: PA Case `block_size` Error
462+ 
463+Different backends have different constraints on `block_size`. Refer to the official documentation of the corresponding backend to confirm the specific restrictions:
464+ 
465+- For GPU, refer to the `flash_attn` library documentation
466+- For NPU, refer to the `flash_attn` operator documentation
467+ 
468+### Q5: How to View the Actual Parameters of a Case
469+ 
470+```bash
471+python -c "
472+from core.case_loader import load_case_modules, normalize_params
473+cases = load_case_modules(['test_cases.functional_stc'])
474+for name, raw in cases.items():
475+ if 'TND_05' in name:
476+ print(normalize_params(raw))
477+"
478+```
Aattention/incre_flash_attention/README_en.md+145-0
@@ -0,0 +1,145 @@
1+# IncreFlashAttention
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+## Description
17+ 
18+- Operator function: For auto-regressive language models, the inference input length keeps increasing as new words are generated. On the basis of the original full inference, **incremental inference is implemented**: the S axis of query is fixed to 1, and key and value are stored in the kvCache, where the state information inferred previously is accumulated. The actual length of the S axis may differ across batches, and the input data is fixed-length data after padding. The following features are supported: **quantization, position encoding, page attention, kvCache dequantization, and KV left padding features.**
19+ 
20+ Compared with the FlashAttention operator in the full scenario ([PromptFlashAttention](../prompt_flash_attention/README_en.md)), the incremental inference process is not completely equivalent to normal full inference, but the precision of incremental inference is not significantly degraded.
21+ 
22+ **Note:**
23+ 
24+ kvCache is a common technique for optimizing the inference performance of large models. During sampling, the Transformer model performs inference with the given prompt/context as the initial input (which can be processed in parallel), and then generates additional tokens one by one to continuously refine the generated sequence (which reflects the auto-regressive nature of the model). During sampling, the Transformer performs self-attention operations. To do so, it needs to extract key-value (KV) vectors for each item in the current sequence (whether a prompt/context or a generated token). These vectors are stored in a matrix, which is commonly referred to as the KV cache (kvCache).
25+ 
26+- Calculation formula:
27+ 
28+ Self-attention builds an attention model by using the relationships within the input sample itself. The principle is as follows: assume an input sample sequence $x$ of length $n$, where each element of $x$ is a $d$-dimensional vector. Each $d$-dimensional vector can be regarded as a token embedding. Passing such a sequence through 3 weight matrices yields 3 matrices of dimension $n*d$.
29+ 
30+ The calculation formula of self-attention is generally defined as follows, where $Q$, $K$, and $V$ are important attribute elements of the input sample. They are obtained by spatial transformation of the input sample and can be unified into one feature space.
31+ 
32+ $$
33+ Attention(Q,K,V)=Score(Q,K)V
34+ $$
35+ 
36+ In this operator, the Score function uses the Softmax function. The calculation formula of self-attention is:
37+ 
38+ $$
39+ Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{d}})V
40+ $$
41+ 
42+ The product of $Q$ and $K^T$ represents the attention of the input $x$. To prevent this value from becoming too large, it is usually scaled by dividing by the square root of $d$, and softmax normalization is performed on each row. After being multiplied by $V$, an $n*d$ matrix is obtained.
43+ 
44+## Parameter Description
45+ 
46+<div style="overflow-x: auto;">
47+<table style="undefined;table-layout: fixed; width: 930px"><colgroup>
48+<col style="width: 150px">
49+<col style="width: 150px">
50+<col style="width: 250px">
51+<col style="width: 230px">
52+<col style="width: 150px">
53+</colgroup>
54+<thead>
55+ <tr>
56+ <th style="text-align: center;">Parameter Name</th>
57+ <th style="text-align: center;">Input/Output/Attribute</th>
58+ <th style="text-align: center;">Description</th>
59+ <th style="text-align: center;">Data Type</th>
60+ <th style="text-align: center;">Data Format</th>
61+ </tr></thead>
62+<tbody>
63+ <tr>
64+ <td>query</td>
65+ <td >Input</td>
66+ <td >Input Q in the formula.</td>
67+ <td >FLOAT, FLOAT16</td>
68+ <td >ND</td>
69+ </tr>
70+ <tr>
71+ <td>key</td>
72+ <td >Input</td>
73+ <td >Input K in the formula.</td>
74+ <td >FLOAT, INT8, FLOAT16</td>
75+ <td >ND</td>
76+ </tr>
77+ <tr>
78+ <td>value</td>
79+ <td >Input</td>
80+ <td >Input V in the formula.</td>
81+ <td >FLOAT, INT8, FLOAT16</td>
82+ <td >ND</td>
83+ </tr>
84+ <tr>
85+ <td>scaleValue</td>
86+ <td >Attribute</td>
87+ <td >The reciprocal of the square root of d in the formula.</td>
88+ <td >DOUBLE</td>
89+ <td >-</td>
90+ </tr>
91+ <tr>
92+ <td>attentionOut</td>
93+ <td >Output</td>
94+ <td >The output in the formula.</td>
95+ <td >FLOAT, INT8, FLOAT16</td>
96+ <td >ND</td>
97+ </tr>
98+</tbody>
99+</table></div>
100+ 
101+## Constraints
102+ 
103+- The shapes of query and attentionOut must be exactly the same, and the shapes of the corresponding tensors in key and value must be exactly the same.
104+- In the non-contiguous scenario, the number of tensors in the tensorlist of key and value equals the B of query (due to the tensorlist limitation, B must be less than or equal to 256 in the non-contiguous scenario). The shapes must be exactly the same except for S, and the batch can only be 1.
105+- The N in query equals numHeads, the N in key and value equals numKeyValueHeads, and numHeads is a multiple of numKeyValueHeads.
106+- <term>Atlas A2 products</term>, <term>Ascend 950PR&950DT products</term>:
107+ - The B axis is less than or equal to 65536, the N axis is less than or equal to 256, and the D axis is less than or equal to 512.
108+- Atlas inference products:
109+ - The B axis is less than or equal to 256, the N axis is less than or equal to 256, the D axis is less than or equal to 512, and the S axes of key and value are less than or equal to 65536.
110+ - The data types of query, key, value, and attentionOut support only FLOAT16.
111+ - When the data layout format is BNSD, the ratio of numHeads to numKeyValueHeads must not be greater than 8. In other cases, only the value 0 is supported.
112+- Comprehensive limitations on the number of INT8 quantization-related input parameters and the input/output data formats:
113+ - In the scenario where the inputs query, key, and value are FLOAT16 and the output is INT8: the input parameter quantScale2 is required, quantOffset2 is optional, and the parameters dequantScale1, quantScale1, and dequantScale2 must not be passed in (that is, they must be nullptr).
114+ 
115+- The data type of pseShift must be the same as the data type of query.
116+ 
117+- Parameter constraints of antiquantScale and antiquantOffset:
118+ - per-channel mode: the shapes of the two parameters support (2, N, 1, D), (2, N, D), and (2, H), where N is numKeyValueHeads. The parameter data type is the same as the data type of query.
119+ - per-tensor mode: the shapes of both parameters are (2), and the data type is the same as the data type of query.
120+ 
121+- The input parameters quantScale2 and quantOffset2 support two formats (per-tensor/per-channel) and two data types (FLOAT32/BFLOAT16). If quantOffset2 is passed in, its type and shape information must be the same as those of quantScale2. When the input is BFLOAT16, both FLOAT32 and BFLOAT16 are supported; otherwise, only FLOAT32 is supported. In per-channel format, when the output layout is BSH, the product of all dimensions of quantScale2 must equal H; for other layouts, the product must equal N*D. (It is recommended that when the output layout is BSH, the shape [1,1,H] or [H] is passed in for quantScale2; when the output is BNSD, it is recommended to pass in [1,N,1,D] or [N,D]; when the output is BSND, it is recommended to pass in [1,1,N,D] or [N,D].)
122+ 
123+- page attention scenario:
124+ 
125+ - The prerequisite for enabling page attention is that blockTable exists and is valid, and key and value are arranged in a contiguous memory area according to the indices in blockTable. The dtypes of key and value support FLOAT16/BFLOAT16/INT8. In this scenario, the inputLayout parameters of key and value are invalid.
126+ - blockSize is a user-defined parameter whose value affects the performance of page attention. In the scenario where page attention is enabled, a non-0 value must be passed in for blockSize, and blocksize must not exceed 512. When the input types of key and value are FLOAT16/BFLOAT16, the value must be aligned to 16; when the input types of key and value are INT8, the value must be aligned to 32. A value of 128 is recommended. In general, page attention can improve throughput but causes performance degradation.
127+ - In the page attention scenario, when the inputLayout of query is BNSD, the kvCache layout supports two formats: (blocknum, blocksize, H) and (blocknum, KV_N, blocksize, D). When the inputLayout of query is BSH or BSND, the kvCache layout supports only the format (blocknum, blocksize, H). blocknum cannot be less than the sum of the number of blocks of each batch calculated based on actualSeqLengthsKv and blockSize. In addition, the shapes of key and value must be the same.
128+ - In the page attention scenario, the performance with the kvCache layout (blocknum, KV_N, blocksize, D) is usually better than that with the kvCache layout (blocknum, blocksize, H). It is recommended to preferentially select the format (blocknum, KV_N, blocksize, D).
129+ - In the scenario where page attention is enabled, when the input kvCache layout format is (blocknum, blocksize, H) and numKvHeads * headDim exceeds 64k, the request is intercepted and an error is reported due to hardware instruction constraints. This problem can be solved by enabling GQA (reducing numKvHeads) or changing the kvCache layout format to (blocknum, numKvHeads, blocksize, D).
130+ - In the page attention scenario, the input actualSeqLengths must be passed in.
131+ - In the page attention scenario, blockTable must be two-dimensional. The length of the first dimension must equal B, and the length of the second dimension cannot be less than maxBlockNumPerSeq (maxBlockNumPerSeq is the number of blocks corresponding to the maximum actualSeqLengthsKv in each batch).
132+ - In the scenario where page attention is enabled, the input S must be greater than or equal to maxBlockNumPerSeq * blockSize in the following scenarios:
133+ - Attention Mask is enabled, for example, the mask shape is (B, 1, 1, S).
134+ - pseShift is enabled, for example, the pseShift shape is (B, N, 1, S).
135+- kv left padding scenario:
136+ - The calculation formula of the start of the kvCache transfer is: Smax - kvPaddingSize - actualSeqLengths; the calculation formula of the end of the kvCache transfer is: Smax - kvPaddingSize. When the start or end of the kvCache transfer is less than 0, the returned data result is all 0.
137+ - kvPaddingSize is set to 0 when it is less than 0.
138+ - This feature must be enabled together with the actualSeqLengths parameter; otherwise, the kv right padding scenario is used by default.
139+ - When this feature is enabled together with the Attention Mask parameter, ensure that the meaning of Attention Mask is correct, that is, invalid data can be correctly masked. Otherwise, precision problems are introduced.
140+ 
141+## Calling Instructions
142+ 
143+| Calling Method | Sample Code | Description |
144+| --------- | ------------------------------------------------------------ | ------------------------------------------------------------ |
145+| aclnn API | [test_aclnn_IncreFlashAttentionV4](./examples/test_aclnn_incre_flash_attention_v4.cpp) | Calls the IncreFlashAttentionV4 operator by using [aclnnIncreFlashAttentionV4](docs/aclnnIncreFlashAttentionV4_en.md). |
Aattention/mla_prolog/docs/MlaProlog_design_en.md+428-0
@@ -0,0 +1,428 @@
1+ 
2+# MlaProlog Operator Design Introduction
3+ 
4+In recent years, with the rapid development of artificial intelligence technology, large models have achieved significant breakthroughs in natural language processing, computer vision, multimodal interaction, and other fields. However, the complexity and scale of models have grown exponentially, placing higher demands on underlying computing frameworks and hardware computing power.
5+When handling ultra-large-scale parameter models, traditional matrix operations (such as GEMM) and attention mechanism implementations gradually expose problems such as low computational efficiency, high memory usage, and insufficient parallel scalability, which restrict the real-time performance and resource utilization of model training and inference.
6+ 
7+Against this background, the DeepSeek team proposed the MLA operator requirement, which calls for coordinated optimization of algorithms and hardware. Through operator fusion, the number of operator calls and the overhead time are reduced, which lowers inference latency and reduces hardware resource consumption while maintaining model accuracy. The MlaProlog operator is exactly the fused preprocessing part of the MLA operator.
8+ 
9+## Implementation Principle
10+ 
11+Figure 1 Computation flow diagram:
12+ 
13+![MlaProlog computation flow diagram](../../../docs/en/figures/mla_prolog_compute_flow.png)
14+ 
15+The implementation follows the computation flow defined by Multi-Head Latent Attention. The overall computation flow is as follows:
16+ 
17+1. The input sequence x passes through the downsampling matrix $W^{DQ}$ for rank-reduction transformation and then undergoes normalization to obtain $c^Q$
18+ 
19+2. The $W^{UQ}$ and $W^{QR}$ matrices are concatenated to perform the rank-elevation and mapping computation on $c^Q$. The computation result is then split for Qn normalization and Rope position encoding respectively, finally obtaining the Query and Query Rope
20+ 
21+3. The $W^{DKV}$ and $W^{KR}$ matrices are concatenated to perform the rank-reduction and mapping computation on the input sequence x. The computation result is then split for Rope position encoding and $C^{KV}$ normalization respectively, finally obtaining the KR Cache and KV Cache
22+ 
23+For the specific computation formulas, see the [Complete Computation Formulas](#complete-computation-formulas) section
24+ 
25+## Data Splitting Design
26+ 
27+The hardware buffer size is limited, while the data volume to be computed is huge, so the computation cannot be completed at one time and tiling is required. Different shapes lead to different splitting axes of the operator, and the splitting axis of the operator affects the function and performance of the template. Consider the following points:
28+ 
29+a. Use all cores to prevent some cores from being idle.
30+ 
31+b. Keep the computation amount assigned to each core relatively even, which avoids the situation where some cores compute too much data while other cores remain idle.
32+ 
33+c. Ensure that the data volume processed between the AIC and the AIV matches their corresponding computing power, which avoids long idle periods of the AIC or the AIV.
34+ 
35+The MlaProlog operator has multiple Matmul operations:
36+ 
37+1. MatmulCq: Merges BS into one axis as the M axis, uses HeadSizeCq as the N axis, and uses HeadSizeX as the K axis. HeadSizeCq is split across cores based on the number of cube cores, and each core is assigned a certain number of sub-blocks
38+ 
39+2. MatmulCkvKr: Merges BS into one axis as the M axis, uses HeadSizeCkvKr as the N axis, and uses HeadSizeX as the K axis. HeadSizeCkvKr is split across cores into 64 blocks
40+ 
41+3. MatmulQcQr: Merges BS into one axis as the M axis, uses HeadSizeQcQr as the N axis, and uses HeadSizeCq as the K axis. HeadSizeQcQr is split across cores based on the number of cube cores, and each core is assigned a certain number of sub-blocks
42+ 
43+4. MatmulQn: Merges BS into one axis as the M axis, uses HeadSizeCkv as the N axis, and uses dimHeadSizeQc as the K axis
44+ 
45+## Pipeline Design
46+ 
47+The MlaProlog fused operator contains both Vector computation and Cube computation. Computations on the Vector side and the Cube side depend on each other, so CV pipeline design is required. Otherwise, the Cube side and the Vector side are likely to form a serial pipeline, which fails to achieve parallel computation and prevents the fused operator from achieving optimal performance:
48+ 
49+Figure 2 Pipeline control diagram:
50+ 
51+![MlaProlog pipeline control diagram](../../../docs/en/figures/mla_prolog_pipeline_control.png)
52+ 
53+1. During MatmulCq computation, Hcq is split across cores, and a single core does not compute a complete token. Therefore, synchronization control between the AIC and the AIV (SYNC_MMCQ_NOMRCQ) is required before RmsNormCq computation
54+ 
55+2. During MatmulCkvKr computation, Hcq is also split across cores, and a single core does not compute a complete token. Therefore, synchronization control between the AIC and the AIV (SYNC_MMCKVKR_NOMRCKV) is also required before RmsNormCkv computation
56+ 
57+3. Since RmsNormCq is split by rows onto different cores, synchronization control between the AIV and the AIC (SYNC_NOMRCQ_MMQCQR) is also required during MatmulQcQr computation
58+ 
59+4. Since the core-splitting strategy differs from that of RopeQr, synchronization control between the AIC and the AIV (SYNC_MMQCQR_ROPEQR) is also required before RopeQr computation
60+ 
61+5. Since MatmulQcQr and MatmulQn use different core-splitting strategies and MatmulQn depends on the output of MatmulQcQr, all-core synchronization on the cube side (SYNC_ALL_CUBE) is required
62+ 
63+6. Before Vector core computation, all-core synchronization on the vector side (SYNC_ALL_VECTOR) is required to ensure pipelined data movement
64+ 
65+## TilingKey Partitioning
66+ 
67+TilingKey is of the uint64 type. Each template parameter corresponds to one or more binary bits in TilingKey. The implementation details are as follows:
68+ 
69+|Binary bits|Variable name|Description|Parameter list|
70+|-------|------|----|-------|
71+|0-3|CACHE_MODE|Storage format of KVCache|0-BNSD (reserved), 1-PA_BSND, 2-PA_NZ|
72+|4-5|SCENARIO|Input scenario|0-FP16 (reserved), 1-non-quantization scenario, 2-quantization scenario|
73+|6-9|QUANT_MODE|Quantization scenario|0-MMQcQr quantization, 1-MMQcQr quantization + KVCache quantization|
74+|10|ENABLE_DEQUANT_OPTIONAL|Dequantization is enabled. It cannot be used together with ENABLE_DEQUANT_OPTIONAL|0-disabled, 1-enabled|
75+|11|ENABLE_GROUP_COMPUTE_OPTIONAL|Computing power grouping optimization for quantization. It cannot be used together with ENABLE_DEQUANT_OPTIONAL|0-disabled, 1-enabled|
76+|12-13|EMPTY_TENSOR_MODE|Empty tensor scenario, used when an input tensor has a dimension of 0|0-no empty tensor, 1-KVCache empty and KRCache empty, 2-Query empty|
77+ 
78+## Main Flow
79+ 
80+```c
81+// Pseudocode for single-core computation
82+void Process() {
83+ loops = compute_step_this_core(); // Compute the number of loops for the current core
84+ //
85+ for (i = 0; i < loops; i++) {
86+ updateCurrentStepSize(i, stepSize, allTokenSize); // Update the M-axis size computed in the current round
87+ 
88+ if ASCEND_IS_AIC {
89+ 
90+ MatmulCq(tokenXOffset, weightDqOffset, cqResOffset);
91+ CrossCoreSetFlag<0x2, PIPE_FIX>(SYNC_MMCQ_NORMSCQ_FLG); // Cube-vector synchronization
92+ 
93+ MatmulCkvKr(tokenXOffset, weightDkvKrOffset, ckvKrResOffset);
94+ CrossCoreSetFlag<0x2, PIPE_FIX>(SYNC_MMCKVKR_NORMROPE_FLG); // Cube-vector synchronization
95+ 
96+ CrossCoreWaitFlag(SYNC_MMCQ_NORMSCQ_FLG); // MatmulQcQr depends on the output of RmsNormCq, so CV inter-core synchronization needs to be inserted
97+ 
98+ MatmulQcQr(weightUqQrOffset, qcQrResOffset);
99+ CrossCoreSetFlag<0x2, PIPE_FIX>(SYNC_MMQCQR_ROPEQR_FLG); // Cube-vector synchronization
100+ 
101+ 
102+ // Since MatmulQn and MatmulQcQr use different core-splitting strategies, and MatmulQn depends on the output of MatmulQcQr,
103+ // MatmulQn can start only after MatmulQcQr completes on all cube cores
104+ CrossCoreSetFlag<0x0, PIPE_FIX>(SYNC_ALL_CUBE_FLG);
105+ CrossCoreWaitFlag(SYNC_ALL_CUBE_FLG);
106+ 
107+ MatmulQn(qcOffset, weightUkOffset, qnResOffset, mmQnLoopTime); // The result of MatmulQn is output directly to queryOut, and qnOffset needs to be offset along the Batch axis
108+ 
109+ }
110+ 
111+ if ASCEND_IS_AIV {
112+ 
113+ GetSinCos(tokenIndex);
114+ 
115+ CrossCoreWaitFlag(SYNC_MMCQ_NORMSCQ_FLG); // wait MatmulCq
116+ 
117+ CrossCoreSetFlag<0x0, PIPE_MTE3>(SYNC_ALL_VECTOR_FLG);
118+ CrossCoreWaitFlag(SYNC_ALL_VECTOR_FLG);
119+ 
120+ RmsNormCq(rmsNormCqOffset);
121+ 
122+ // Since RmsNormCq and MatmulQcQr use different core-splitting strategies, MatmulQcQr can start only after RmsNormCq completes on all vector cores
123+ // MatmulQcQr starts only after RmsNormCq completes on all vector cores
124+ CrossCoreSetFlag<0x0, PIPE_MTE3>(SYNC_ALL_VECTOR_FLG);
125+ CrossCoreWaitFlag(SYNC_ALL_VECTOR_FLG);
126+ 
127+ CrossCoreSetFlag<0x2, PIPE_MTE3>(SYNC_MMCQ_NORMSCQ_FLG); // Ensures that MatmulQcQr waits for RmsNormCq
128+ 
129+ 
130+ CrossCoreWaitFlag(SYNC_MMCKVKR_NORMROPE_FLG); // wait MatmulCkvKr
131+ 
132+ CrossCoreSetFlag<0x0, PIPE_MTE3>(SYNC_ALL_VECTOR_FLG);
133+ CrossCoreWaitFlag(SYNC_ALL_VECTOR_FLG);
134+ RmsNormRopeScatterCkvKr(tokenIndex, rmsNormCkvOffset, ropeKrOffset);
135+ 
136+ CrossCoreWaitFlag(SYNC_MMQCQR_ROPEQR_FLG); // wait MatmulQcQr
137+ 
138+ CrossCoreSetFlag<0x0, PIPE_MTE3>(SYNC_ALL_VECTOR_FLG);
139+ CrossCoreWaitFlag(SYNC_ALL_VECTOR_FLG);
140+ RopeQr(ropeQrOffset, ropeQrResOffset);
141+ }
142+ }
143+}
144+```
145+ 
146+## Complete Computation Formulas
147+ 
148+Combined with the flow diagrams in the preceding sections, the complete computation formulas of the entire MlaProlog are as follows:
149+ 
150+$$
151+q^N = RmsNorm(x \cdot W^{DQ}) \cdot W^{UQ} \cdot W^{UK} \tag{a}
152+$$
153+ 
154+$$
155+q^R = ROPE(RmsNorm(x \cdot W^{DQ}) \cdot W^{QR}) \tag{b}
156+$$
157+ 
158+$$
159+k^R = ROPE(x \cdot W^{KR}) \tag{c}
160+$$
161+ 
162+$$
163+c^{KV} = RmsNorm(x \cdot W^{DKV}) \tag{d}
164+$$
165+ 
166+The complete computation flow can be decomposed into the following basic computation units.
167+ 
168+### MatmulCq
169+ 
170+The input $x$ is multiplied by the Query downsampling matrix $W^{DQ}$ to perform the downsampling operation and obtain the compressed Query matrix $c^Q$.
171+$$
172+c^Q = x \cdot W^{DQ} \tag{1}
173+$$
174+ 
175+The matrix multiplication modules involved in this section (and subsequent sections) are implemented using the Matmul high-level API in the AscendC Kernel API. For the usage of related APIs, refer to the official website development guide [Operator Implementation -> Matrix Programming (High-level API)](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC3alpha003/devguide/opdevg/ascendcopdevg/atlas_ascendc_10_0041.html).
176+ 
177+### RMSNormCq
178+ 
179+RMSNorm (root mean square normalization) is performed row by row on the compressed $Q$ matrix. The RMSNorm operation requires two hyperparameters, $\gamma$ and $\epsilon$, which correspond to $rmsnormGammaCq$ and $rmsnormEpsilonCq$ in the API document.
180+$$
181+c_{norm}^Q = RMSNorm(c^Q) \tag{2}
182+$$
183+The computation formula of RMSNorm is as follows:
184+$$
185+RMSNorm(x) = \gamma \cdot \frac{x_i}{RMS(x)} \tag{3}
186+$$
187+$$
188+RMS(x) = \sqrt{\frac{1}{N} \sum_{i=1}^{N} x_i^2 + \epsilon} \tag{4}
189+$$
190+ 
191+### MatmulCkvKr
192+ 
193+The input $x$ is multiplied by the Key/Value downsampling matrix (Key and Value share the same downsampling matrix) to obtain the compressed Key/Value matrix $c^{KV}$.
194+$$
195+c^{KV} = x \cdot W^{DKV} \tag{5}
196+$$
197+The input $x$ is multiplied by the $W^{KR}$ matrix to obtain the $K^R$ matrix used for position-encoding computation.
198+$$
199+k^R = x \cdot W^{KR} \tag{6}
200+$$
201+Based on the properties of matrix multiplication, the $W^{DKV}$ and $W^{KR}$ matrices can be horizontally concatenated into one matrix $[W^{DKV}|W^{KR}]$ for computation. This concatenated matrix corresponds to the $weightDkvKr$ parameter in the API document.
202+$$
203+c^{KV}k^R = x \cdot [W^{DKV}|W^{KR}] = [x \cdot W^{DKV}|x \cdot W^{KR}] \tag{7}
204+$$
205+ 
206+### RMSNormCkv
207+ 
208+RMSNorm (root mean square normalization) is performed row by row on the compressed $KV$ matrix. The RMSNorm operation requires two hyperparameters, $\gamma$ and $\epsilon$, which correspond to $rmsnormGammaCkv$ and $rmsnormEpsilonCkv$ in the API document.
209+$$
210+c_{norm}^{KV} = RMSNorm(c^{KV}) \tag{8}
211+$$
212+For the computation of RMSNorm, refer to formulas (3)-(4).
213+ 
214+### PostQuant/Dequant
215+ 
216+- Basic concepts
217+ - Quantization: generally refers to using low-bit-width data (usually INT8) instead of high-bit-width data (usually floating-point numbers) to complete tasks such as storage, transmission, and computation. Sometimes it specifically refers to the computation process of converting floating-point numbers into integers.
218+ - Dequantization: restores low-bit-width data into high-bit-width data.
219+- The basic quantization formula is:
220+ $$
221+ x\_float≈round(x\_float/scale)∗scale
222+ $$
223+ Here, PostQuant corresponds to $round(x\_float/scale)$, which scales each value in the data by scale and then performs rounding. Multiplying by the scale parameter corresponds to the dequantization operation.
224+ 
225+Two quantization methods are currently supported:
226+ 
227+ - Dynamic quantization: computes the scale based on the input and uses smooth_scale to smooth the input. The corresponding formula is
228+ $$
229+ scale=max(abs(x\_float∗smooth\_scale))/127
230+ $$
231+ smooth means making the data "smoother" so that the data is easier to quantize, as shown in the following figure:
232+ 
233+ ![smooth concept](../../../docs/en/figures/smooth_concept.png)
234+ 
235+ - Static quantization: currently only Perchannel quantization is supported. This quantization quantizes by column, using one scale to quantize each column of the input tensor.
236+ 
237+### MatmulQcQr
238+ 
239+The normalized $C^q$ matrix is multiplied by the Query upsampling matrix $W^{UQ}$ to obtain the $q^C$ matrix.
240+$$
241+q^C = c_{norm}^Q \cdot W^{UQ} \tag{9}
242+$$
243+The normalized $C^q$ matrix is multiplied by the $W^{QR}$ matrix to obtain the $q^R$ matrix used for position-encoding computation.
244+$$
245+q^R = c_{norm}^Q \cdot W^{QR} \tag{10}
246+$$
247+Similarly, based on the properties of matrix multiplication, the $W^{UQ}$ and $W^{QR}$ matrices can be horizontally concatenated into one matrix $[W^{UQ}|W^{QR}]$ for computation. This concatenated matrix corresponds to the $weightDkvKr$ parameter in the API document.
248+$$
249+q^Cq^R = c_{norm}^Q \cdot [W^{UQ}|W^{QR}] = [c_{norm}^Q \cdot W^{UQ}|c_{norm}^Q \cdot W^{QR}] \tag{11}
250+$$
251+ 
252+### RotaryPosEmb
253+ 
254+Rotary Position Embedding (RoPE) is a position-encoding algorithm proposed in the paper [Roformer](https://arxiv.org/abs/2104.09864). Its original computation formula is as follows:
255+$$
256+f_{\{q,k\}}(x_m, m) = R_{\Theta,m}^{d}W_{\{q,k\}}x_m \tag{12}
257+$$
258+ 
259+$$
260+R_{\Theta,m}^{d} =
261+\left(
262+\begin{matrix}
263+\cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\
264+\sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\
265+0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0 & 0 \\
266+0 & 0 & \sin m\theta_2 & \cos m\theta_2 & \cdots & 0 & 0 \\
267+\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\
268+0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2} & -\sin m\theta_{d/2} \\
269+0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2} & \cos m\theta_{d/2} \\
270+\end{matrix}
271+\right)
272+\tag{13}
273+$$
274+$$
275+\Theta = \{ \theta_i = 10000^{-2(i-1)/d}, i \in [1,2,...,d/2] \}
276+$$
277+ 
278+Here, $f_{\{q,k\}}(x_m, m)$ represents the Query/Key vector after the word vector $x_m$ corresponding to the $m$-th token integrates the position information $m$. The Query and Key vectors use the same computation formula. $d$ is the dimension of the word vector.
279+ 
280+After expansion and transformation, formula (12) can be converted into the form of formula (14), where $\otimes$ denotes the element-wise cross multiplication.
281+ 
282+$$
283+ROPE(x) = R_{\Theta,m}^{d} x =
284+\left(
285+\begin{matrix}
286+x_0 \\
287+x_1 \\
288+x_2 \\
289+x_3 \\
290+\vdots \\
291+x_{d-2} \\
292+x_{d-1}
293+\end{matrix}
294+\right)
295+\otimes
296+\left(
297+\begin{matrix}
298+\cos m\theta_0 \\
299+\cos m\theta_0 \\
300+\cos m\theta_1 \\
301+\cos m\theta_1 \\
302+\vdots \\
303+\cos m\theta_{d/2-1} \\
304+\cos m\theta_{d/2-1}
305+\end{matrix}
306+\right)
307++
308+\left(
309+\begin{matrix}
310+-x_1 \\
311+x_0 \\
312+-x_3 \\
313+x_2 \\
314+\vdots \\
315+-x_{d-1} \\
316+x_{d-2}
317+\end{matrix}
318+\right)
319+\otimes
320+\left(
321+\begin{matrix}
322+\sin m\theta_0 \\
323+\sin m\theta_0 \\
324+\sin m\theta_1 \\
325+\sin m\theta_1 \\
326+\vdots \\
327+\sin m\theta_{d/2-1} \\
328+\sin m\theta_{d/2-1} \\
329+\end{matrix}
330+\right)
331+\tag{14}
332+$$
333+ 
334+To save the storage space of $cos$/$sin$ and simplify the internal implementation logic, note that the final $q^R*k^R$ matrix multiplication performs a Reduce operation: adjusting the positions of elements in rows and columns simultaneously does not affect the final accumulation result. Therefore, the computation of the internal $ROPE$ implementation is adjusted into the form of formula (15). The input $x$ is split into two parts by odd and even positions, and the inputs of the corresponding $cos$ and $sin$ parts are consecutive.
335+ 
336+$$
337+ROPE(x) = R_{\Theta,m}^{d} x =
338+\left(
339+\begin{matrix}
340+x_0 \\
341+x_2 \\
342+\vdots \\
343+x_{d-2} \\
344+x_{1} \\
345+x_{3} \\
346+\vdots \\
347+x_{d-1}
348+\end{matrix}
349+\right)
350+\otimes
351+\left(
352+\begin{matrix}
353+\cos m\theta_0 \\
354+\cos m\theta_1 \\
355+\vdots \\
356+\cos m\theta_{d/2-1} \\
357+\cos m\theta_0 \\
358+\cos m\theta_1 \\
359+\vdots \\
360+\cos m\theta_{d/2-1} \\
361+\end{matrix}
362+\right)
363++
364+\left(
365+\begin{matrix}
366+-x_1 \\
367+-x_3 \\
368+\vdots \\
369+-x_{d-1} \\
370+x_0 \\
371+x_2 \\
372+\vdots \\
373+x_{d-2}
374+\end{matrix}
375+\right)
376+\otimes
377+\left(
378+\begin{matrix}
379+\sin m\theta_0 \\
380+\sin m\theta_1 \\
381+\vdots \\
382+\sin m\theta_{d/2-1} \\
383+\sin m\theta_0 \\
384+\sin m\theta_1 \\
385+\vdots \\
386+\sin m\theta_{d/2-1} \\
387+\end{matrix}
388+\right)
389+\tag{15}
390+$$
391+ 
392+In the entire MlaProlog operator, $ROPE$ computation is performed on the $q^R$ matrix and the $k^R$ matrix to obtain the final $queryRope$ and $keyRope$ information, and the $keyRope$ information is finally updated into the $KRCache$.
393+ 
394+$$
395+q^R = ROPE(q^R)
396+$$
397+ 
398+$$
399+k^R = ROPE(k^R)
400+$$
401+ 
402+### MatmulQn
403+ 
404+The $q^C$ matrix is multiplied by the Key upsampling matrix $W^{UK}$ to obtain the final Query matrix $q^N$.
405+$$
406+q^N = q^C \cdot W^{UK} \tag{16}
407+$$
408+ 
409+### KVCache
410+ 
411+After the Key/Value results corresponding to the input $x$ are computed, the Key/Value results are updated to the corresponding positions in the KVCache. Currently, cacheIndex is introduced to identify the storage position of the computation result in the KVCache. cacheIndex is a 2-dimensional tensor with the shape [B, S], identifying the target update position of each Token in the Query. Currently, KVCache mainly supports the non-PA (Page Attention) scenario and the PA scenario (ND-format storage and NZ-format storage).
412+ 
413+PA scenario
414+ 
415+- KVCache uses ND-format storage, and its update flow is shown in Figure 3.
416+ - Assume that the Token/sequence length in the Query is S1. A specific value in the second dimension of cacheIndex indicates the row of the KVCache to which the KVCache computed in the current round is to be updated. For example, 11 in the figure indicates updating to the position of row 11, which is located in Block2.
417+ - Figure 3 KVCacheScatter operation (PA scenario - ND format)
418+ ![KVCacheScatter_PA_ND](../../../docs/en/figures/KVCacheScatter_PA_ND.jpg)
419+- KVCache uses NZ-format storage, and its update flow is shown in Figure 4. For the introduction to the fractal data format, refer to the official website guide - [Data Layout Format](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC3alpha003/devguide/opdevg/ascendcopdevg/atlas_ascendc_10_0104.html).
420+ - Currently, when KVCache is stored in NZ format, the entire Block is stored in NZ format. Within a small block, the layout is row-major, forming a Z shape; between blocks, the layout is column-major, forming an N shape. After the complete H (Hidden States) is divided into small blocks of 32B length, small blocks that are stored contiguously in ND format need to be stored with a jump of $32B * BlockSize$ in NZ format, which means the stride needs to be set to $32B * BlockSize$.
421+ - Figure 4 KVCacheScatter operation (PA scenario - NZ format)
422+ ![PA_NZ](../../../docs/en/figures/PA_NZ.jpg)
423+ 
424+### KRCache
425+ 
426+After the KeyRope result corresponding to the input $x$ is computed, the KeyRope result is updated to the corresponding position in the KRCache. Currently, cacheIndex is introduced to identify the storage position of the computation result in the KRCache.
427+ 
428+The update logic of KRCache is the same as that of KVCache.