1 of 23

οΏ½NeuronMM: High-Performance οΏ½Matrix MultiplicationοΏ½for LLM Inference on AWS TrainiumοΏ½οΏ½

Dinghong Song

University of California, Merced

PASA Lab

PASA Lab

2 of 23

Introduction

  • As the capabilities of LLMs improve, the model sizes continue to grow accordingly.

​

  • Large-scale model sizes impose significant barriers to practical application and deployment, especially in resource-constrained environments

​

  • These challenges have motivated extensive research into model compression and efficient hardware.

​

    • Singular Value Decomposition (SVD) has recently proven effective for compressing the weight matrices of LLMs.

2

PASA Lab

3 of 23

Introduction

  • Emerging AI accelerators are attracting growing attention, as they offer new opportunities for efficient LLM inference.

​

  • AWS Trainium, a specialized chip for deep learning workloads.
  • EachΒ chip consists of:
    • Two NeuronCore
    • 32 GiB HBM
    • DMA
    • NeuronLink
    • Host PCIe

​

3

PASA Lab

4 of 23

Introduction

  • NeuronCore
    • 4 main engines
      • Tensor Engine: matrix multiplication (2D systolic array)
      • Vector Engine: vector computation
      • Scalar Engine: element-wise operations
      • GPSIMD Engine: Β General Purpose SIMDΒ 
    • SRAM memory
      • State Buffer (SBUF): main on-chip software-managed memory
      • Partial Sum Buffer (PSUM): a dedicated accumulation buffer

for the tensor engine

4

PASA Lab

5 of 23

Introduction

  • Matmul on Tensor Engine
      • load matrices from HBM into SBUF
      • divide large matrix into tiles
      • multiply two tiles A (128, 128) and B(128, 512)
        • matmul(stationary = 𝐴𝑇 , moving = 𝐡)
      • accumulate intermediate results to PSUM
      • the final result is written back to HBM

5

PASA Lab

6 of 23

Motivation

  • Challenge 1: I/O Bottleneck
      • Directly computing on the SVD-compressed weight matrices sequentially leads to low on-chip SBUF utilization and reduced Tensor Engine efficiency.
      • Frequent idle periods in the Estimated MFU (Model FLOPs Utilization) indicate that Tensor Engine is underutilized while waiting for data transfers and data preparation to complete.

​

6

PASA Lab

7 of 23

Motivation

  • Β 

7

PASA Lab

8 of 23

Motivation

  • Challenge 3: Transpose Overhead
    • tensor engine’s systolic array architecture requires stationary matrix in a matmul to be provided in a transposed layout
      • matmul(stationary = 𝐴𝑇 , moving = 𝐡)
    • Two forms of overhead: I/O transposes during data loading to SBUF and transposes of the intermediate result.
      • Y = matmul(stationary = X𝑇 , moving = U)
      • O = matmul(stationary = Y𝑇 , moving = V)

​

8

PASA Lab

9 of 23

NeuronMM

  • Compress large weight-matrices of MLP layers via block-aligned SVD
  • Restore model accuracy with a LoRA-based fine-tuning
  • Implement high-performance Neuron Kernel Interface (NKI) kernels to efficiently execute matmul in compressed LLMs.

9

PASA Lab

10 of 23

Block-Aligned SVD

  • Β 

10

PASA Lab

11 of 23

Block-Aligned SVD

  • Β 

11

PASA Lab

12 of 23

LoRA Fine-Tuning

  • Β 

12

PASA Lab

13 of 23

XUV NKI Kernel

  • caching : cache entire row of the intermediate matrix π‘‹π‘ˆ in a dedicated buffer.
  • Implicit transposition: swapping the order of the inputs to obtain the transposed intermediate result.
  • Blocking: compute π‘‹π‘ˆπ‘‰ by blocks.

​

13

PASA Lab

14 of 23

MLP NKI Kernel

  • Β 

14

PASA Lab

15 of 23

Evaluation

  • Implementation
      • Hardware: a trn1.2xlarge instance of AWS Trainium on Amazon Elastic Compute Cloud (Amazon EC2)
      • Software: leverage Neuron Kernel Interface (NKI) to develop NeuronMM and integrate it into the NeuronX Distributed Inference library
      • Only replace the MLP layer of LLM with NKI kernel and leave the rest of model unchanged.
  • Models: Llama-3.2-1B, Llama-3.2-3B, Qwen3-1.7B, and Qwen3-4B.
  • Datasets:
    • 3 language modeling datasets (WikiText-2, PTB, C4)
    • 6 common sense reasoning datasets (OpenBookQA, WinoGrande, HellaSwag, PIQA, ARC-e, and ARC-c).

​

15

PASA Lab

16 of 23

Evaluation of π‘‹π‘ˆπ‘‰ kernel

  • Β 

16

PASA Lab

17 of 23

Evaluation of π‘‹π‘ˆπ‘‰ kernel

  • Execution time and Memory Footprint:
    • The highest speedup is 2.2Γ—, which is driven by a 4.8Γ— reduction in HBM memory footprint

​

​

17

PASA Lab

18 of 23

Evaluation of π‘‹π‘ˆπ‘‰ kernel

  • Tensor engine active time and Model Float Utilization (MFU):
    • NeuronMM exhibits long tensor engine active time (defined in terms of the percentage of overall time) and MFU
    • Most cycles in the tensor engine are actively engaged in actual matmul rather than transpose, leading to a lower total execution time.

18

PASA Lab

19 of 23

Evaluation of MLP Kernel

  • Β 

19

PASA Lab

20 of 23

Impact of Block Size on Kernel Performance

  • Larger blocks reduce the latency by increasing Arithmetic Intensity.
  • Exceeding the 24 MiB SBUF capacity triggers costly memory spills to

HBM and degrade performance.

20

PASA Lab

21 of 23

Case Study

  • Llama 3.2-1B: input/output length 1K
  • End-to-end latency is significantly reduced 1.86x (from 41.22s to 22.14s)
  • Throughputs increase from 49.69 to 92.52 tokens/s.

21

PASA Lab

22 of 23

Conclusion

  • NeuronMM applies SVD and LoRA to decompose LLM weights into low-rank matrices, reducing parameters while preserving accuracy.
  • NeuronMM exploits Trainium’s Tensor Engine and on-chip memories (PSUM and SBUF) to minimize data transfers between off-chip high-bandwidth memory (HBM) and on-chip SRAM.
  • Implemented with NKI, NeuronMM achieves up to 1.67Γ— and 2.49Γ— speedups over standard inference on Qwen-3-4B and Llama-3-3B, respectively.

​

22

PASA Lab

23 of 23

οΏ½οΏ½οΏ½Thank you!

Dinghong Song

University of California, Merced

PASA Lab

PASA Lab