AWS Trainium Accelerator Kernels Overview

• AWS Trainium accelerator kernels are specialized low-level operators designed for high-performance execution of LLM and deep neural network workloads.
• They employ systematic autonomous optimizations via agentic LLM workflows, integrating techniques like loop tiling, fusion, and vector engine utilization for significant performance gains.
• Benchmark suites such as NKIBench demonstrate improvements from 45% to 71% peak throughput on Trainium 1 and 2, highlighting cost-effective, scalable AI acceleration.

AWS Trainium accelerator kernels are specialized low-level operators designed for execution on AWS Trainium, a family of AI accelerators architected by Amazon Web Services for high-performance training and inference of LLMs and related deep neural networks. The properties, optimization strategies, and benchmarking of these kernels are critical for maximizing throughput and cost-effectiveness. Recent research has established systematic approaches for their analysis and autonomous optimization, particularly in the context of agentic LLM-guided systems such as AccelOpt and the construction of comprehensive kernel benchmark suites such as NKIBench (Zhang et al., 19 Nov 2025).

1. AWS Trainium Architecture and Kernel Types

AWS Trainium exists in at least two generations, denoted “Trainium 1” and “Trainium 2”, each exposing a heterogeneous architecture on a per-core basis. Each core contains:

• Tensor Engine: optimized for large matrix multiplications, operating on 128×128 (and similar) tiles.
• Vector Engine: responsible for elementwise, reduction, and non-matmul floating-point operations.
• Scalar Engine.
• SBUF: an on-chip kernel-managed scratch buffer (~192 KB per partition, up to 128 partitions) interfacing with high-bandwidth memory (HBM) via DMA.

Performance characteristics for each generation are summarized as:

Metric	Trainium 1 (T1)	Trainium 2 (T2)
Peak HBM Bandwidth (GB/s)	440.2	640.0
FP32 Tensor Peak (TFLOPS)	23.75	19.75
FP32 Vec+Scalar Peak (GFLOPS)	286.8	550.0

Kernels are commonly categorized by workload patterns. The principal operators extracted from real LLM workloads include:

• GEMM (General Matrix-Matrix Multiply): e.g., $m=4096$, $k=5120$, $n=12288$, FP32; FLOPs $=2 m n k$.
• BatchMatMul, often with batch $B=16$, and lower $K$ dimensions.
• Fused BatchMatMul+Softmax: matrix multiply fused with softmax, utilizing both tensor and vector engines.
• Grouped Query Attention and LoRA update kernels with structured, high-dimensional access patterns.
• Mamba block kernels: vector/mesh-bound operators with nontrivial memory and compute balance.
• Others: Add + RMSNorm + MatMul, RoPE, SiLU/SwiGLU activations, Transpose+MatMul, AdamW.

Each kernel records parameter shapes, dtypes (primarily FP32), and binding information to match Trainium’s hardware constraints.

2. Benchmarking and the Construction of NKIBench

The NKIBench suite was established by extracting 14 representative kernels from high-impact LLM repositories including DeepSeek-MoE, Qwen3, Falcon, and Mamba. For each operator invocation, metadata is recorded:

• GEMM dimensions $(m, n, k)$, convolution-like window sizes $(H, W)$, channel counts, data types.
• Each benchmark is tailored to stress different hardware components and expose representative performance bottlenecks.

A roofline throughput model is used to estimate theoretical upper bounds, with core quantities:

• $$\text{Throughput}~\Phi = \frac{\text{FLOPs}_{\text{executed}}}{\text{execution\_time}}$$
• $$\text{Traffic}_{\text{Min}} = \sum_{\text{tensors}} |\text{bytes}|$$
• $$\text{T}_{\text{peak}} = \max \left(\frac{\text{Traffic}_{\text{Min}}}{\text{Bandwidth}}, \frac{\text{FLOPs}_{\text{MM}}}{\text{Peak}_{\text{MM}}}, \frac{\text{FLOPs}_{\text{Vec}}}{\text{Peak}_{\text{Vec}}}\right)$$
• $$\%~\text{of peak} = \text{T}_{\text{peak}} / t_{\text{measured}}$$

Performance is assessed by absolute throughput, percentage of peak, and improvement ratio $r = \Phi_{\text{opt}} / \Phi_{\text{base}} - 1$.

3. Autonomous Kernel Optimization via AccelOpt

AccelOpt is an LLM-based, agentic optimizer that iteratively refines AWS Trainium kernels without requiring explicit hardware knowledge, leveraging experiential memory and prompt-driven planning (Zhang et al., 19 Nov 2025). The system is organized as a loop over candidate kernels, driven by a trio of LLM agents (Planner, Executor, Summarizer) and an experience memory:

AccelOpt Agentic Loop (excerpted pseudocode):

Input: E_{i−1} = optimization memory, C_i = B candidate kernels, θ_p, θ_e, θ_s = LLM weights for Planner, Executor, Summarizer r() = profiler, σ() = memory curation, β() = candidate selector For each c in C_i: P ← {plan p ~ θ_p(p | c, E_{i−1})} # N plans per kernel For each p in P: A_p ← { (a, p, r(a)) | a ~ θ_e(a | p, c) } # Exec attempts K ← K ∪ A_p E_i ← σ(K, E_{i−1}; θ_s) # Summarization & memory update C_{i+1} ← β(K ∪ C_i, B) # Top B for next round

• Planner: Receives kernel code, profiling results, and prior memory to propose atomic optimizations (e.g., loop tiling, fusion, reordering).
• Executor: Applies plans as code rewrites, compiles, profiles, and performs correctness validation.
• Summarizer: Detects slow–fast kernel pairs with significant speedup and distills transformations into generalizable rules.
• Memory store: FIFO buffer holding {slow_code, fast_code, summary}, enabling recurrent learning and avoidance of previously suboptimal schedules.

This methodology allows unsupervised, iterative improvement, guided only by empirical performance feedback and memory-augmented LLM prompt conditioning.

4. LLM Prompt Engineering and Hyperparameterization

Prompt templates are designed to elicit targeted, atomic transformations from LLMs. Examples include:

• Planner prompt (template):

You are given a problem and baseline kernel. Profile: {profile_metrics} Baseline code: ```{kernel_code}``` Propose one 1‐step optimization plan: loop‐tile, reorder, fusion, tiling, etc.

• Summarizer prompt (template):

Slow kernel:

{slow} Fast kernel: {fast} Speedup: {r:.2f}× [Summarize](https://www.emergentmind.com/topics/jetson-nano-s-summarize) the single‐step rewrite that converts slow→fast.

Empirically tuned hyperparameters for optimization runs on Trainium 1 & 2:

Parameter	Value
Beam width B	6 candidates
# Plans/candidate (N)	12
# Executor tries (K)	2 (or 4 for hard cases)
Iterations (T)	16
Memory capacity (ExpN)	16
Memory update batch (TopK)	8
LLM temperatures	1.0 (open-source)
Token budgets (prop. LLMs)	up to 20K in / 10K out

Prompts leverage few-shot exemplars extracted from actual kernel optimization scenarios, enhancing model specificity and generalization within the agentic workflow.

5. Performance Results and Cost Analysis

Experimental evaluations with AccelOpt on NKIBench demonstrate:

• On Trainium 1: average percentage of peak throughput improved from 49% (baseline) to 61% (optimized).
• On Trainium 2: improvement from 45% to 59%.

Granular per-kernel-class improvements, with representative figures:

Kernel Class	Base (T1)	Opt (T1)	Base (T2)	Opt (T2)
GEMM	52%	68%	50%	66%
BatchMatMul+Softmax	45%	62%	42%	59%
GroupQueryAttention	49%	64%	47%	61%
LoRA	55%	71%	53%	69%
Mamba block	48%	60%	46%	58%

Cost-effectiveness is benchmarked by LLM-inference expenditures, normalized by kernel-hour:

• Open-source LLM (gpt-oss-120B): $0.15/\$$0.60 per million tokens.
• Proprietary LLM (Claude Sonnet 4): $3/\$$15 per million tokens.

At matched throughput targets, open-source AccelOpt achieves comparable kernel speedups while being 26× less expensive in LLM-inference cost.

6. Key Optimization Insights and Limitations

• Memory tiling: Aligning loop tiling strategies (e.g., 128×128, 128×512) with SBUF/PSUM constraints yields up to 10–15% performance gains.
• Vector engine utilization: Fusing elementwise operations (Softmax, SiLU) into adjacent loops can double vector-engine active time (30%→60%).
• Loop fusion/order optimization: Merging small loops and reordering to exploit data locality reduces HBM traffic by up to 20%.
• Algebraic peephole rewrites: Replacing reciprocal(sqrt()) with rsqrt(), or hand-fusing SiLU→$x\cdot\sigma(x)$, often leads to further latency reductions.

Identified limitations include reduced optimization success for small/irregular shapes (e.g., $K=64$ BatchMatMul) and restrictiveness of existing NKI DSL coarse loop primitives, which currently preclude certain register-level micro-schedules.

7. Open Challenges and Future Directions

Future enhancements for AWS Trainium kernel optimization and benchmarking include:

• Extension to multi-core scheduling with cross-core synchronization.
• Automated convolutional kernel pattern discovery and matching.
• Integration of analytical cost-models to restrict LLM-driven search spaces.

A plausible implication is that agentic memory-augmented LLM workflows can replicate and often exceed manual or expert-driven Trainium kernel optimizations, with significant cost advantages and scalability to future hardware targets (Zhang et al., 19 Nov 2025).