Empowering Hybrid Attention Models on NPUs

Yinyuan Zhang; Daliang Xu; Xiaolong Huang; Wangsong Yin; Yun Ma; Mengwei Xu; Gang Huang

cs.DCSep 26, 2026

Empowering Hybrid Attention Models on NPUs

Authors: Yinyuan Zhang, Daliang Xu, Xiaolong Huang, Wangsong Yin, Yun Ma, Mengwei Xu, Gang Huang

Organizations: Peking University · Beijing University of Posts and Telecommunications

Abstract

Hybrid attention models have emerged as a crucial architecture for Large Language Models (LLMs) (e.g., the Qwen3.5 and Kimi series). Their memory and computational efficiency make them highly attractive for on-device inference, forming a promising synergy with edge Neural Processing Units (NPUs). However, naive execution of these hybrid models on edge NPUs fails to deliver these benefits, often bottlenecking the prefill stage due to severe memory-system inefficiencies and architectural mismatches within the linear attention (LA) layers. We present HA-NPU, the first system to enable efficient hybrid attention LLM inference on edge NPUs without modifying the underlying algorithms. HA-NPU enhances execution efficiency by reorganizing the dataflow of the LA components across three levels: (1) At the core level, it partitions workloads by the head dimension and fuses dependent operators, eliminating cross-core global memory accesses; (2) At the operator level, it reorders execution to consume intermediate tensors immediately, drastically minimizing local-buffer pressure; (3) At the tensor level, it employs dataflow-aware layout planning to minimize transformation overhead between matrix and vector processing units. Compared to competitive baselines, HA-NPU achieves up to 35.95 $\times$ LA kernel speedup and 36.14 $\times$ energy reduction, delivering up to 2.03 $\times$ faster end-to-end request latency. The source code will be made publicly available at https://github.com/yinyuanzhang/HA-NPU

Figures & tables

Panel of Figure 1 — Figure 1. Architecture and chunkwise execution of the Qwen3.5 (GDN) hybrid linear-attention LLM.

Figure 2. Representative Edge NPU architecture (Ascend 310 series).

Panel of Figure 3 — Figure 3. Prefill performance on Ascend 310P.

	Full Attention			Linear Attention
Tokens	MACs [-1pt] (GMAC)	Cube [-1pt] util.	K/V cache [-1pt] (MiB)	MACs [-1pt] (GMAC)	Cube [-1pt] util.	State [-1pt] (MiB)
512	1.08	16.8%	1.0	0.75	0.20%	1.0
1,024	4.33	22.0%	2.0	1.51	0.29%	1.0
2,048	17.31	26.7%	4.0	3.02	0.37%	1.0
4,096	69.25	29.0%	8.0	6.04	0.41%	1.0

Table 1. Matrix processing unit execution and per-layer history memory of native Full- and Linear-Attention cores in Qwen3.5-2B prefill on one Ascend 310P.

Panel of Figure 4 — Figure 4. Comparison of memory traffic between ideal tensor residency and actual spill/reload on Ascend 310P ( $T=2048$ , $H=10$ , $C=64$ ).

Panel of Figure 5 — Figure 5. Dense and Blocked memory layouts of a $128\times 128$ FP16 matrix. Red arrows show element order in memory.

Figure 7. Illustration of Chunk-Head-Chunk, Chunk-Seg-Seg, and Head-Head-Head execution mappings of chunkwise LA.

Panel of Figure 8 — Figure 8. Comparison of workload partitioning strategies.

Figure 9. Workflow of HA-NPU’s three-tiered scheduling.

Panel of Figure 10 — Figure 10. Operator (Op) decoupling and reordering for in-place state updates.

Panel of Figure 11 — Figure 11. Operator rewriting and reordering enable GEMM-first execution for on-core tensor reuse.

Figure 12. Illustration of dataflow-aware layout planning.

Panel of Figure 13 — Figure 13. Forward latency of GDN and KDA under different sequence lengths and head counts on Ascend 310P and 310B. Green annotations show the range of speedups achieved by HA-NPU over the baselines for each configuration.

Figure 14. Energy consumption of GDN and KDA under different sequence lengths and head counts on Ascend 310P.

Panel of Figure 15 — Figure 15. End-to-end request latency and measured device energy per request on Ascend 310P.

	Qwen3.5-2B		Qwen3.5-4B
Method	HotpotQA	TriviaQA	HotpotQA	TriviaQA
vLLM-Ascend	48.90	83.70	65.39	89.17
HA-NPU	49.06	83.53	65.14	89.17

Table 2. Model output quality on Ascend 310P (F1

\uparrow

Figure 16. Peak memory consumption of GDN and KDA on Ascend 310P (

H=32

Panel of Figure 17 — Figure 17. Memory traffic and Matrix utilization of GDN and KDA on Ascend 310P ( $H=32$ ).

Figure 18. Ablation study of HA-NPU on Ascend 310P (

T=1{,}024

;

H=16

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,090 papers

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU
Jul 10, 2026Victor J. B. Jung, Gagandeep Singh, Joseph Melber +3Efficient Transformer Inference Energy-Efficient ML
Integrated Systems Laboratory (IIS), ETH Zürich, Switzerland. · AMD Research and Advanced Development (RAD). · Department of Electrical, Electronic and Information Engineering (DEI), University of Bologna, Italy.
When NPUs Are Not Always Faster: A Stage-Level Analysis of Mobile LLM Inference
May 22, 2026Pu Li, Jiawen Qi, Qinyu ChenAI Accelerator Inference On-Device Language Model Inference
Leiden Institute of Advanced Computer Science (LIACS), Leiden University, The Netherlands.
Is Your NPU Ready for LLMs? Dissecting the Hidden Efficiency Bottlenecks in Mobile LLM Inference
Jul 6, 2026Guanyu Cai, Ruiming Tian, Lang Yang +4LLM Inference Efficiency AI Accelerator Inference
Tsinghua University · Beijing Jiaotong University

Jul 10, 2026J/K move · Enter open · S save

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU

Victor J. B. Jung, Gagandeep Singh, Joseph Melber, Kristof Denolf+2

Integrated Systems Laboratory (IIS), ETH Zürich, Switzerland. · AMD Research and Advanced Development (RAD). · Department of Electrical, Electronic and Information Engineering (DEI), University of Bologna, Italy.

Efficient Transformer Inference Energy-Efficient ML Efficient Attention On-Device Inference

Open paper AskPDF arXiv

The growing adoption of large language model-based agents within operating system workflows has increased the importance of energy-efficient inference on laptop-class systems-on-chip (SoCs). While cloud offloading remains common, it introduces reliability and privacy concerns that are particularly problematic for agentic workloads. Recent laptop SoCs, therefore, incorporate neural processing engines (NPUs) optimized for energy efficiency; however, effectively mapping attention mechanisms onto NPUs remains challenging due to architectural diversity and explicit data-movement programming models. In this work, we present STEEL, the first open-source implementation of FlashAttention targeting XDNA-like NPUs. STEEL introduces a dataflow formulation of prefill attention, enabling efficient exploitation of spatial parallelism and on-chip memory. Furthermore, STEEL addresses the load imbalance induced by the causal mask by leveraging a sparsity-aware pipeline placement onto the NPU array, reducing synchronization overhead and improving utilization. We evaluate STEEL on the AMD Ryzen AI 9 HX 370 SoC and compare its performance against optimized CPU and GPU implementations. Experimental results show that STEEL reduces energy consumption by an average of 9.17x and 1.75x relative to CPU and GPU baselines, respectively. On XDNA 1, STEEL achieves an average 9.6x latency reduction over the prior state of the art, and delivers a 22.8x speedup on average compared to a layer-by-layer attention implementation on XDNA 2.

Abstract

LLM Inference Acceleration

Explore similar work

STEEL: Sparsity-Aware Fused Attention for Energy-Efficient Long-Sequence Inference on AMD's XDNA NPU

When NPUs Are Not Always Faster: A Stage-Level Analysis of Mobile LLM Inference

Is Your NPU Ready for LLMs? Dissecting the Hidden Efficiency Bottlenecks in Mobile LLM Inference