Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs

arXiv:2611.01936v1 [cs.DC] 1 Nov 2026

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs Jiahe Wang

Yingxia Shao∗

Shuhan Wang

Beijing University of Posts and Telecommunications Beijing, China meaningful@bupt.edu.cn

Beijing University of Posts and Telecommunications Beijing, China shaoyx@bupt.edu.cn

Beijing University of Posts and Telecommunications Beijing, China wangshuhan@bupt.edu.cn

Qingbo Ge

Hongzheng Li

Yucheng Zhang∗

Beijing University of Posts and Telecommunications Beijing, China geqingbo@bupt.edu.cn

Beijing University of Posts and Telecommunications Beijing, China Ethan_Lee@bupt.edu.cn

Huawei Co., Ltd Beijing, China zhangyucheng23@huawei.com

Abstract Sparse Matrix-Dense Matrix Multiplication (SpMM) is a fundamental kernel in graph neural networks and scientific computing, yet its optimization on Ascend NPUs remains largely unexplored. The physically decoupled Cube/Vector architecture and single-threaded scalar instruction dispatch of modern NPUs introduce unique challenges. We present Ascend-SpMM, a novel and systematic framework tailored to the physically decoupled architecture of Ascend NPUs. We introduce a memory-centric cost model that quantifies hardware affinity and guides a conflict-bounded, densityaware task-mapping strategy. The strategy maps dense regions to the high-throughput Cube path and sparse residuals to the flexible Vector path, with a theoretical bound on cross-processor conflicts. To execute these workloads with high efficiency, we further develop highly customized sub-kernels. These kernels leverage hierarchy-aware pipelining and round-robin ring buffer-based implicit synchronization to effectively mask memory latency and reduce instruction pipeline stalls. Extensive experiments demonstrate that Ascend-SpMM achieves a geometric-mean speedup of 12.70– 20.55× over the strongest evaluated NPU baseline for each matrix, accelerates end-to-end GNN training by up to 6.8×, and delivers performance competitive with state-of-the-art GPU implementations. The source code of Ascend-SpMM is publicly available at https://github.com/IDMG-Lab/AscendSpMM. ACM Reference Format: Jiahe Wang, Yingxia Shao, Shuhan Wang, Qingbo Ge, Hongzheng Li, and Yucheng Zhang. 2027. Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs. In Proceedings of 32nd ACM SIGPLAN Annual Symposium on Principles and ∗ Corresponding author.

PPoPP ’27, Salt Lake City, Utah, USA 2027. ACM ISBN 978-x-xxxx-xxxx-x/YYYY/MM https://doi.org/XXXXXXX.XXXXXXX

Practice of Parallel Programming (PPoPP ’27). ACM, New York, NY, USA, 12 pages. https://doi.org/XXXXXXX.XXXXXXX

1

Introduction

Sparse Matrix-Dense Matrix Multiplication (SpMM) is a fundamental operator in modern computing, serving as the computational backbone for Graph Neural Networks (GNNs)[1, 26, 44], scientific computing[3, 22], sparse attention[19, 40], and modern sparsity-aware Large Language Model (LLM) inference[15, 17, 38, 45]. As the scale and complexity of these applications continue to grow, accelerating SpMM has become increasingly critical. While Graphics Processing Units (GPUs) have traditionally dominated the acceleration of dense and sparse linear algebra, domain-specific Neural Processing Units (NPUs)—represented by the Ascend AI processor[30]—have emerged as a powerful computing force in data centers, offering massive computational throughput and high memory bandwidth tailored for deep learning. However, high-performance SpMM on Ascend NPUs cannot be achieved by simply porting existing state-of-the-art GPU kernels, primarily due to two fundamental architectural mismatches. The first mismatch lies in the physically decoupled architecture and its independent bandwidth allocation. Modern NPUs isolate dense matrix operations into AI Cube (AIC) units and fine-grained operations into AI Vector (AIV) units, each possessing independent memory access paths and bandwidth capacities. Since SpMM is a memory-bound operator, execution methods that utilize only one type of hardware resource are inherently insufficient. Cube-centric approaches like Moustafa’s [33] force irregular data into the AIC through zero-padding, saturating the AIC bandwidth with invalid zeros while the AIV’s bandwidth remains unused. While NeutronSparse [2] recently attempted to address this via heterogeneous workload coordination, its heuristicbased task mapping lacks rigorous theoretical foundations, leading to sub-optimal generalizability across highly diverse density patterns.

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

The second mismatch stems from the single-threaded execution model of NPUs. Unlike GPUs that rely on hardwaremanaged warp schedulers and implicit caching to transparently hide memory latency, NPUs employ a single-threaded Scalar Unit for explicit instruction dispatching. However, different hardware units (i.e., memory access units and compute units) possess independent instruction pipelines, relying on explicit synchronization among them to achieve latency hiding. This execution model introduces distinct architectural challenges: on the AIC, the memory hierarchy is highly complex, demanding explicit data and instruction orchestration across deep, multi-level memory buffers; on the AIV, while the architecture is relatively simple, throughput is severely bottlenecked by the limited number of hardware synchronization flags. Without carefully leveraging these NPU-specific architectural features, instruction streams can become serialized, leaving hardware resources severely underutilized. To address the aforementioned mismatches, we propose Ascend-SpMM, a systematic framework explicitly tailored for the decoupled architecture and single-threaded execution model of Ascend NPUs. First, to exploit the physically decoupled architecture, we introduce a conflict-bounded densityaware heterogeneous task mapping strategy. Guided by a rigorous memory-centric cost model, our approach achieves load balancing with a theoretical guarantee by mapping structured dense workload to the high-throughput AICs while delegating highly irregular workload to the flexible AIVs with a bounded number of atomic-write conflicts. Second, to overcome the absence of hardware warp schedulers and eliminate pipeline stalls, we develop highly customized sub-kernels featuring hierarchy-aware pipelining and implicit synchronization. By scheduling the scalar-dispatched instruction sequence across the explicit multi-level memory hierarchy, and transforming hardware-limited synchronization flags into spatial ring-buffer implicit synchronization, our design achieves increased overlap between memory access and computation. In summary, our main contributions are as follows: • We identify the main architectural bottlenecks of Ascend NPUs caused by physically decoupled heterogeneous resources and strict scalar instruction dispatching. • We propose Ascend-SpMM, a systematic hardwareaware software framework explicitly tailored to the NPU’s decoupled architecture, effectively transforming hardware limitations into system optimization opportunities. • We propose a conflict-bounded density-aware heterogeneous task mapping strategy. Under a continuous relaxation, we show that the optimal assignment follows a density-threshold structure with limited atomic conflicts.

Jiahe Wang et al. AIC

Data Flow

MTE2

Instruction Flow

L1 Buffer

Memory Transfer Engine

Cube Unit

L0C Buffer

Cube Instr. Queue

Data Cache

Memory

MTE1 L0B Buffer FixPipe

Cache Compute Unit

MTE1 L0A Buffer

Global Memory

L2 Cache

Instr. Cache

Scalar Unit

MTE1 Instr. Queue MTE2 Instr. Queue FixPipe Instr. Queue

AIV MTE3

MTE2

Data Cache Instr. Cache

Unified Buffer Scalar Unit

×C

Vector Unit Vector Instr. Queue MTE2 Instr. Queue MTE3 Instr. Queue ×V

Figure 1. The Ascend NPU Decoupled Architecture. • We design highly optimized NPU-native SpMM subkernels featuring hierarchy-aware pipelining and implicit synchronization, which increase bandwidth utilization and overlap between memory access and computation without relying on hardware multi-threading. • Extensive evaluations on diverse real-world datasets demonstrate that Ascend-SpMM achieves a geometricmean speedup of up to 20.55× over the strongest evaluated NPU baseline. Furthermore, it delivers highly competitive performance relative to the state-of-theart GPU implementations.

2

Background and Motivations

2.1

Ascend NPU Architecture

Ascend is a domain-specific NPU designed for deep learning workloads[30, 31]. As shown in Figure 1, it adopts a physically decoupled architecture: the AI Cube (AIC) for matrix operations and the AI Vector (AIV) for vector operations are separate cores with independent computing units, on-chip memory, and memory access engines[48, 49]. Compute Units. An NPU contains three types of compute units: the Scalar Unit, which is responsible for scalar data processing and program control flow; the Vector Unit, which performs vector operations in the Single Instruction Multiple Data (SIMD) paradigm; and the Cube Unit, which supports FP16 and FP32 matrix multiplications. Memory Hierarchy. The memory hierarchy comprises software-managed Global Memory (GM), Unified Buffer (UB), and L0/L1 buffers exposing data placement to developers, and hardware-managed caches. Data movement is orchestrated by Memory Transfer Engines (MTEs), each serving a specific memory level. Transfers within the same MTE are executed sequentially, while transfers across different MTEs can proceed in parallel. Execution Model. Ascend employs an explicit instruction pipelining mechanism to hide latency. Both compute operations and memory transfers are abstracted as instructions and dispatched by the Scalar Unit to different execution queues corresponding to the Scalar, Vector, Cube, and MTE units. Instructions within the same queue are executed in

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs

non-zero element

0

zero element

1

2

5

8

0

0

1

1

0

2

2

1

3

3

0

1

2

3

4

5

6

7

8

9 10 11

Window 0 Window 1 Window 2

2

TC Block 0

3

3

4

4

9 10 *

TC Block 1

6 11

4

SGT

5

5

6

6

7

7

8

TC Block 2

2

9

5

7 10

8

10

9

11

10

TC Block 3

11

Figure 2. Sparse Graph Translation (SGT).

FIFO order, whereas instructions in different queues can be executed concurrently. 2.2

Sparse Matrix-Dense Matrix Multiplication

SpMM is mathematically defined as 𝐶 = 𝐴 × 𝐵, where 𝐴 is a 𝑀 × 𝐾 sparse matrix, 𝐵 is a 𝐾 × 𝑁 dense matrix, and 𝐶 is the resulting 𝑀 × 𝑁 dense matrix. Vectorized SpMM Vectorized SpMM can be implemented through various approaches[20, 27, 39]. The row-wise approach, also known as Gustavson’s algorithm[21], has been shown to be highly effective. It computes each output row by accumulating parallel products of 𝐴[𝑖, 𝑘] with 𝐵 [𝑘, :], as follows: 𝐶 [𝑖, :] =

∑︁

𝐴[𝑖, 𝑘]𝐵 [𝑘, :]

(1)

𝑘

The row-wise approach can be easily vectorized since each element of 𝐴 is multiplied with all the elements of a row in matrix 𝐵. This yields a vector of partial results, which naturally aligns with the AIV execution paradigm. Cube SpMM Methods designed for GPU Tensor Cores [10] are partially applicable to Cube Units. They require converting sparse matrices into dense TC blocks via Sparse Graph Translation (SGT) introduced by TC-GNN[42]. As shown in Figure 2, we illustrate a simplified example with a block size of 4 × 4. The nonzeros (blue squares) in the sparse matrix are partitioned into row windows. Within each window, every matrix column containing at least one nonzero is represented as a fixed-length column vector padded with zeros (white square). Each window may require additional zero-padding columns on the right side to form a complete TC block. In SpMM, these blocks serve as the left operand; the corresponding dense B blocks are loaded according to column indices of the A blocks. Results are accumulated on-chip and written back to global memory after completing a whole row window.

2.3

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Bottleneck and Opportunity

To reveal the architectural bottlenecks of SpMM on Ascend NPUs, we examine the following two aspects. Write Conflicts in Heterogeneous Execution. Prior work has established that the physically decoupled AIC and AIV can execute concurrently and that coordinating their workloads is important for SpMM performance [2], yet the overhead caused by heterogeneous execution has not been considered: when the AIC and AIV concurrently write partial sums to the same rows of the output matrix 𝐶, atomic accumulation is required for correctness, and the resulting write conflicts can reduce performance. To quantify this effect, we construct a write-only microbenchmark with four configurations: AIC-only, AIV-only, concurrent AIC-AIV execution with overlapping destinations, and concurrent AIC-AIV execution with disjoint destinations. Each run writes the same 640 KB of output data. In both concurrent configurations, the AIC and AIV each write half of the data; the configurations differ only in whether their destination regions overlap. We report the aggregate time of 1,000 consecutive runs. As shown in Figure 3(a), while concurrent execution reduces the total time compared to single-processor execution, overlapping destinations incur an additional 14.9% overhead compared to disjoint destination. This result demonstrates the cost of atomic write conflicts and motivates limiting shared output ownership between the two types of processors. It suggests that heterogeneous workload mapping should consider not only load balance but also output ownership. By assigning each output row exclusively to one processor whenever possible and confining unavoidable AIC-AIV sharing to a bounded set, we can preserve the benefits of concurrent execution with less overhead. MTE Starvation Caused by Inefficient Software Pipelining. The NPU relies on a single-threaded Scalar Unit to explicitly issue instructions and manage data movement across various computing units and multi-level memory hierarchy[30, 31]. Because of this single-threaded execution model, an inefficient software pipeline inflicts severe MTE starvation, exposing different architectural bottlenecks on AIC and AIV. In AIC, data movement and computation across the memory hierarchy exhibit substantially different latencies, while the associated buffers possess vastly different capacities, as summarized in Table 1. Without hierarchy-aware tiling and pipelining, transfer and computation stages become partially serialized: short-latency MTE1, Cube, and Scalar instructions delay the issuance of subsequent MTE2 requests, preventing MTE2 from maintaining enough in-flight transfers. As a result, MTE2 is frequently starved of requests, creating pipeline stalls and underutilizing its bandwidth. In AIV, the hardware imposes a strict limit of eight synchronization flags for coordinating the MTE and compute

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Jiahe Wang et al.

Table 1. AIC memory-hierarchy parameters on Ascend Atlas 800I

Sparse Matrix A

AIC Workload

Heterogeneous Task Mapping

Resource or Data Path

Value

Capacity

L1 buffer L0A/L0B operand buffers L0C accumulation buffer

512 KB 64 KB each 256 KB

MTE2 (GM → L1) MTE1 (L1 → L0A/L0B) Cube (Compute Unit) Scalar (Dispatch Instructions)

> 100 ns ~20 ns ~5 ns ~2 ns

Total time (μs)

3000

2908

20

3096 +14.9% latency 1972 1717

2500 2000 1500 1000 500 0

Effective bandwidth utilization (%)

3500

Cost Model

16.17% 9.92%

10 5.28% 5

2.75%

0

D

ve O

in t

p

(a) Writeback microbenchmark

jo

is

rla

y

y

nl

nl

-o

o VAI

C AI

1

2

4

8

Number of synchronization flags (b) AIV bandwidth utilization

Figure 3. (a) Write-only microbenchmark. (b)Effective bandwidth utilization of the naive CSR SpMM with different numbers of synchronization flags on mip1 (𝑁 = 256, FP32).

units. For the highly irregular AIV execution, eight in-flight memory requests are utterly insufficient to hide the long-tail memory latency, forcing the instruction stream into serialized execution. To quantify this limitation, we execute the naive CSR SpMM on mip1 with 𝑁 = 256 and FP32 while varying the number of synchronization flags, as shown in Figure 3(b). However, even at the hardware limit of eight flags, only 16.17% of the effective bandwidth is utilized. This result shows that the limited number of synchronization flags bounds the number of in-flight memory requests and prevents the AIV from exposing sufficient memory-level parallelism. The above architecture-specific bottlenecks demand NPUfriendly sub-kernel designs. By orchestrating explicit hierarchical pipeline on the AIC and introducing implicit synchronization technique on the AIV, we can effectively mask memory latency and substantially reduce pipeline bubbles.

3

Ascend-SpMM

3.1

Overview

Bounded Atomic Conflicts

AIV Workload

Implicit Synchronization

A nonzeros

MTE2 Instr. Queue

B Rows

Vector Unit

Ring Buffer

Figure 4. Overview of Ascend-SpMM.

8-flag limit

15

Hierarchy-Aware Pipeline

A Blocks

B Tiles

Parameter

Latency

Sub-Kernel Runtime

Output Matrix C

A2 NPU (per AIC).

Ascend-SpMM is a novel systematic framework designed for sparse matrix-dense matrix multiplication, which maximizes the utilization of the heterogeneous resources on Ascend NPUs. Figure 4 provides an overview of Ascend-SpMM. To

achieve efficient hybrid SpMM computation on the decoupled NPU architecture, we first establish a memory-centric cost model to quantify the performance characteristics of the AI Cube (AIC) and AI Vector (AIV) when processing matrices with different density patterns. Based on the cost model, we propose a conflict-bounded density-aware heterogeneous task mapping strategy, which accurately partitions the sparse matrix into Cube-friendly regular blocks and Vector-friendly residuals with light coordination overhead. To execute the resulting partitioned workload efficiently, we further design two NPU-friendly sub-kernels: the AIC kernel exploits a hierarchy-aware pipeline with double buffering to overlap data movement and Cube computation, and the AIV kernel adopts an implicit synchronization based on a ring buffer, transforming irregular memory fetches of Gustavson’s algorithm[21] into a continuous, overlapped instruction stream that bypasses the hardware synchronization-flag limit. 3.2

Memory-Centric Cost Model

SpMM is a fundamentally memory-bound operator, and accessing the dense matrix 𝐵 dominates the data movement cost. We therefore model the AIC and AIV execution times using the amount of 𝐵 data they access. Other hardwaredependent effects, including cache behavior and instruction overhead, are absorbed into two constants, 𝛼 and 𝛽. We partition the sparse matrix 𝐴 into 𝑚 = ⌈𝑀/𝑊 ⌉ row windows, where 𝑊 = 16 matches the row dimension of the Cube unit. Let 𝑏𝑖,𝑘 = 𝐴[𝑖𝑊 : (𝑖 +1)𝑊 , 𝑘] denote a non-empty 𝑊 ×1 block within 𝑖-th row window. Let 𝑑𝑖,𝑘 ∈ [1,𝑊 ] denote its number of nonzeros. The execution costs of assigning this block to the AIC and AIV are modeled as 𝑇AIC (𝑏𝑖,𝑗 ) = 𝛼𝑁 ,

𝑇AIV (𝑏𝑖,𝑗 ) = 𝛽𝑑𝑖,𝑗 𝑁 .

(2)

In our implementation, 𝛼 and 𝛽 are calibrated once using the effective AIC bandwidth 𝐵𝐴𝐼𝐶 and AIV bandwidth 𝐵𝐴𝐼𝑉 : 1 1 𝛼∝ , 𝛽∝ . (3) 𝐵 AIC 𝐵 AIV The AIC pays a fixed cost because the block is zero-padded, whereas the AIV processes only its nonzeros. A dense block amortizes one AIC access to 𝐵 [𝑘, :] over more nonzeros, whereas a sparse block benefits from the AIV’s fine-grained

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs AIC Workload

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Double Buffer Schedule

Instruction Sequence

Sorted by Density

A B A B C A B A B C A B

W=3

Scalar Dispatch Latency

(a)

Marginal Window Split

A Tile A Tile A Tile A Tile A Tile B Tile B Tile B Tile B Tile A A A A B B B B C Tile MMA

Sparse Matrix A AIV Workload

C Tile MMA

C Tile MMA

MTE2 MTE1 Cube

A B B A B B C C A B B A

Scalar Dispatch Latency

Figure 5. Cross-Core Task Mapping. Row windows are sorted by (b)

density, and only the marginal window is split to generate AIC and AIV workloads.

A Tile × g B Tile B Tile

A

A Tile × g B Tile B Tile B

A

B

C Tile C Tile MMA MMA

execution. Consequently, dense blocks should be prioritized for the AIC and sparse blocks for the AIV.

A Tile × g B Tile B Tile B

MTE2 MTE1

B

C Tile C Tile MMA MMA

Cube

Figure 6. AIC pipeline schedules. (a) Naive tile-by-tile scheduling 3.3

Conflict-bounded Density-Aware Task Mapping

At the finest granularity, each 𝑊 × 1 block in 𝐴 could be ranked by its density, with dense blocks assigned to the AIC and sparse blocks assigned to the AIV according to our cost model. However, such a mapping may split many row windows to different types of processors. The AIC and AIV would then update the same rows of 𝐶, introducing atomic operation conflicts for every split window. To bound the atomic conflicts, we formulate task mapping at the rowwindow granularity. Let window 𝑖 contain 𝑐𝑖 nonempty blocks and 𝑧𝑖 nonzeros, and define its effective density as 𝜌𝑖 = 𝑧𝑖 /𝑐𝑖 . Under a continuous relaxation, let 𝑥𝑖 ∈ [0, 1] be the fraction of window 𝑖 assigned to the AIC. Assuming that nonzeros are uniformly distributed within each divisible window, the mapping problem is ( ) 𝑚 𝑚 ∑︁ ∑︁ min 𝑚 max 𝛼𝑁 𝑥𝑖 𝑐𝑖 , 𝛽𝑁 (1 − 𝑥𝑖 )𝑧𝑖 . (4) x∈ [0,1]

𝑖=1

𝑖=1

The two terms are the concurrent AIC and AIV execution costs, respectively. Theorem 1 (Threshold Optimality). For the relaxation in Eq. (4), suppose the windows are ordered such that 𝜌 1 ≥ 𝜌 2 ≥ · · · ≥ 𝜌𝑚 . There exists an optimal solution with an index 𝑘 satisfying 𝑥𝑖 = 1 (𝑖 < 𝑘),

𝑥𝑘 ∈ [0, 1],

𝑥𝑖 = 0 (𝑖 > 𝑘).

interleaves MTE2 requests with MTE1 and Cube instructions, introducing more MTE2 issue latency. (b) The hierarchy-aware pipeline with 𝑔 = 2 batches sparse tiles and issues two dense-tile transfers in advance, allowing MTE2 transfers to overlap with inner-level instruction dispatch and execution.

After sorting the windows, a linear scan identifies the boundary window 𝑘. We discretize only this window. Let its block occupancies be sorted as 𝑑𝑘,1 ≥ · · · ≥ 𝑑𝑘,𝑐𝑘 and define 𝐷𝑘 (𝑞) =

𝑞 ∑︁ 𝑗=1

𝑑𝑘,𝑗 ,

𝐶𝑘 =

∑︁ 𝑖<𝑘

𝑐𝑖 ,

𝑍𝑘 =

∑︁

𝑧𝑖 .

𝑖>𝑘

The number of boundary blocks assigned to the AIC is selected by n  o 𝑞★ = arg min 𝑁 max 𝛼 (𝐶𝑘 + 𝑞), 𝛽 𝑍𝑘 + 𝑧𝑘 − 𝐷𝑘 (𝑞) , 0≤𝑞 ≤𝑐𝑘

(5) Hence, cross-processor output sharing remains bounded to at most one row window. Figure 5 illustrates the task mapping strategy with 𝑊 = 3. The selected AIC blocks are stored in 𝑊 × 1 BSR format[7], whereas the AIV workload remains in CSR format. Window profiling takes O (NNZ) time, sorting takes O (𝑚 log 𝑚), and the boundary window can be refined by bucket sorting in O (𝑐𝑘 + 𝑊 ). The total preprocessing complexity is therefore O (NNZ +𝑚 log 𝑚), and the resulting mapping can be reused across repeated SpMM executions.

Thus, at most one row window is shared by the AIC and AIV.

3.4

Proof. Consider two windows 𝑖 < 𝑗 with 𝑥𝑖 < 1 and 𝑥 𝑗 > 0. Moving an equal AIC block budget 𝜖 from 𝑗 to 𝑖 preserves the AIC cost and changes the AIV cost by −𝛽𝑁 𝜖 (𝜌𝑖 − 𝜌 𝑗 ) ≤ 0. Repeating this exchange produces a density-ordered prefix assignment. Consolidating any fractional assignments with equal density leaves at most one fractional window. □

To improve utilization of the Matrix Multiply-Add (Cube) Unit on the Ascend NPU, we design a highly customized AIC sub-kernel tailored to the explicit multi-level memory hierarchy (Global Memory → L1 → L0A/L0B → L0C). The kernel uses 𝑊 × 1 BSR blocks as the basic storage in Global Memory. We first select a tile shape that reduces padding

AIC Kernel with Hierarchy-Aware Pipelining

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

overhead while respecting the on-chip buffer constraints, then introduce a hierarchy-aware pipeline that decouples the coarse-grained MTE2 transfers from the finer-grained MTE1 and Cube operations. Hierarchy-Aware Tiling. Let (𝑚, 𝑘, 𝑛) denote the tile dimensions for multiplying an 𝑚 × 𝑘 sparse tile with a 𝑘 × 𝑛 dense tile. As observed by previous work [36], reducing 𝑚 decreases the number of padded zeros within a sparse tile and therefore reduces redundant Cube computation. We consequently set 𝑚 = 𝑊 = 16, which is the smallest Cubecompatible tile height and matches the height of one row window. For a fixed 𝑘, increasing 𝑛 reduces the number of tiles along the feature dimension of 𝐵. This reduces the number of MTE2 transfer instructions for 𝐵, improves per-instruction throughput by amortizing the fixed launch overhead. We therefore maximize 𝑛 under the on-chip buffer constraints. Since the minimum Cube-compatible reduction dimension is 𝑘 = 16, the double-buffered L0B capacity imposes 2𝑘𝑛𝑠𝐵 ≤ 𝑆 L0B , where 𝑠𝐵 is the element size and 𝑆 L0B = 64 KB, which indicates 𝑛 = 1024 under FP16 and 𝑛 = 512 under FP32. Hierarchy-Aware Pipeline. After fixing the tile size, we decouple the scheduling granularities of MTE2 and the inner pipeline. Instead of processing one tile at a time, MTE2 stages a batch of tiles in L1, while MTE1 and the Cube Unit continue to consume the staged data tile by tile. Figure 6(a) illustrates a naive tile-by-tile pipeline. For each tile, the Scalar Unit dispatches its MTE2, MTE1, and Cube instructions to corresponding hardware units before proceeding to the next tile. Because instruction dispatch is singlethreaded, the MTE1 and Cube instructions inserted between consecutive MTE2 transfers delay the issue of the next GMto-L1 requests. These issue latencies limit the number of in-flight MTE2 transfers and result in low MTE2 bandwidth utilization. As shown in Figure 6(b), we instead batch multiple tiles at the MTE2 level. The AIC workload consists of 𝑊 × 1 blocks, which can be concatenated into a larger 𝑊 × (𝑔𝑘) tile for a batch of 𝑔 reduction tiles. This reduces the number of MTE2 instructions accessing 𝐴. The corresponding B-tile transfers are also issued consecutively before dispatching their MTE1 and Cube instructions. By issuing multiple MTE2 requests in advance, their transfers overlap with the subsequent dispatch and execution of MTE1 and Cube instructions, reducing MTE2 issue latency and improving its effective bandwidth utilization. Once the data is staged in L1, MTE1 and the Cube Unit consume it at the original 𝑘 = 16 computation granularity. We further restrict 𝑔 to be even to accommodate double buffering. Taking into account the constraint of L1=512KB, we choose 𝑔 = 6 in practice. 3.5

AIV Kernel with Implicit Synchronization

The AIV kernel processes irregular long-tail rows directly in the CSR format. Following Gustavson’s algorithm [21], each

Jiahe Wang et al.

Algorithm 1 AIV Round-Robin Schedule within a row. Input: row_ptr, col_id, value, 𝐵, row_id Output: Result Matrix 𝐶 1: __ub__ b[𝐿 × 𝑁 ], c[𝑁 ]; 2: start = row_ptr[row_id], end = row_ptr[row_id + 1]; 3: duplicate(c, 0, 𝑁 ); 4: for t = start to end do 5: slot = t mod L; 6: copy_gm_to_ub(b[slot ×𝑁 ], 𝐵[col_id[t] ×𝑁 ] , 𝑁 ); 7: set_flag(PIPE_MTE2, PIPE_V, 0); 8: wait_flag(PIPE_MTE2, PIPE_V, 0); 9: vaxpy(c, b[slot ×𝑁 ], value[t], 𝑁 ); 10: end for 11: copy_ub_to_gm(𝐶[row_id ×𝑁 ], c, 𝑁 );

Cache Hit Latency

time 0

2

4

1

(a) Mutiple Buffer

5

7

4 1

5 2

Bandwidth

6 3

0

(c) Implicit Synchronization

Batch

6

3

0

(b) Multiple Buffer + Batched Process

Cache Miss Latency

7 4

1

5 2

6 3

7

Figure 7. Simplified MTE2 bandwidth utilization timeline processing eight row fetches from the dense matrix. (a) Traditional Multiple Buffer strategy with buffer number = 2. (b) Multiple Buffer + Batched Process strategy with buffer num = 2 and batch size = 2. (c) Implicit Synchronization strategy.

nonzero element triggers a fine-grained row fetch 𝐷𝑡 from the dense matrix 𝐵, followed by a Vector vaxpy operation 𝑉𝑡 . Such short and random transfers expose insufficient memorylevel parallelism to fully utilize MTE2. Synchronization Bottleneck. A conventional multi-buffer schedule protects each buffer in UB with both a MTE2-toVector forward dependency, which ensures that computation starts after data arrival, and a Vector-to-MTE2 backward dependency, which prevents the next MTE2 instruction from overwriting the current computation operands. Since the hardware provides only eight synchronization flags, explicit synchronization restricts the number of memory transactions that can be processed simultaneously. Batching can increase the number of issued transfers, but a batch cannot progress until its slowest transfer completes. Consequently, the latency variation between L2 hits and misses introduces tail-induced MTE2 bubbles, as illustrated in Figure 7(a) and (b).

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

N=128, FP32

2000

Ascend-SpMM NeutronSparse Naive CSR Moustafa

3000

Throughput (GFLOPS)

Throughput (GFLOPS)

N=128, FP16 Ascend-SpMM NeutronSparse Naive CSR Moustafa

2500

1500 1000 500

2500 2000 1500 1000 500 0

0

K

101

K

224

K

438

K

810

NNZ

1.6M

3.8M

M

11.6

K

.1M

101

283

K

224

K

438

K

810

N=256, FP16

NNZ

11.6

M

.1M 283

3.8M

11.6

M

.1M 283

Ascend-SpMM NeutronSparse Naive CSR Moustafa

5000

Throughput (GFLOPS)

Throughput (GFLOPS)

4000

3.8M

N=256, FP32

Ascend-SpMM NeutronSparse Naive CSR Moustafa

5000

1.6M

3000 2000 1000

4000 3000 2000 1000 0

0

K

101

K

224

K

438

K

810

NNZ

1.6M

3.8M

M

11.6

K

.1M

101

283

K

224

K

438

K

810

NNZ

1.6M

Figure 8. SpMM performance results on NPU. The matrices are sorted in ascending order by NNZ and each point in the figure represents the average GFLOPS for a contiguous set of 5 matrices. NeutronSparse has some missing data points due to OOM.

Implicit Backward Synchronization. The fundamental bottleneck of the conventional schedule is not the UB capacity, but the limited number of synchronization flags. Each explicitly managed buffer requires forward synchronization to ensure data readiness and backward synchronization to protect its reuse, restricting the number of in-flight buffers to the eight available flags. We therefore manage data transfer and computation by implicit backward synchronization, removing the explicit per-buffer Vector-to-MTE2 backward synchronization and instead guaranteeing safe reuse through the maximum depth of the Vector instruction queue. Let 𝑄𝑉 denote the maximum depth of the Vector instruction queue. We use a ring buffer of 𝐿 buffers for row fetch and vector operation. When the Vector instruction queue is full, the Scalar unit will stop issuing subsequent instructions, until the first instruction to enter the queue is completed [30]. The hardware also ensures that 𝑉𝑡 is completed before 𝑉𝑡 +1 completes, therefore 𝑉𝑡 is completed before 𝑉𝑡 +𝑄𝑉 issued. Since 𝐷𝑡 +𝑄𝑉 +1 is issued after 𝑉𝑡 +𝑄𝑉 , choosing 𝐿 ≥ 𝑄𝑉 + 1 ensures that the previous Vector instruction is completed before it is reused by MTE2, satisfying the backward synchronization requirement. Since instructions complete in program order, the remaining forward dependencies can sequentially reuse a single flag [31]. Consequently, one flag can coordinate a 𝐿-slot ring buffer with safe reuse, decoupling the prefetch distance from the number of hardware synchronization flags. As shown in Algorithm 1, we organize the UB as a ring buffer with 𝐿 slots (Line 1) and map the 𝑡-th dense-row fetch to slot 𝑠 (𝑡) = 𝑡 mod 𝐿 (Line 5). MTE2 continuously issues transfers to successive slots, while the Vector pipeline performs the corresponding vaxpy operations after the forward dependency is satisfied (Line 6-9). By removing batch-wide barriers and backward dependencies, our design converts

fine-grained random fetches into a dense MTE2 instruction stream, reducing bandwidth bubbles and improving sustained MTE2 utilization, as shown in Figure 7(c).

4

Evaluation

4.1

Experimental Setup

Baselines. We compare Ascend-SpMM with state-of-theart methods on both NPUs and GPUs. For SpMM on Ascend NPUs, we compare against (i) an AIV-based SpMM implementation of Gustavson’s algorithm[21], denoted as Naive CSR; (ii) Moustafa[33], an AIC-based SpMM method; and (iii) NeutronSparse[2], an AIC-AIV collaborative SpMM method. Since no ready-to-use SpMM implementations for naive CSR and Moustafa are available on Ascend NPUs, we implemented both baselines following their published descriptions and tuned them on the target platform. For SpMM on GPUs, we compare with (i) CUDA Core-based methods, cuSPARSE[11] and RoDe[34]; (ii) Tensor Core-based methods, DTC-SpMM[14] and FlashSparse[36]; (iii) hybrid CUDA-core/Tensor-core method, HC-SpMM[28]. For end-toend GNN frameworks, we only compare with PyG[16] and MindSpore-GL[32], as DGL[41] is not supported on Ascend NPUs. We evaluate the latest open-source implementations of these methods as strong baselines. The kernels are evaluated in single (FP32) and half precision (FP16), with the width of the dense matrices set to 128 and 256 (N = 128, 256). Datasets. We select 820 square matrices from the SuiteSparse [12] collection with at least 10K rows, 10K columns, and 100K nonzeros. These matrices cover a variety of sparsity patterns and are used to evaluate kernel performance. In addition, we select classic graph datasets Flickr[46], Yelp[46], and ogbn-arxiv[24] for end-to-end GNN performance evaluation.

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Jiahe Wang et al. Ascend-SpMM

Speedup vs cuSparse (log scale)

Environment. We conducted our experiments on an Ascend Atlas 800I A2 with CANN 8.5.0 [29] for NPU experiments, and an NVIDIA A100-SXM-40GB GPU [9] with CUDA 12.2 for GPU experiments. Detailed hardware specifications [4, 8] are listed in Table 2. All tests were run ten consecutive times with warm-ups to obtain the average values using msprof [5]. Parameter Selection. We select the parameters on the target platform: the measured AIC/AIV bandwidths are 295/505 GB/s, yielding 𝛼 = 1/295 and 𝛽 = 1/505. Since the Vector queue depth is undocumented, we characterize it using the official msprof [5] as 𝑄𝑉 = 32 and conservatively set 𝐿 = 64 > 𝑄𝑉 .

FlashSparse

HC-SpMM

RoDe

DTC-SpMM

102

101

4.5 4.4

10 1

4.0

3.4 3.9 1.4

100

128 FP16

128 FP32

2.8

3.9 3.1

256 FP16

2.7 3.1

3.6 1.4

2.0

256 FP32

Figure 9. The speedup distribution of Ascend-SpMM and baselines over cuSPARSE. The number in the box plot represents the median speedup.

4.2 SpMM Performance Comparison on Ascend NPU Figure 8 compares the throughput of Ascend-SpMM with the NPU baselines across 820 matrices sorted by NNZ. AscendSpMM maintains higher throughput over the full NNZ range. For a conservative comparison, Table 3 reports the speedup over the strongest NPU baseline selected separately for each Table 2. Detailed hardware specifications. Specification

Ascend Atlas 800I A2

NVIDIA A100SXM-40GB

Architecture Compute Units

Da Vinci 60 Cores (20 AIC + 40 AIV)

Memory Capacity Memory Bandwidth Peak Performance

32 GB 800 GB/s 280 TFLOPS @ FP16

Ampere 108 SMs (6912 CUDA Cores + 432 Tensor Cores) 40 GB 1555 GB/s 312 TFLOPS @ FP16

Table 3. Speedup distribution of Ascend-SpMM over the strongest NPU baseline for each matrix. Speedup

128 FP16

128 FP32

256 FP16

256 FP32

< 1× (%) [1, 3) (%) [3, 10) (%) ≥ 10× (%) Geo. Mean (×)

0.0% 0.4% 4.9% 94.8% 19.03×

0.0% 0.1% 6.1% 93.8% 20.55×

0.0% 0.5% 17.8% 81.7% 14.98×

0.0% 0.2% 30.6% 69.1% 12.70×

Table 4. SpMM Profile Analysis on Matrix mip1 (66,463 × 66,463, NNZ= 10,352,819), 𝑁 = 256, FP32

matrix and configuration. Ascend-SpMM outperforms this per-matrix baseline on every evaluated matrix, achieving geometric-mean speedups of 19.03×, 20.55×, 14.98×, and 12.70× under 128 FP16, 128 FP32, 256 FP16, and 256 FP32, respectively. Across the four configurations, 99.5–99.9% of the matrices achieve at least a 3× speedup, while 69.1–94.8% exceed 10×; only 0.1–0.5% fall between 1× and 3×. Although the relative gain decreases as 𝑁 increases, Ascend-SpMM still delivers a 12.70× geometric-mean speedup in the most conservative 256 FP32 configuration, demonstrating robust performance across matrix scales, feature widths, and data types. To explain the above speedup, we profile the kernels on the matrix “mip1”, which contains both long and short rows and is representative of the sparse patterns of the majority of matrices. We report execution time and the AIC/AIV MTE2 utilization ratios. As shown in Table 4, naive CSR executes exclusively on the AIV (Vector) cores with an AIV MTE2 ratio of merely 34.8%, while Moustafa runs solely on the AIC (Cube) cores with an AIC MTE2 ratio of only 24.0%, indicating that both baselines leave the MTE severely under-utilized. In contrast, our method employs a carefully designed software pipeline that overlaps computation and data movement within both the AIC and AIV pipelines. This overlap raises the AIC MTE2 ratio to 94.6% and the AIV MTE2 ratio to 93.4%, which in turn delivers an order-of-magnitude increase in effective memory bandwidth. Meanwhile, the cooperative use of both Cube and Vector units fully exploits the heterogeneous hardware resources. The combined effect explains the speedup observed over the baselines.

Metric

Naive CSR

Moustafa

AscendSpMM

4.3

Exec. Time (𝜇s) AIC MTE2 Ratio (%) AIV MTE2 Ratio (%)

95854.44 — 34.8%

58795.06 24.0% —

1350.50 94.6% 93.4%

We further compare our NPU-based method against state-ofthe-art GPU implementations, including FlashSparse, RoDe, DTC-SpMM, HC-SpMM, and cuSPARSE, on the same 820

SpMM Performance Comparison with GPU Methods

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Table 5. Detailed Speedup Distribution vs cuSPARSE. Some GPU methods do not support FP16 computation. N=128 FP16 N=128 FP32 Ascend- Flash Ascend- Flash HCRoDe SpMM Sparse SpMM Sparse SpMM

< 1× (%) [1, 1.5) (%) [1.5, 2) (%) ≥ 2× (%) Geo. Mean (×) Min (×) Max (×)

9.9% 11.0% 18.1% 8.2% 7.9% 11.7% 8.5% 6.2% 7.4% 73.4% 74.9% 62.7% 4.46× 4.82× 3.08× 0.199× 0.027× 0.181× 163.8× 174.1× 65.2×

with hierarchy-aware pipeline

39.7% 12.9% 7.1% 40.3% 1.37× 0.004× 61.1×

6.1% 13.8% 10.7% 15.0% 22.2% 5.9% 14.4% 9.5% 12.8% 14.6% 11.0% 9.4% 8.6% 9.4% 6.2% 77.1% 62.3% 71.1% 62.9% 57.0% 3.63× 2.76× 3.61× 2.96× 2.48× 0.016× 0.472× 0.228× 0.046× 0.225× 126.4× 66.0× 196.2× 239.0× 207.2×

Throughput (GFLOPS)

N=128, FP32

N=128, FP16

2000

2000

1500

1500

1000

1000 500

K K K K M M M M 101 224 438 810 1.6 3.8 11.6 283.1

NNZ

500

K K K K M M M M 101 224 438 810 1.6 3.8 11.6 283.1

NNZ

Throughput (GFLOPS)

2500

2500

11.8% 41.5% 16.0% 11.6% 9.0% 7.8% 63.2% 39.1% 2.94× 1.13× 0.058× 0.007× 310.6× 123.5×

Round-Robin Ring Buffer Batched Process + Multi-Buffer

without hierarchy-aware pipeline

N=128, FP16

Throughput (GFLOPS)

11.7% 11.9% 8.3% 68.1% 3.79× 0.035× 119.2×

N=256 FP16 N=256 FP32 DTC- Ascend- Flash Ascend- Flash HCRoDe SpMM SpMM Sparse SpMM Sparse SpMM

500 400 300 200 100

K K K K M M M M 101 224 438 810 1.6 3.8 11.6 283.1

NNZ

8.4% 4.6% 15.5% 71.5% 3.03× 0.020× 202.4×

DTCSpMM 19.8% 14.3% 16.4% 49.4% 2.12× 0.461× 181.0×

Multiple Buffer Naive CSR

N=128, FP32 Throughput (GFLOPS)

Metric

500 400 300 200 100 0 K K K K M M M M 101 224 438 810 1.6 3.8 11.6 283.1

NNZ

Figure 10. The throughput of AIC sub-kernel with/without Hierarchy-Aware Pipeline.

Figure 11. The throughput of AIV sub-kernel across various overlapping mechanisms.

square-matrix subset. Figure 9 presents the speedup distribution. As shown in Table 5, our method achieves a geometricmean speedup of 2.48×–4.46× over cuSPARSE across the four configurations, with 57.0%–73.4% of datasets exceeding 2× speedup. FlashSparse stands out as the strongest GPU baseline with a geometric-mean of 2.94×–4.82× and 62.9%– 74.9% of datasets above 2×, benefiting from higher raw GPU memory bandwidth and warp-level parallelism on regular sparse patterns. Our method surpasses FlashSparse in the 256 FP16 configuration (3.61× vs. 2.96×), where the larger computational intensity better saturates the NPU’s cube units and amortizes the software pipeline overhead. Compared to the other GPU baselines, our method consistently outperforms DTC-SpMM and HC-SpMM, and achieves comparable results to RoDe. Notably, our method exhibits lower performance variance than FlashSparse: the minimum speedup of 0.199× (128 FP16) versus FlashSparse’s 0.027×, and only 9.9% of cases fall below 1× compared to 11.0% for FlashSparse, indicating our software pipeline provides more robust performance across diverse sparse patterns. These results demonstrate that a carefully designed NPU-based SpMM implementation can achieve competitive and, on certain workload profiles, higher throughput than several stateof-the-art GPU SpMM methods. 4.4

Ablation Study

Effectiveness of Hierarchy-Aware Pipeline. We isolate the contribution of our hierarchy-aware pipeline on the

AIC by comparing two AIC-only implementations: one with the optimized pipeline and one without. Figure 10 presents the GFLOPS comparison under N=128 FP16 and FP32 configurations on the 820 square-matrix subset. Enabling the hierarchy-aware pipeline achieves a geometric-mean speedup of 1.33× under FP16 and 1.31× under FP32. The uniform improvement across FP16 and FP32 further demonstrates that the pipeline’s benefit is precision-agnostic and derives from reducing memory stalls rather than from compute-side optimizations. Effectiveness of Implicit Synchronization. We evaluate the effectiveness of implicit synchronization on the AIV by comparing four AIV-only implementations: (i) Naive CSR, which processes sequentially with one buffer; (ii) Multiple Buffer with buffer number = 8, which decouples memory accesses from scalar execution at the per-element level; (iii) Batched Process + Multiple Buffer with batch size = 8 and buffer number = 8, which further amortizes loop overhead by grouping elements; and (iv) our AIV sub-kernel with 𝐿 = 64 slots. Figure 11 shows the GFLOPS comparison under N=128 FP16 and FP32. Compared with Naive CSR, our AIV sub-kernel achieves geometric-mean speedups of 12.87× under FP16 and 13.43× under FP32. Compared with Multiple Buffer, it remains 4.77× and 5.11× faster, respectively. Even over Batched Process + Multiple Buffer with the same total number of buffers, the round-robin ring buffer provides speedups of 1.81× and

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

Jiahe Wang et al.

Table 6. Geometric-Mean Speedup of Ascend-SpMM over Hetero-

Ascend-SpMM

geneous Load Balancing Ablations.

Naive Load Balancing NeutronSparse Load Balancing Block-Level Split

128 FP16 1.56× 1.36× 1.19×

128 FP32 1.46× 1.29× 1.20×

256 FP16 1.47× 1.26× 1.14×

256 FP32 1.25× 1.23× 1.10×

Training time (s)

Ours vs

16

MindSpore-GL 240

18

8

206.5

19.0

12

4

PyG

24

13.8

5.8

180

12

4.0

6

0

120 5.3

7.4

0

Flickr

71.9 60

30.5

0

ogbn-arxiv

Yelp

Figure 12. End-to-end two-layer GCN training time (200 epochs). Preprocessing time is included.

1.88×, reaching averages of 433 and 423 GFLOPS, respectively. The critical enabler of the speedup is implicit synchronization: the ring buffer allows each buffer slot to independently progress data transfer and computation without explicit barriers. This design eliminates inter-phase pipeline bubbles, which is essential on the AIV processor where the scalar-rich execution pattern of sparse row operations creates irregular compute latencies. We also verified the correctness of implicit synchronization against the Naive CSR implementation with explicit forward and backward synchronization using a single buffer, and observed bitwise-identical results on all 820 matrices for 𝑁 ∈ {16, 32, 64, 128, 256} with random values in standard normal distribution 𝑁 (0, 1). Effectiveness of Conflict-Bounded Density-Aware Task Mapping. To quantify the contribution of conflictbounded density-aware task mapping, we compare our strategy against three ablations: (i) a naive heterogeneous AIC+AIV hybrid that partitions rows by equally distributing the total NNZ count across cores; (ii) NeutronSparse[2] workload partition strategy; and (iii) Block-Level split that maps dense blocks to AIC and sparse blocks to AIV, ignoring atomic write conflicts. Table 6 summarizes the geometric-mean speedup of ours over each ablation across all three ablations. The proposed task mapping strategy consistently outperforms all ablations. Compared with naive NNZ-balanced mapping, it achieves geometric-mean speedups of 1.25×– 1.56×, as equal NNZ partitioning does not capture the distinct execution costs of AIC and AIV. Moreover, our strategy achieves a 1.23×–1.36× speedup over the NeutronSparse workload partitioning strategy, and a 1.10×–1.20× speedup over the Block-Level split strategy because these methods neglect the atomic write conflicts on output matrix. These results demonstrate that effective heterogeneous task mapping should consider both heterogeneous execution costs and write conflicts rather than workload volume alone. 4.5

Case Study: End-to-end GNN Performance

To evaluate the benefits of our proposed Ascend-SpMM, we integrate it into PyG[16] via the PyTorch C++ Extension[6, 35] interface. We compare our method against the original

PyG[16] implementation and MindSpore-GL[32] implementation on the Ascend Atlas 800I A2 NPU with the same twolayer Graph Convolutional Neural Network (GCN) model, replacing PyG’s message-passing operator with Ascend-SpMM. Due to the lack of native SpMM support on the NPU, both baselines resort to Scatter-based message passing rather than SpMM, which leads to substantial memory overhead and causes out-of-memory (OOM) failures on larger-scale graphs. We evaluate end-to-end model training time on three medium-scale datasets with a hidden dimension of 128. As shown in Figure 12, our method achieves speedups of 3.5×– 6.8× over PyG and 1.4×–2.4× over MindSpore-GL across Flickr[46], ogbn-arxiv[24], and Yelp[46]. The performance benefit widens considerably on Yelp—the largest graph of the three—where our approach attains a 6.8× speedup over PyG, demonstrating the scalability advantages of SpMM-based message passing over Scatter-based alternatives on NPUs.

5

Related Work

SpMM Research on NVIDIA GPUs. On traditional GPU CUDA cores, cuSPARSE[11] provides a widely used and highperformance baseline. Several CUDA-core methods employ one-dimensional or adaptive tiling, row caching, and row decomposition to improve memory access efficiency and load balance [18, 23, 25, 34]. With hardware advancements, TC-GNN [42] pioneered the use of Tensor Cores [10] to accelerate sparse computations, and several subsequent studies explored optimization through reordering, compact formats, and fine-grained packing [14, 36, 43, 47]. Recent studies have explored heterogeneous cooperation between CUDA cores and Tensor Cores [13, 28, 37]. However, these methods are designed for the GPU execution model and do not address Ascend’s independent AIC/AIV memory paths, cross-core coordination, or scalar-issued pipelines that require explicit synchronization. SpMM Research on Ascend NPUs. Early NPU SpMM studies targeted coupled architectures. Moustafa et al.[33] explore multi-level tiling for AIC-based SpMM, but irregular blocks still incur padding and the AIV path remains unused in the modern decoupled architecture of Ascend NPUs. Recently, NeutronSparse[2] addressed this misalignment by

Ascend-SpMM: Accelerating Sparse Matrix-Dense Matrix Multiplication on Ascend NPUs

introducing sparsity-aware workload coordination between heterogeneous engines and locality-aware tile orchestration. However, it does not optimize the kernel based on the underlying execution model of Ascend NPUs, which limits its overall performance.

6

Conclusion

This paper presents Ascend-SpMM, a systematic sparse matrixdense matrix multiplication framework tailored to the decoupled architecture of Ascend NPUs. It combines a heterogeneous task mapping strategy with specialized AIC and AIV sub-kernels to exploit decoupled computing resources while reducing output conflicts and improving latency hiding. Extensive evaluations demonstrate that Ascend-SpMM achieves 12.70–20.55× geometric-mean speedup over the strongest evaluated NPU baseline, remains competitive with GPU implementations, and improves end-to-end GCN training performance. These results demonstrate the importance of jointly considering workload structure and architectural characteristics when optimizing sparse operators on heterogeneous accelerators.

Acknowledgments In this paper, we use AI models: GPT 5.6, Gemini 3.1 Pro and Deepseek-V4-Pro, to polish writing and code debugging. All AI-generated or AI-revised content was rigorously reviewed, validated, and refined by the authors, who take full responsibility for the integrity, originality, and technical validity of the work presented.

References [1] Sergi Abadal, Akshay Jain, Robert Guirado, Jorge López-Alonso, and Eduard Alarcón. 2021. Computing graph neural networks: A survey from algorithms to accelerators. ACM Computing Surveys (CSUR) 54, 9 (2021), 1–38. [2] Xin Ai, Zeyu Ling, Hao Yuan, Qiange Wang, Yanfeng Zhang, Yutao Peng, and Ge Yu. 2026. NeutronSparse: Coordinating Heterogeneous Engines for Sparse Matrix Multiplication on NPUs. arXiv preprint arXiv:2606.22482 (2026). [3] Hartwig Anzt, Stanimire Tomov, and Jack J Dongarra. 2015. Accelerating the LOBPCG method on GPUs using a blocked sparse matrix vector product.. In SpringSim (HPS). 75–82. [4] Ascend. 2024. Atlas 800I A2. Ascend 800I A2 specification. https:// support.huawei.com/enterprise/zh/doc/EDOC1100360587/b2592910 [5] Ascend. 2024. msProf. https://www.hiascend.com/document/ detail/zh/CANNCommunityEdition/850alpha001/devaids/optool/ atlasopdev_16_0082.html [6] Ascend. 2024. PyTorch Adaptation. Ascend Extension for PyTorch. https://www.hiascend.com/en/developer/software/aiframeworks/pytorch [7] NVIDIA Corporation. 2023. Block Sparse Row (BSR). https://docs. nvidia.com/cuda/cusparse/index.html#block-sparse-row-bsr [8] NVIDIA Corporation. 2023. NVIDIA A100 | Tensor Core GPU. A100 specification. https://www.nvidia.com/content/dam/enzz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-usnvidia-1758950-r4-web.pdf

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA

[9] NVIDIA Corporation. 2023. NVIDIA Ampere GPU Architecture Tuning Guide. Version 12.x. https://docs.nvidia.com/cuda/ampere-tuningguide/index.html [10] NVIDIA Corporation. 2023. Tensor Core. Version 12.x. https://www. nvidia.com/en-us/data-center/tensor-cores/ [11] NVIDIA Corporation. 2024. cuSPARSE Library User Guide. Version 12.x. https://docs.nvidia.com/cuda/cusparse/ [12] Timothy A Davis and Yifan Hu. 2011. The University of Florida sparse matrix collection. Acm transactions on mathematical software (toms) 38, 1 (2011), 1–25. [13] Yukang Dong, Wenbin Jiang, Xinhai Shen, Haihong Guo, Zhiyuan Shao, and Hai Jin. 2025. BRP-SpMM: Block-Row Partition Based Sparse Matrix Multiplication with Tensor and CUDA Cores. In 2025 IEEE International Parallel and Distributed Processing Symposium (IPDPS). IEEE, 901–912. [14] Ruibo Fan, Wei Wang, and Xiaowen Chu. 2024. Dtc-spmm: Bridging the gap in accelerating general sparse matrix multiplication with tensor cores. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 253–267. [15] Ruibo Fan, Xiangrui Yu, Peijie Dong, Zeyu Li, Gu Gong, Qiang Wang, Wei Wang, and Xiaowen Chu. 2025. Spinfer: Leveraging low-level sparsity for efficient large language model inference on gpus. In Proceedings of the Twentieth European Conference on Computer Systems. 243–260. [16] Matthias Fey, Jinu Sunil, Akihiro Nitta, Rishi Puri, Manan Shah, Blaž Stojanovič, Ramona Bendias, Alexandria Barghi, Vid Kocijan, Zecheng Zhang, Xinwei He, Jan Eric Lenssen, and Jure Leskovec. 2025. PyG 2.0: Scalable learning on real world graphs. arXiv preprint arXiv:2507.16991 (2025). [17] Elias Frantar and Dan Alistarh. 2023. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International conference on machine learning. PMLR, 10323–10337. [18] Trevor Gale, Matei Zaharia, Cliff Young, and Erich Elsen. 2020. Sparse gpu kernels for deep learning. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 1–14. [19] Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi, Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. 2024. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410.13276 (2024). [20] John R Gilbert, Cleve Moler, and Robert Schreiber. 1992. Sparse matrices in MATLAB: Design and implementation. SIAM journal on matrix analysis and applications 13, 1 (1992), 333–356. [21] Fred G Gustavson. 1978. Two fast algorithms for sparse matrices: Multiplication and permuted transposition. ACM Transactions on Mathematical Software (TOMS) 4, 3 (1978), 250–269. [22] Mohammed Heyouni and Azeddine Essai. 2005. Matrix Krylov subspace methods for linear systems with multiple right-hand sides. Numerical Algorithms 40, 2 (2005), 137–156. [23] Changwan Hong, Aravind Sukumaran-Rajam, Israt Nisa, Kunal Singh, and Ponnuswamy Sadayappan. 2019. Adaptive sparse tiling for sparse matrix multiplication. In Proceedings of the 24th Symposium on Principles and Practice of Parallel Programming. 300–314. [24] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. 2020. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems 33 (2020), 22118–22133. [25] Guyue Huang, Guohao Dai, Yu Wang, and Huazhong Yang. 2020. Gespmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 1–12. [26] Thomas N Kipf and Max Welling. 2016. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907

PPoPP ’27, March 20–24, 2027, Salt Lake City, Utah, USA (2016). [27] Jiayu Li, Fugang Wang, Takuya Araki, and Judy Qiu. 2019. Generalized sparse matrix-matrix multiplication for vector engines and graph applications. In 2019 IEEE/ACM Workshop on Memory Centric High Performance Computing (MCHPC). IEEE, 33–42. [28] Zhonggen Li, Xiangyu Ke, Yifan Zhu, Yunjun Gao, and Yaofeng Tu. 2025. HC-SpMM: Accelerating Sparse Matrix-Matrix Multiplication for Graphs with Hybrid GPU Cores. In 2025 IEEE 41st International Conference on Data Engineering (ICDE). IEEE, 501–514. [29] Xiaoyao Liang. 2020. Ascend AI Processor Architecture and Programming: Principles and Applications of CANN. Elsevier. [30] Heng Liao, Jiajin Tu, Jing Xia, Hu Liu, Xiping Zhou, Honghui Yuan, and Yuxing Hu. 2021. Ascend: a scalable and unified architecture for ubiquitous deep neural network computing: Industry track paper. In 2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA). IEEE, 789–801. [31] Heng Liao, Jiajin Tu, Jing Xia, and Xiping Zhou. 2019. Davinci: A scalable architecture for neural network computing. In 2019 IEEE Hot Chips 31 Symposium (HCS). IEEE, 1–44. [32] Mindspore. 2022. MindSpore Graph Learning. mindgl. https://github. com/mindspore-ai/graphlearning [33] Salli Moustafa. 2023. Accelerating sparse matrix-matrix multiplication with the Ascend AI core. In Proc. Workshop Accelerated Mach. Learn. 33–41. [34] Meng Pang, Xiang Fei, Peng Qu, Youhui Zhang, and Zhaolin Li. 2024. A row decomposition-based approach for sparse matrix multiplication on GPUs. In Proceedings of the 29th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming. 377–389. [35] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems 32 (2019). [36] Jinliang Shi, Shigang Li, Youxuan Xu, Rongtian Fu, Xueying Wang, and Tong Wu. 2025. Flashsparse: Minimizing computation redundancy for fast sparse matrix multiplications on tensor cores. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming. 312–325. [37] Jinliang Shi, Shigang Li, Youxuan Xu, Xueying Wang, Rongtian Fu, Zhi Ma, and Tong Wu. 2025. Libra: Synergizing CUDA and Tensor Cores for High-Performance Sparse Matrix Multiplication. arXiv e-prints (2025), arXiv–2506. [38] Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. 2024. A simple and effective pruning approach for large language models. In International Conference on Learning Representations, Vol. 2024. 4942–4964. [39] Vasileios Titopoulos, Kosmas Alexandridis, Christodoulos Peltekis, Chrysostomos Nicopoulos, and Giorgos Dimitrakopoulos. 2025. Optimizing structured-sparse matrix multiplication in RISC-V vector processors. IEEE Trans. Comput. 74, 4 (2025), 1446–1460. [40] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. 2017. Graph attention networks. arXiv preprint arXiv:1710.10903 (2017). [41] Minjie Yu Wang. 2019. Deep graph library: Towards efficient and scalable deep learning on graphs. In ICLR workshop on representation learning on graphs and manifolds. [42] Yuke Wang, Boyuan Feng, Zheng Wang, Guyue Huang, and Yufei Ding. 2023. {TC-GNN}: Bridging sparse {GNN} computation and dense tensor cores on {GPUs}. In 2023 USENIX Annual Technical Conference (USENIX ATC 23). 149–164. [43] Yaoyu Wang, Xiao Guo, Junmin Xiao, De Chen, and Guangming Tan. 2025. {GeneralSparse}: Bridging the Gap in {SpMM} for Pruned Large Language Model Inference on {GPUs}. In 2025 USENIX Annual

Jiahe Wang et al. Technical Conference (USENIX ATC 25). 417–432. [44] Jun Wu, Jingrui He, and Jiejun Xu. 2019. Net: Degree-specific graph neural networks for node and graph classification. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining. 406–415. [45] Haojun Xia, Zhen Zheng, Yuchao Li, Donglin Zhuang, Zhongzhu Zhou, Xiafei Qiu, Yong Li, Wei Lin, and Shuaiwen Leon Song. 2023. Flash-llm: Enabling cost-effective and highly-efficient large generative model inference with unstructured sparsity. arXiv preprint arXiv:2309.10285 (2023). [46] Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. 2019. Graphsaint: Graph sampling based inductive learning method. arXiv preprint arXiv:1907.04931 (2019). [47] Haisha Zhao, San Li, Jiaheng Wang, Chunbao Zhou, Jue Wang, Zhikuang Xin, Shunde Li, Zhiqiang Liang, Zhijie Pan, Fang Liu, Yan Zeng, Yangang Wang, and Xuebin Chi. 2025. Acc-spmm: Accelerating general-purpose sparse matrix-matrix multiplication with gpu tensor cores. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming. 326–338. [48] Yuhang Zhou, Zhibin Wang, Guyue Liu, Shipeng Li, Xi Lin, Zibo Wang, Yongzhong Wang, Fuchun Wei, Jingyi Zhang, Zhiheng Hu, Yanlin Liu, Chunsheng Li, Ziyang Zhang, Yaoyuan Wang, Bin Zhou, Wanchun Dou, Guihai Chen, and Chen Tian. 2025. Squeezing operator performance potential for the ascend architecture. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2. 1156–1171. [49] Yuhang Zhou, Zibo Wang, Zhibin Wang, Ruyi Zhang, Chen Tian, Xiaoliang Wang, Wanchun Dou, Guihai Chen, Bingqiang Wang, Yonghong Tian, Yan Zhang, Hui Wang, Fuchun Wei, Boquan Sun, Jingyi Zhang, Bin She, Teng Su, Yifan Yao, Chunsheng Li, Ziyang Zhang, Yaoyuan Wang, Bin Zhou, and Guyue Liu. 2025. Accelerating Model Training on Ascend Chips: An Industrial System for Profiling, Analysis and Optimization. In 2025 USENIX Annual Technical Conference (USENIX ATC 25). 1387–1408.

Related documents

• SpSYRK: Half the Work in Distributed Sparse Matrix Multiplication#441743
• FalconGEMM: Surpassing Hardware Peaks with Lower-Complexity Matrix Multiplication#168269
• P-GADMM: Parallel Group-Based ADMM for Asynchronous Optimization in Heterogeneous Edge Networks#1033788
• P-GADMM: Parallel Group-Based ADMM for Asynchronous Optimization in Heterogeneous Edge Networks#978393
• A Neural Hierarchical-Matrix Preconditioner for Real-Time GPU Solves#1077888