Vectorization Of Narrow Matrix Multiplication for Ascend AI Inference Acceleration

I. INTRODUCTION

• Narrow MatMul can leave the Cube Unit underutilized while the Vector Unit remains idle, motivating vector-unit offloading.The authors report this observation from operator profiling on an architectural simulator.Source: §I · p. 1
• Offloading narrow MatMul to AIV can let independent later matrix calculations start earlier on AIC, overlapping work across the units.The algorithm uses vector instructions available in the AscendC API.Source: §I · p. 1
• 20% mean performance gain was reported for single-token MLA DeepSeek-V3 inference.The paper says the result was confirmed by real-hardware measurements.Source: §I · p. 1

A. Hardware Architecture

• Ascend AI Core includes Cube, Vector, and Scalar resources, with Cube performing matrix-related operations and Vector handling vector operations.The Scalar Unit handles scalar-data operations and program control flow.Source: A. Hardware Architecture · pp. 1–2

B. Cube Unit

• The Cube Unit provides parallel multiplications and additions, with throughput of one 16 × 16 multiplication per cycle on Ascend 910B.The passage lists support for INT4, INT8, FP16, BF16, and FP32 at the time of the research.Source: B. Cube Unit · p. 2

C. Vector Unit

• Each chip has two independent Vector Units, each with a 192-kilobyte Unified Buffer, and each instruction processes at most 256 consecutive bits.The passage notes that basic arithmetic operations are not supported for BF16.Source: C. Vector Unit · p. 2

D. Matrix Tiling

• Matrix tiling divides tensors into blocks to reuse data and distribute computation across AI cores; AscendC uses FRACTAL_ZZ for A, FRACTAL_ZN for B, and FRACTAL_NZ for C.The tiling figure presents one possible partitioning of A × B = C on the Cube Unit.Source: D. Matrix Tiling · p. 2, Figure 2 · p. 2

E. Problem Statement

• 16× lower Cube efficiency can result in matrix-vector multiplication because padding leaves 15 of a 16 × 16 block idle.The paper identifies small tile remnants and excessive padding as motivations for considering alternative narrow-MatMul computation.Source: E. Problem Statement · p. 2

III. VECTOR MATMUL OPTIMIZATION

• Offloading is worthwhile only when AIV work can overlap with subsequent AIC computations; a single Vector Unit is estimated to be about 64 times slower than a single Cube Unit.The paper warns that naive offloading is not beneficial.Source: §III · p. 2

A. Optimization Scope

• Offloading requires independent subsequent Cube computations, a vector dimension with min(M, N) = 1, and approximately 2P idle Vector Units when P AI cores execute the operator.These are the paper's data-independence, narrow-dimension, and idling conditions.Source: A. Optimization Scope · pp. 2–3

B. Proposed Algorithm

• The AIV algorithm initializes by loading matrix blocks, computes vector products and partial sums, then synchronizes Vector Units to gather and accumulate the results.During computation, it splits loaded blocks to fit the Vector Unit instruction limit and stores partial sums in a temporary buffer.Source: B. Proposed Algorithm · p. 3

• The optimization was successfully applied on Ascend NPU to DeepSeek-V3 inference for single-token processing.Source: §IV · p. 3
• Each AIV executor loads its assigned blocks of A and B from Global Memory into the Unified Buffer during initialization.Source: §IV · p. 3
• The compute stage multiplies partitioned A and B blocks, reduces each product, and stores partial sums in a temporary buffer.Source: §IV · p. 3
• The gather stage waits for vector work, accumulates stored partial sums, and writes the resulting tensor block to C.Source: §IV · p. 3

A. MLA Operator Optimization

• MLA reduces the KV cache footprint during inference, which is critical for efficient long-context processing [14] [15].Source: A. MLA Operator Optimization · p. 4
• The target multiplications MM[cKV] and MM[kR] are combined into one operation called MatMul 2.Source: A. MLA Operator Optimization · p. 4
• Cube calculations do not depend on MatMul 2, while its results feed Vector Unit operations such as RMSNorm and RoPE.Source: A. MLA Operator Optimization · p. 4
• Successful AIV/AIC overlap could let RMS, RoPE, and caching normalization start after MatMul 2 without Cube synchronization.Source: A. MLA Operator Optimization · p. 4
• The optimization targets processing one token, or batch size = 1.Source: A. MLA Operator Optimization · p. 4

B. Algorithm Specialization

• The target case specializes MatMul for one token, with hidden dimension d=7168, KV compression dimension dc=512, and per-head h=64.The specialization concerns multiplication by W_DKV and W_KR.Source: B. Algorithm Specialization · p. 4

C. Data Layout Issue

• FRACTAL_NZ weights are arranged in 16×16 blocks, so the weight matrix is transposed in fractalSize-sized blocks to obtain the row-column multiplication layout.The algorithm requires the weight matrix to be transposed before applying the proposed row-column multiplication.Source: C. Data Layout Issue · p. 4
• AscendC Transpose does not support bfloat16_t, so the loaded memory block is reinterpreted as half before transposition.The resulting fractal column is transposed at size fractalSize · kSpitSize.Source: C. Data Layout Issue · p. 4
• After transposition, the required kSplitSize column is scattered through the UB buffer with a block stride of fractalSize.Cast repeat-stride and block-stride parameters are chosen to address this layout for vector multiplication.Source: C. Data Layout Issue · p. 5

D. Data Type Issue

• At the compute stage, the weight matrix and token row buffers are converted from bfloat16_t to float using RoundMode::CAST_NONE.The conversion applies to buffers of sizes fractalSize · kSpitSize and kSpitSize, respectively.Source: D. Data Type Issue · p. 5
• Cast gathers one weight-matrix column of size kSplitSize while converting its elements to the required format.The column is collected over kSplitSize mask (opNum) repeats.Source: D. Data Type Issue · p. 5
• Partial sums are stored in Global Memory as float, then the gathered result is cast back to bfloat16_t with RoundMode::CAST_ROUND.The final cast is specified to match the accuracy obtained on the Cube Unit.Source: D. Data Type Issue · p. 5

E. MLA Operator MatMul Algorithm

• The target MLA calculation assumes P=24 physical AI Cores, and the specialized algorithm casts the accumulation buffer back to bfloat16_t to match Cube accuracy and subsequent calculations.The subsequent calculations named are RoPE and RMS Norm.Source: E. MLA Operator MatMul Algorithm · pp. 5–6
• The empirically chosen kSplitSize overlaps next-iteration weight loading via the MTE2 bus with current-iteration calculations.The AIV pool size V is the maximum number of vector executors into which the weight matrix can be evenly divided into fractalNum · kSplitSize blocks; its optimal value minimizes nPartW.Source: E. MLA Operator MatMul Algorithm · pp. 5–6
• Each vector executor loads its token block from GM to UB only once, and each weight-column multiplication and reduction is represented as a dot product.These are listed as optimizations to the specialized algorithm.Source: E. MLA Operator MatMul Algorithm · p. 6
• Because T is a row, fractalNum vector executors gather the result matrix, which the algorithm notes is acceptable because fractalNum < V.The gather step is assigned to a sufficient number of vector executors.Source: E. MLA Operator MatMul Algorithm · p. 6

F. Evaluation

• 19% Vector Unit performance improvement and 17% Cube Unit performance improvement were measured in the timing-simulator profiles, alongside a 71% increase in MTE3 instruction count without contention.The AIC still completes the operator’s final computations.Source: F. Evaluation · p. 6
• Real-hardware execution-time distributions were measured on Ascend 910B over 100 runs of 100 iterations, dropping each run’s first iteration.The first iteration was dropped to eliminate cold-start effects such as non-warmed-up caches.Source: F. Evaluation · p. 6

V. DISCUSSION AND FUTURE WORK

• ≈150 KB of Unified Buffer per AIV core can matter when that buffer already holds Cube-needed data; MTE pressure may slow Cube transfers, and formats or layouts can add work.The added work may include redundant casts, DataCopy operations, or transposes.Source: §V · p. 6
• The specialized algorithm maps bfloat16 T and W to bfloat16 C, divides K into 1024-element chunks and N into 16-wide fractals, then gathers partial sums.It accumulates in float before casting the output to bfloat16.Source: §V · p. 6
• Future offloading is proposed for the narrow remainder when dmin mod 16 = 1 and dmin ≠ 1, with broader applicability hypothesized for certain data types and memory layouts.The broader hypothesis depends on extending Condition III.2.Source: §V · p. 7

VI. CONCLUSION

• Offloading matrix-vector multiplication from the Cube Unit to the Vector Unit lets its work overlap with subsequent Cube Unit computations.Source: §VI · p. 7
• 20% mean performance gain for single-token processing on the DeepSeek-V3 MLA operator was confirmed by real hardware measurements.Source: §VI · p. 7
• The authors focus on server NPUs and suggest extending the small-token multiplication optimization to similar operators on other NPU chips.Source: §VI · p. 7