---
title: AWS Trainium Accelerator Kernels Overview
url: https://www.emergentmind.com/topics/aws-trainium-accelerator-kernels
type: topic
---

# AWS Trainium Accelerator Kernels Overview

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 large language models (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 [2511.15915].

## 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 [2511.15915]. 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):
```python
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 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 [2511.15915].

Source: https://www.emergentmind.com/topics/aws-trainium-accelerator-kernels