Empowering Hybrid Attention Models on NPUs

1 Introduction

• At 4,096 input tokens, a Qwen3.5-2B LA layer performs 8.7× fewer multiply-accumulate operations than full attention yet runs 3.02× slower on Ascend 310P.LA consumes 75.16%–76.99% of total prefill latency across Qwen3.5 models; global-memory accesses account for 56.2% of LA latency on Ascend 310P.Source: §1 · p. 1
• LA’s default execution order generates intermediate tensors long before their final use, exceeding on-core cache capacity and causing repeated spills and reloads that inflate logical GM traffic by 93.4% for GDN and 63.8% for KDA.Source: §1 · p. 1
• HA-NPU partitions work by head and fuses dependent operators on one core, reorders operations to consume intermediates promptly, and plans layouts to limit transformation overhead.It preserves the underlying attention algorithms and their recurrent-state dependencies.Source: §1 · p. 2
• Tail-phase sub-head parallelism splits remainder heads column-wise along the non-reduction dimension to use otherwise idle cores.For example, mapping 16 heads onto 10 cores leaves six tail heads; splitting them column-wise allows all 10 cores to be used.Source: §1 · p. 2
• Lifetime-aware reordering uses chunk-first scheduling to enable in-place state updates and GEMM-first scheduling to group shared-input operations for on-core reuse.The latter decouples matrix inversions from GEMMs, reducing the time operands remain idle.Source: §1 · p. 2
• 1.69–35.95× kernel speedups and 1.76–36.14× kernel energy reductions, with up to 2.03× faster end-to-end inference and 48.0% lower device energy per request.The evaluation covers GDN and KDA kernels on two Ascend NPUs and six LLMs, compared with four competitive baselines.Source: §1 · p. 2

2.1 Linear-Attention-based LLMs

• LA reduces full attention’s O(T^2) compute and O(T) memory costs to O(T) and O(1), respectively, by maintaining a fixed-size recurrent state per head.State propagation introduces a strict serial dependency, trading sequence-level parallelism for resource efficiency.Source: §2.1 · p. 3
• Chunkwise Map–Reduce–Map execution confines sequence dependence to the Reduce phase, while both Map phases run independently across chunks.For GDN, Map precomputes chunk-local transformations, Reduce updates state, and the final Map computes token outputs.Source: §2.1 · p. 3

2.2 Edge NPUs

• NPUs combine Scalar, Vector, and Matrix units with global memory and explicitly managed on-core buffers; Ascend cores use UB, L0, and L1 buffers.On Ascend 310P, profiled UB→L1, GM→L1, and UB→GM→L1 bandwidths are 205, 95, and 48 GB/s, respectively.Source: §2.2 · p. 3

2.3 Gaps between Linear Attention and Edge NPUs

• At 512–4K tokens, LA is 3.02–4.98× slower than full attention on Ascend 310P and consumes 75.16%–80.12% of prefill time.The profiling covers Qwen3.5-2B/4B; the bottleneck is traced to data movement and mismatches at core, operator, and tensor levels.Source: §2.3 · p. 4
• Chunk-head mapping fractures tensor locality on NPUs: global-memory accesses take 56.2% of LA latency, while Matrix-unit utilization is 0.20%–0.41%.Full attention achieves 16.8%–29.0% Matrix-unit utilization in the cited comparison.Source: §2.3 · p. 4
• Limited on-core buffers cannot retain long-lived intermediates, so spilling increases logical GM traffic by 93.4% for GDN and 63.8% for KDA.The cited comparison is for Ascend 310P; Figure 4 describes ideal tensor residency versus actual spill/reload at T=2048, H=10, C=64.Source: §2.3 · p. 4, Figure 4 · p. 4
• Matrix units require blocked layouts while Vector units favor contiguous layouts, so handoffs cause transformations that consume 41.58% of cross-unit execution time and add 16.7% memory traffic.These measurements are profiling results for Qwen3.5-2B.Source: §2.3 · pp. 4–5

3 Design Overview

• HA-NPU targets edge-NPU prefill latency and energy without altering original attention semantics, and supports GDN and KDA while being extendable to RetNet and GLA.It is designed to integrate with existing frameworks such as vLLM-Ascend.Source: §3 · p. 5
• The system shifts from sequence-dimension parallelism to operator-dataflow locality across core, operator, and tensor levels.It maps and fuses head workloads to avoid cross-core GM access, reorders operators to consume intermediates promptly, and coordinates layouts between Matrix and Vector units.Source: §3 · p. 5
• Offline, HA-NPU reorders the computation graph and assigns tensor formats and transformation sites; online, it executes the plan with head-level mapping and fusion.The online stage pins each head to an AI Core and elastically splits tail execution to improve hardware saturation while preserving on-core locality.Source: §3 · p. 5

4 Head-Level Mapping with Cross-Operator Fusion

• Chunk-Head-Chunk and Chunk-Seg-Seg mappings independently schedule granularities, severing on-core dependencies and forcing intermediate tensors and states to be materialized in GM.The mappings pursue GPU-style concurrency; on bandwidth-constrained edge NPUs, the resulting traffic can starve Matrix units.Source: §4 · pp. 5–6
• HA-NPU fuses each head’s Map–Reduce–Map pipeline on one core and processes chunks sequentially, keeping intermediates and recurrent state in on-core buffers.Independent heads provide cross-head concurrency without cross-head synchronization.Source: §4 · p. 6
• When residual heads underfill cores, HA-NPU splits them column-wise rather than row-wise, exposing concurrent work without cross-core reductions.A row-wise split of the recurrent state would require costly cross-core reductions.Source: §4 · p. 6
• The scheduler maps one head per core when head parallelism suffices, splits residual heads along state columns for underfilled tails, and uses chunk-head-chunk only for extreme tails.This three-tier pipeline balances on-core reuse with core saturation.Source: §4 · pp. 6–7

5 Lifetime-Aware Operator Reordering

• Deferring recurrent-state updates until the old state is last read enables in-place overwrite, halving state memory while preserving the algorithmic recurrence.The chunk-first order splits the update into a state read, chunk output, and state update.Source: §5 · p. 7, Figure 10 · p. 7
• A 244-KiB kernel footprint plus a 64-KiB state exceeds the Ascend 310P’s 256-KB buffer; immediate consumption of intermediates limits peak memory.Reordering targets cross-chunk state updates and within-chunk shared-input GEMMs.Source: §5 · p. 7
• Factoring A^-1 from W = A^-1K decouples a GEMM from inversion, allowing shared-input GEMMs to run first and reuse L0-buffered inputs.The grouped GEMMs are K S, Q S, K K^T, and Q K^T; the inputs are then released for inversion workspace.Source: §5 · p. 7

6 Dataflow-Aware Layout Transformation Planning

• Layout planning minimizes total inference latency by balancing layout-specific execution costs against transformation overheads for matrix and vector units.Matrix units favor Blocked layouts, while vector units favor contiguous Dense layouts.Source: §6 · p. 7
• Cross-chunk state propagation, shared tensors in branches, and long vector chains couple layout choices, making general PBQP search NP-hard and naive enumeration intractable for long contexts.The search space grows as 2^(m*N_c), where N_c scales linearly with input sequence length.Source: §6 · p. 8
• HA-NPU prunes layouts using chunk isolation, fixed Blocked matrix interfaces, and two candidate layouts for vector chains, then applies dynamic programming in O(N_c) time.The planner backtracks from recorded predecessor states to assign tensor layouts and transformation sites, yielding a static plan without runtime layout-scheduling overhead.Source: §6 · p. 8

7 Implementation and Evaluation

• HA-NPU is a general edge-NPU inference system; its implementation is a roughly 13K-line C/C++ kernel-level drop-in for vLLM-Ascend.Projection fusion and convolution-metadata reuse provide a 1.62× speedup over native vLLM-Ascend, with that backend shared across baselines.Source: §7 · pp. 8–9, Figure 12 · p. 9

7.1 Experiment Setups

• Evaluation covers GDN and KDA across six models, two Ascend NPU devices, and four baselines, with shared projection and decoding across methods.The baselines include native vLLM-Ascend, FLA, FLA-MegaKernel, and FlashQLA; experiments average 3–5 runs.Source: §7.1 · p. 9

7.2 Linear-Attention Kernel Performance

• At 4,096 tokens, HA-NPU is 11.60–19.28× faster than vLLM-Ascend on Ascend 310P.It also reduces energy by 11.92–18.50× versus vLLM-Ascend for 4,096-token sequences on Ascend 310P.Source: §7.2 · pp. 9–10, Figure 13 · p. 10, Figure 14 · p. 10

7.3 End-to-End Performance

• Across 18 model-dataset configurations, HA-NPU achieves the lowest end-to-end latency, with up to 2.03x speedup over vLLM-Ascend on Qwen3.5 models.It reduces total device energy per request by 32.8–48.0% versus vLLM-Ascend.Source: §7.3 · p. 11, Figure 15 · p. 11

7.4 Accuracy

• HA-NPU preserves model output quality, with negligible F1 deviations (≤0.25) on HotpotQA and TriviaQA versus vLLM-Ascend.Across 360 GDN instances, maximum NMSEs are 5.23×10−8 for outputs and 9.38 × 10−12 for final states.Source: §7.4 · p. 11, Table 2 · p. 11

7.5 Peak Device Memory

• HA-NPU reduces peak memory versus FLA by 31.7–52.6% for GDN and 38.7–54.6% for KDA.At 65,536 tokens, its peak memory is 6.0% lower for GDN and 22.8% lower for KDA than FlashQLA.Source: Figure 16 · p. 11, §7.5 · p. 12

7.6 Memory Traffic and Matrix Utilization

• Against FLA, HA-NPU reduces GM traffic by 65.5–84.8% and L1-to-L0 transfers by 30.5–46.4%.Matrix utilization reaches 5.33–7.27% for GDN and 4.27–6.38% for KDA, up to 5.24× FLA and 1.74× FlashQLA.Source: §7.6 · p. 12, Figure 17 · p. 12

7.7 Ablation Study

• Starting from FLA, head-level mapping gives the largest initial speedups: 2.86× for GDN and 3.15× for KDA.Reordering further cuts latency by 23.8% and 19.7%, and layout planning trims it by 15.1% and 11.2%, respectively.Source: §7.7 · p. 12, Figure 18 · p. 12

LLM Inference on Edge NPUs

• Prior edge-NPU LLM work includes fixed-size chunking, CPU-NPU pipelining, sparsity-guided token selection, and table-lookup low-bit inference.These studies target standard dense or sparse architectures, whereas HA-NPU co-designs LA dataflow for edge NPUs.Source: §8 · p. 12

Linear Attention in LLMs

• RetNet and GLA introduce temporal decay to control memory retention, while DeltaNet, GDN, and KDA are also identified among recent linear-attention works [54] [64].Source: §8 · p. 12

Linear Attention on GPU/NPU

• GPU LA systems use chunkwise parallelism, enhanced concurrency, and warp specialization, while basic NPU paths and ported GPU designs fail to exploit memory-constrained edge devices.HA-NPU instead co-designs LA execution with edge-NPU hardware.Source: §8 · p. 13

• HA-NPU combines head-level mapping, lifetime-aware operator reordering, and layout planning to improve data locality for edge-NPU linear-attention inference.Source: §9 · p. 13
• Experiments on GDN and KDA demonstrate the effectiveness of HA-NPU for edge-NPU linear-attention inference.Source: §9 · p. 13
• Up to 2.03× end-to-end speedup and 48.0% lower device energy are reported for HA-NPU.Source: §9 · p. 13