From 3a1e27a3adfef4b61be1059559a23dedaa87006c Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Sat, 21 Mar 2026 09:16:30 +0000
Subject: [PATCH 01/15] (ci) bump CANN version

---
 .github/workflows/build.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 6ce48d49..59494cbd 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -70,7 +70,7 @@ jobs:
             python_version: 3.12
 
     container:
-      image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-beta.1-910-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release
@@ -104,7 +104,7 @@ jobs:
           - arch: x86_64
             python_version: 3.12
     container:
-      image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-beta.1-910-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From 83ff20a2364d39de7b3a803dafe2f2530328f028 Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Sat, 21 Mar 2026 09:18:12 +0000
Subject: [PATCH 02/15] (cmake) bump  pto-isa tag ver

---
 CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index f24285ff..ee466a24 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -81,7 +81,7 @@ include(FetchContent)
 FetchContent_Declare(
   libpto_isa_headers
   GIT_REPOSITORY https://gitcode.com/cann/pto-isa.git
-  GIT_TAG 8.5.0)
+  GIT_TAG v9.0.0-beta.1)
 
 FetchContent_Populate(libpto_isa_headers)
 

From d4dda45ba2873f069751a46e1887c464f5fba194 Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Sat, 21 Mar 2026 12:31:48 +0100
Subject: [PATCH 03/15] Update Docker image tag in build workflow

---
 .github/workflows/build.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 59494cbd..72f9c1d7 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -70,7 +70,7 @@ jobs:
             python_version: 3.12
 
     container:
-      image: quay.io/ascend/cann:9.0.0-beta.1-910-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-beta.1-910b-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From f48e8d1548958292e08e403837c3005db102390a Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Tue, 14 Apr 2026 12:06:03 +0200
Subject: [PATCH 04/15] Update build.yml

---
 .github/workflows/build.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 72f9c1d7..79d4d727 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -104,7 +104,7 @@ jobs:
           - arch: x86_64
             python_version: 3.12
     container:
-      image: quay.io/ascend/cann:9.0.0-beta.1-910-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-beta.2-910b-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From 4307a97e3deecb434c08edad738870053ee8eb87 Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Tue, 14 Apr 2026 12:06:41 +0200
Subject: [PATCH 05/15] Update Docker image version in build workflow

---
 .github/workflows/build.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 79d4d727..0557ad13 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -70,7 +70,7 @@ jobs:
             python_version: 3.12
 
     container:
-      image: quay.io/ascend/cann:9.0.0-beta.1-910b-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-beta.2-910b-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From fb4a17b1f5a6756fa13085175a7cf31e1c58bdfc Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Thu, 7 May 2026 14:36:50 +0200
Subject: [PATCH 06/15] Update Docker image tag in build workflow

---
 .github/workflows/build.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index ee8a6547..dad85bf7 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -78,7 +78,7 @@ jobs:
             npu_device: a3
 
     container:
-      image: quay.io/ascend/cann:9.0.0-beta.2-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From f1a5f5024fbb575b36ca9f2a4c16cdb707cb18ac Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Thu, 7 May 2026 15:11:21 +0200
Subject: [PATCH 07/15] Update GIT_TAG for libpto_isa_headers to v9.0.0

---
 CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6331888c..bb392747 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -85,7 +85,7 @@ include(FetchContent)
 FetchContent_Declare(
   libpto_isa_headers
   GIT_REPOSITORY https://gitcode.com/cann/pto-isa.git
-  GIT_TAG v9.0.0-beta.2)
+  GIT_TAG v9.0.0)
 
 FetchContent_Populate(libpto_isa_headers)
 

From e9f7f5cff22d3e545aa5bce1dfde34adc0ac62ec Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Thu, 7 May 2026 15:13:26 +0200
Subject: [PATCH 08/15] Update build.yml to remove Python 3.12 entries

Removed Python 3.12 configurations from the build matrix.
---
 .github/workflows/build.yml | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index dad85bf7..c0290285 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -68,14 +68,9 @@ jobs:
             python_version: 3.11
             npu_device: 910b
           - arch: x86_64
-            python_version: 3.12
-            npu_device: 910b
           - arch: x86_64
             python_version: 3.11
             npu_device: a3
-          - arch: x86_64
-            python_version: 3.12
-            npu_device: a3
 
     container:
       image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
@@ -110,17 +105,11 @@ jobs:
           - arch: x86_64
             python_version: 3.11
             npu_device: 910b
-          - arch: x86_64
-            python_version: 3.12
-            npu_device: 910b
           - arch: x86_64
             python_version: 3.11
             npu_device: a3
-          - arch: x86_64
-            python_version: 3.12
-            npu_device: a3
     container:
-      image: quay.io/ascend/cann:9.0.0-beta.2-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From e9831049c2bc6f3f09e594c02b4c5daefca255ad Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Thu, 7 May 2026 15:14:48 +0200
Subject: [PATCH 09/15] Remove duplicate architecture entry in build.yml

---
 .github/workflows/build.yml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index c0290285..2d34cd09 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -67,7 +67,6 @@ jobs:
           - arch: x86_64
             python_version: 3.11
             npu_device: 910b
-          - arch: x86_64
           - arch: x86_64
             python_version: 3.11
             npu_device: a3

From 426cf85026761449a7a6344562f887b321ab823e Mon Sep 17 00:00:00 2001
From: Anastasios Zouzias <zouzias@users.noreply.github.com>
Date: Sun, 10 May 2026 09:07:09 +0200
Subject: [PATCH 10/15] Update npu_device from 'a3' to '950' in build.yml

---
 .github/workflows/build.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 2d34cd09..8da8eb35 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -69,7 +69,7 @@ jobs:
             npu_device: 910b
           - arch: x86_64
             python_version: 3.11
-            npu_device: a3
+            npu_device: 950
 
     container:
       image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
@@ -106,7 +106,7 @@ jobs:
             npu_device: 910b
           - arch: x86_64
             python_version: 3.11
-            npu_device: a3
+            npu_device: 950
     container:
       image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
     steps:

From f6a9bdfb8b93d30da5f27a40c3dbc51c8836fcae Mon Sep 17 00:00:00 2001
From: anastasios <anastasios.zouzias@huawei.com>
Date: Sun, 10 May 2026 11:01:46 +0000
Subject: [PATCH 11/15] fix TGATHER

---
 csrc/kernel/kernel_csr_gather.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/csrc/kernel/kernel_csr_gather.cpp b/csrc/kernel/kernel_csr_gather.cpp
index 75eae22d..0221d690 100644
--- a/csrc/kernel/kernel_csr_gather.cpp
+++ b/csrc/kernel/kernel_csr_gather.cpp
@@ -116,10 +116,13 @@ AICORE void runTCsrGather(__gm__ T* values, __gm__ int32_t* indices,
     TileDataVal wTiles(remaining_elements);
     TileDataVal zTiles(remaining_elements);
     TileDataIdx idxTiles(remaining_elements);
+    TileDataIdx tmpTiles(remaining_elements);
 
     // Assign the UB address for each tile
     TASSIGN(valTiles, V_T_ADDR + stage * TILE_SIZE_IN_BYTES);
     TASSIGN(wTiles, W_T_ADDR + stage * 2 * TILE_SIZE_IDX_IN_BYTES);
+    TASSIGN(tmpTiles, W_T_ADDR + stage * 2 * TILE_SIZE_IDX_IN_BYTES +
+                          TILE_SIZE_IDX_IN_BYTES);
     TASSIGN(zTiles, Z_T_ADDR + stage * TILE_SIZE_IN_BYTES);
     TASSIGN(idxTiles, IDX_T_ADDR + stage * TILE_SIZE_IDX_IN_BYTES);
 
@@ -147,7 +150,7 @@ AICORE void runTCsrGather(__gm__ T* values, __gm__ int32_t* indices,
     pipe_barrier(PIPE_V);
 
     // Gather
-    TGATHER(wTiles, xTiles, idxTiles);
+    TGATHER(wTiles, xTiles, idxTiles, tmpTiles);
 
     // Signal end of gather to MTE2 (next load)
     set_flag(PIPE_V, PIPE_MTE2, ev1);

From 13e3fef3279e32eed125a9d64f528ea5fbeacba0 Mon Sep 17 00:00:00 2001
From: anastasios <anastasios.zouzias@huawei.com>
Date: Sun, 10 May 2026 11:59:29 +0000
Subject: [PATCH 12/15] fix aicpu examples

---
 .github/workflows/build-examples-aicpu.yml | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/build-examples-aicpu.yml b/.github/workflows/build-examples-aicpu.yml
index 90f65b52..2e0eb8c1 100644
--- a/.github/workflows/build-examples-aicpu.yml
+++ b/.github/workflows/build-examples-aicpu.yml
@@ -22,17 +22,17 @@ jobs:
             python_version: 3.11
             npu_device: 910b
           - arch: x86_64
-            python_version: 3.12
-            npu_device: 910b
+            python_version: 3.11
+            npu_device: 950
           - arch: x86_64
             python_version: 3.11
-            npu_device: a3
+            npu_device: 910b
           - arch: x86_64
-            python_version: 3.12
-            npu_device: a3
+            python_version: 3.11
+            npu_device: 950
 
     container:
-      image: quay.io/ascend/cann:8.5.1-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
+      image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}
     steps:
       - name: Show OS
         run: cat /etc/os-release

From bb1bdee6a875c5004f4cac7c1104d464883bf281 Mon Sep 17 00:00:00 2001
From: anastasios <anastasios.zouzias@huawei.com>
Date: Tue, 12 May 2026 06:47:20 +0000
Subject: [PATCH 13/15] update internal ci

---
 .gitlab-ci.yml | 30 ++++++++----------------------
 1 file changed, 8 insertions(+), 22 deletions(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 76b4b2d2..f64f04d5 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -5,7 +5,6 @@ stages:
    - build_in_docker
    - test
    - test_aarch64_910B2_py3.11
-   - test_aarch64_910B2_py3.12 # Serialize the py3.11 and py3.12 tests otherwise CI fails!
 
 
 # Define the workflow rules manually to make sure only a single pipeline
@@ -23,11 +22,11 @@ workflow:
 
 
 .build_config: &build_config
-  image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-${PYTHON_VERSION}
+  image: quay.io/ascend/cann:9.0.0-${NPU}-ubuntu22.04-py3.11
   parallel:
     matrix:
-      - {PYTHON_VERSION: py3.11}
-      - {PYTHON_VERSION: py3.12}
+      - {NPU: 910b}
+      - {NPU: 950}
   before_script:
     - source /usr/local/Ascend/ascend-toolkit/set_env.sh
     - pip3 install pyyaml setuptools pytest packaging
@@ -42,20 +41,9 @@ build_docker:
     - make build_wheel
 
 build_and_test_x86_910B4:
+  image: quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
   tags:
     - docker_npu # Build and test on-device NPU (x86-64 / 910B4)
-  <<: *build_config
-  script:
-    - make setup_once
-    - make build_wheel
-    - pip install pto_kernels*.whl
-    - python3 -m pytest -v tests/
-
-build_and_test_aarch64_910B2_py3.11:
-  image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py3.11
-  stage: test_aarch64_910B2_py3.11
-  tags:
-    - docker_npu_910B2 # Build and test on-device (aarch64 / 910B2)
   before_script:
     - source /usr/local/Ascend/ascend-toolkit/set_env.sh
     - pip3 install pyyaml setuptools pytest packaging
@@ -65,12 +53,10 @@ build_and_test_aarch64_910B2_py3.11:
     - make build_wheel
     - pip install pto_kernels*.whl
     - python3 -m pytest -v tests/
-  when: manual
-  allow_failure: true
 
-build_and_test_aarch64_910B2_py3.12:
-  image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py3.12
-  stage: test_aarch64_910B2_py3.12
+build_and_test_aarch64_910B2_py3.11:
+  image: quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
+  stage: test_aarch64_910B2_py3.11
   tags:
     - docker_npu_910B2 # Build and test on-device (aarch64 / 910B2)
   before_script:
@@ -89,7 +75,7 @@ build_and_test_aarch64_910B2_py3.12:
 # On-device test of examples/jit_cpp #
 ######################################
 test_examples_on_x86_910B4:
-  image: quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py3.11
+  image: quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
   tags:
     - docker_npu # Build and test on-device NPU (x86-64 / 910B4)
   parallel:

From b041d647bae9b4a2059387c19d468b2f92db9059 Mon Sep 17 00:00:00 2001
From: anastasios <anastasios.zouzias@huawei.com>
Date: Tue, 12 May 2026 07:23:49 +0000
Subject: [PATCH 14/15] fix dup envs

---
 .github/workflows/build-examples-aicpu.yml | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/.github/workflows/build-examples-aicpu.yml b/.github/workflows/build-examples-aicpu.yml
index 2e0eb8c1..f39bd339 100644
--- a/.github/workflows/build-examples-aicpu.yml
+++ b/.github/workflows/build-examples-aicpu.yml
@@ -24,12 +24,6 @@ jobs:
           - arch: x86_64
             python_version: 3.11
             npu_device: 950
-          - arch: x86_64
-            python_version: 3.11
-            npu_device: 910b
-          - arch: x86_64
-            python_version: 3.11
-            npu_device: 950
 
     container:
       image: quay.io/ascend/cann:9.0.0-${{ matrix.npu_device }}-ubuntu22.04-py${{ matrix.python_version }}

From 9eadcc615732259e5617c888612fd8bf85cddf6c Mon Sep 17 00:00:00 2001
From: anastasios <anastasios.zouzias@huawei.com>
Date: Tue, 12 May 2026 07:29:39 +0000
Subject: [PATCH 15/15] bump torch-npu version

---
 .github/workflows/build.yml | 4 ++--
 Makefile                    | 2 +-
 pyproject.toml              | 4 ++--
 requirements.txt            | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 8da8eb35..217efe4c 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -90,7 +90,7 @@ jobs:
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib # to avoid libascend_hal.so issue
           pip3 install pyyaml setuptools pytest packaging pybind11[global]
           pip3 install -r requirements.txt
-          pip3 install torch-npu==2.8.0.post2 --extra-index-url https://download.pytorch.org/whl/cpu
+          pip3 install torch-npu==2.8.0.post4 --extra-index-url https://download.pytorch.org/whl/cpu
           make clean build_cmake
 
   build_wheel_docker:
@@ -126,5 +126,5 @@ jobs:
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib # to avoid libascend_hal.so issue
           pip3 install pyyaml setuptools pytest packaging pybind11[global]
           pip3 install -r requirements.txt
-          pip3 install torch-npu==2.8.0.post2 --extra-index-url https://download.pytorch.org/whl/cpu
+          pip3 install torch-npu==2.8.0.post4 --extra-index-url https://download.pytorch.org/whl/cpu
           make clean build_wheel
diff --git a/Makefile b/Makefile
index 71142dd2..0bf7a3e8 100644
--- a/Makefile
+++ b/Makefile
@@ -15,7 +15,7 @@ clean:
 
 setup_once:
 	pip3 install -r requirements.txt
-	pip3 install torch-npu==2.8.0.post2 --extra-index-url https://download.pytorch.org/whl/cpu
+	pip3 install torch-npu==2.8.0.post4 --extra-index-url https://download.pytorch.org/whl/cpu
 
 build_cmake: clean
 	bash scripts/build.sh
diff --git a/pyproject.toml b/pyproject.toml
index ecc06667..3d84ec17 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ requires = [
   "psutil==6.1.0",
   "PyYAML==6.0.2",
   "attrs==24.2.0",
-  "torch-npu==2.8.0.post2",
+  "torch-npu==2.8.0.post4",
 ]
 build-backend = "scikit_build_core.build"
 
@@ -51,7 +51,7 @@ dependencies = [
 
 [project.optional-dependencies]
 torch = [
-  "torch-npu==2.8.0.post2",
+  "torch-npu==2.8.0.post4",
 ]
 
 [tool.scikit-build]
diff --git a/requirements.txt b/requirements.txt
index e2fcff27..ee43f35c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,7 +9,7 @@ setuptools==80.9.0
 packaging
 wheel
 scikit-build-core==0.11.6
-pip==26.0.1
+pip==26.1.1
 pytest==8.3.4
 numpy==1.26.4
 scipy==1.13.1