Skip to content

[Docs] Add Ascend NPU installation section to README #938

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zheliuyu wants to merge 2 commits into from
Closed

[Docs] Add Ascend NPU installation section to README #938

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6648f0c
[Docs] Add Ascend NPU installation section to README
zheliuyu Jun 5, 2026
4fd0a2f
Merge branch 'main' into main
zheliuyu Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,16 @@
</div>

<p>
💥 Flash Linear Attention brings together hardware-efficient building blocks, training-ready layers, and components for modern sequence models, spanning linear attention, sparse attention, state space models, and hybrid LLM architectures. All implementations are platform-agnostic and verified on NVIDIA, AMD, and Intel hardware. Pull requests are welcome!
💥 Flash Linear Attention brings together hardware-efficient building blocks, training-ready layers, and components for modern sequence models, spanning linear attention, sparse attention, state space models, and hybrid LLM architectures. All implementations are platform-agnostic and verified on NVIDIA, AMD, Intel and NPU hardware. Pull requests are welcome!
</p>

--------

* [News](#news)
* [Models](#models)
* [Installation](#installation)
* [NVIDIA / AMD / XPU](#nvidia--amd--xpu)
* [Ascend NPU](#ascend-npu)
* [Usage](#usage)
* [Token Mixing](#token-mixing)
* [Fused Modules](#fused-modules)
Expand Down Expand Up @@ -114,6 +116,8 @@

## Installation

### NVIDIA / AMD / XPU

[![nvidia-h100-ci](https://github.com/fla-org/flash-linear-attention/actions/workflows/nvidia-h100.yml/badge.svg?branch=main&event=push)](https://github.com/fla-org/flash-linear-attention/actions/workflows/nvidia-h100.yml)

The following requirements should be satisfied
Expand Down Expand Up @@ -151,6 +155,28 @@ pip install einops ninja datasets transformers numpy
pip uninstall fla-core flash-linear-attention -y && pip install -U --no-use-pep517 git+https://github.com/fla-org/flash-linear-attention --no-deps
```

### Ascend NPU

Ascend NPUs use [triton-ascend](https://github.com/triton-lang/triton-ascend), not the upstream [triton](https://github.com/triton-lang/triton) package from PyPI. Installing both in one environment can cause import or kernel dispatch issues.

**Recommended order**

1. Install CANN and set environment variables (`source` the CANN `set_env.sh` script).
2. Clone this repo and run `pip install -e ./`.
3. Install `torch` and `torch_npu` builds that match your CANN version.
4. Uninstall triton: `pip uninstall -y triton`.
5. Install `triton-ascend`.

**Example**

```sh
git clone https://github.com/fla-org/flash-linear-attention.git
cd flash-linear-attention
pip install -e ./
pip install torch==2.7.1 torch_npu==2.7.1
pip uninstall -y triton
pip install triton-ascend==3.2.0
```
Comment on lines +162 to +179

@gemini-code-assist gemini-code-assist Bot Jun 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Running pip install -e ./ first will trigger the download and installation of the standard torch and triton packages from PyPI (which are very large). Immediately after, the instructions ask the user to install torch_npu (which overwrites torch) and uninstall triton to install triton-ascend. This results in redundant downloads of several gigabytes and significantly slower installation times.

Instead, we can recommend installing the correct dependencies (torch, torch_npu, triton-ascend, and other requirements) first, and then installing the repository with --no-deps.

Here is a suggested improvement for the recommended order and example:

Recommended order

  1. Install CANN and set environment variables (source the CANN set_env.sh script).
  2. Install torch and torch_npu builds that match your CANN version.
  3. Install triton-ascend.
  4. Install other required dependencies (pip install einops transformers datasets).
  5. Clone this repo and install it with --no-deps: pip install -e . --no-deps.

Example

git clone https://github.com/fla-org/flash-linear-attention.git
cd flash-linear-attention
pip install torch==2.7.1 torch_npu==2.7.1
pip install triton-ascend==3.2.0
pip install einops transformers datasets
pip install -e . --no-deps

@zheliuyu zheliuyu Jun 5, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We considered this installation approach. In the end, we decided that while installing triton and then uninstalling it is a bit redundant, it serves as a reminder that triton and triton-ascend are incompatible with each other.


## Usage

Expand Down
Loading