Tensor engine in AI systems
Engineering notebook — August 20, 2026
The big picture
An AI model — a transformer, a CNN, a diffusion model — is at its core a sequence of tensor operations: matrix multiplies, convolutions, element-wise functions, normalizations. The tensor engine (the systolic array MAC grid we've been building) exists to accelerate exactly one of those operations: general matrix multiply (GEMM).
The reason GEMM gets its own piece of silicon is arithmetic intensity. A single forward pass through GPT-3 requires ~350 billion floating-point multiply-accumulates. At 100 GFLOP/s on a CPU that would take 3.5 seconds per token. A tensor engine running at 300+ TOPS (tera-operations per second) brings that under a millisecond.
AI model (e.g. transformer)
│
├── Layer 0: embedding lookup ← memory bandwidth bound, no GEMM
├── Layer 1: attention ← QKV projection = 3× GEMM ← tensor engine
│ softmax ← element-wise, not GEMM
│ output projection ← GEMM ← tensor engine
├── Layer 2: feed-forward (FFN) ← 2× GEMM ← tensor engine
│ activation (GELU) ← element-wise
├── ...
└── Layer N: logits projection ← GEMM ← tensor engine
GEMM dominates. In a 70B-parameter transformer, roughly 85–90% of FLOPs are matrix multiplies. Everything else is overhead from the tensor engine's perspective.
Where the tensor engine sits in the hardware stack
┌─────────────────────────────────────────────────────────┐
│ Host system │
│ CPU ←→ PCIe / NVLink ←→ AI accelerator chip │
└─────────────────────────────────────────────────────────┘
Inside the AI accelerator:
┌─────────────────────────────────────────────────────────┐
│ On-chip system │
│ │
│ ┌──────────────┐ ┌─────────────────────────────┐ │
│ │ Scalar / SIMD│ │ Tensor engine │ │
│ │ cores │ │ (systolic array, MMA) │ │
│ │ (softmax, │ │ handles all GEMM /conv │ │
│ │ layernorm, │ └────────────┬────────────────┘ │
│ │ gelu, etc.) │ │ │
│ └──────────────┘ ┌────────────▼────────────────┐ │
│ │ On-chip SRAM (SMEM/L2) │ │
│ │ weight tiles + activation │ │
│ │ tiles staged here │ │
│ └────────────┬────────────────┘ │
│ │ DMA engine │
│ ┌───────────────────────────────────▼────────────────┐ │
│ │ HBM / LPDDR (off-chip DRAM) │ │
│ │ full model weights + KV-cache + activations │ │
│ └────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────┘
Key insight: the tensor engine almost never touches DRAM directly. Data must be staged into on-chip SRAM (by a DMA engine) before the systolic array can consume it. The bottleneck in production is almost always memory bandwidth, not compute throughput.
Mapping transformer operations to hardware
Attention layer
The multi-head self-attention block is the heart of a transformer. Here is how each sub-operation maps to hardware:
| Operation | Math | Hardware target |
|---|---|---|
| Q, K, V projection | X @ Wq, X @ Wk, X @ Wv |
Tensor engine (GEMM) |
| Attention scores | Q @ Kᵀ / √d |
Tensor engine (GEMM) |
| Softmax | exp(x) / Σexp(x) |
Scalar/SIMD cores |
| Context aggregation | scores @ V |
Tensor engine (GEMM) |
| Output projection | context @ Wo |
Tensor engine (GEMM) |
One attention layer = 5 GEMMs (for multi-head; more for grouped-query attention).
Feed-forward network (FFN)
FFN(x) = GELU(x @ W1 + b1) @ W2 + b2
| Operation | Hardware target |
|---|---|
x @ W1 |
Tensor engine |
GELU(...) |
Scalar/SIMD |
(result) @ W2 |
Tensor engine |
Two GEMMs per FFN layer, with scalar work sandwiched between them.
Normalizations and residuals
LayerNorm, RMSNorm, and residual additions are all element-wise or reduction operations. They run on scalar/SIMD cores and are fast — but they stall the tensor engine while they complete, which is why fused kernels (combining norm + GEMM in one dispatch) matter for efficiency.
The roofline model: compute vs. memory bound
Every operation on an accelerator falls somewhere on the roofline:
FLOP/s achieved
│
│ ████ roofline (peak compute)
│ ████████████
│████ ████████████████████ ← memory bandwidth ceiling
│
└────────────────────────────────────────── arithmetic intensity
low high
(memory bound) (compute bound)
Arithmetic intensity = FLOPs / bytes of memory traffic.
| Operation | Intensity | Bound |
|---|---|---|
| Large GEMM (training) | ~500 FLOP/byte | Compute |
| Small GEMM (inference, short seq) | ~10 FLOP/byte | Memory bandwidth |
| Softmax, layernorm | ~1–4 FLOP/byte | Memory bandwidth |
| KV-cache reads (long context) | Very low | Memory bandwidth |
This is why inference on large models is almost always memory bandwidth limited, not compute limited — you're reading huge weight matrices for every token, reusing them only a handful of times. The tensor engine is often sitting idle waiting for data.
Precision and number formats
The tensor engine's datapath is designed around a specific numerical format. Choosing precision is a first-order design decision:
| Format | Bits | Typical use | Notes |
|---|---|---|---|
| FP32 | 32 | Training (gradients) | High precision, expensive |
| BF16 | 16 | Training, inference | Same exponent range as FP32, less mantissa |
| FP16 | 16 | Inference | Smaller exponent range, needs loss scaling |
| INT8 | 8 | Inference (quantized) | 2× throughput vs FP16, needs calibration |
| FP8 | 8 | Training frontier | NV H100/B100, requires careful scaling |
| INT4 | 4 | Aggressive inference | 2–4× memory savings, accuracy loss |
Modern tensor engines support mixed precision: INT8 or FP16 inputs with FP32 accumulation, then quantize the output back down. Our systolic array design reflects this — DATA_W=8 inputs, ACC_W=32 accumulator.
Tiling: bridging the tensor engine and DRAM
A 70B model's weight matrices are enormous — a single FFN weight can be 8192×28672 elements. The systolic array (say, 128×128 PEs) can only hold a tiny tile of that at once. Tiling is the software/firmware strategy that sequences the work:
For each output tile (M_tile × N_tile) of C:
Zero the PE accumulators
For each K-chunk of width K_tile:
DMA: load A[m_tile, k_chunk] → SRAM buffer A
DMA: load B[k_chunk, n_tile] → SRAM buffer B
Systolic array: compute partial C += A_tile × B_tile
Write completed output tile back to DRAM
Double buffering overlaps these stages: while the tensor engine is computing tile N, the DMA is pre-fetching tile N+1 into the other SRAM bank. If the DMA is fast enough, the tensor engine never stalls — this is the target operating regime.
Cycle timeline:
DMA: [load tile 0] [load tile 1] [load tile 2] ...
Engine: [compute tile 0] [compute tile 1] ...
Overlap: ████████████████████████
Real silicon examples
These are commercial chips and how they implement the concept:
| Chip | Tensor engine | Key design point |
|---|---|---|
| NVIDIA H100 | 4th-gen Tensor Cores (MMA units) | FP8 support, 3958 TOPS (FP8 sparse) |
| Google TPU v4 | 128×128 systolic array per core | Closest to our design; 275 TOPS/chip |
| Apple M-series Neural Engine | 16-core array (Apple proprietary) | INT8/INT16, tightly integrated with CPU |
| Groq LPU | SIMD array + compiler-scheduled SRAM | Deterministic latency, no caches |
| Cerebras WSE-3 | 900K cores, full wafer | Eliminates off-chip memory entirely |
The TPU is architecturally the closest to our systolic array design — Google's original TPU paper (Jouppi et al., 2017) describes essentially what we've been building, scaled to 256×256.
The full inference pipeline
Putting it all together — what happens when you run a transformer inference pass:
1. Host CPU dispatches layer schedule to accelerator
2. DMA engine loads embedding weights → on-chip SRAM
3. Embedding lookup (scalar core)
4. For each transformer layer (L times):
a. DMA pre-fetches Wq, Wk, Wv weight tiles
b. Tensor engine: Q=X@Wq, K=X@Wk, V=X@Wv (3 GEMMs)
c. Tensor engine: scores = Q@Kᵀ (1 GEMM)
d. Scalar core: softmax(scores / √d)
e. Tensor engine: context = scores@V (1 GEMM)
f. Tensor engine: out = context@Wo (1 GEMM)
g. Scalar core: residual add + layernorm
h. DMA pre-fetches W1, W2
i. Tensor engine: hidden = X@W1 (1 GEMM)
j. Scalar core: GELU(hidden)
k. Tensor engine: output = hidden@W2 (1 GEMM)
l. Scalar core: residual add + layernorm
5. Tensor engine: logits = output@W_embed (1 GEMM)
6. Scalar core: argmax / sampling
7. Return token to host
Total GEMMs per layer: 8. For a 96-layer model (GPT-4 scale): ~768 GEMM calls per token generated.
Design bottlenecks in practice
| Bottleneck | Symptom | Mitigation |
|---|---|---|
| HBM bandwidth | Tensor engine utilization < 30% | Larger batch size, quantization, KV-cache compression |
| On-chip SRAM capacity | Frequent DMA stalls | Larger tile sizes, better prefetch scheduling |
| Scalar core throughput | Tensor engine waiting on softmax/norm | Kernel fusion (FlashAttention fuses QKᵀ + softmax + V) |
| PCIe bandwidth (host↔device) | First-token latency | Keep weights resident on device, batching |
| Tensor engine utilization | Low for small batches | Continuous batching, speculative decoding |
FlashAttention: a case study in hardware-aware algorithm design
FlashAttention (Dao et al., 2022) is the canonical example of redesigning an algorithm around hardware constraints rather than adding more hardware:
Problem: The naive attention implementation writes the full N×N score matrix to HBM and reads it back for softmax. For a 4096-token sequence, that's 4096² = 16M elements — dominated by memory traffic, not compute.
Solution: Tile the Q, K, V matrices to fit in SRAM. Compute attention scores, softmax, and context aggregation all in one pass without ever writing the intermediate score matrix to HBM.
Result: 5–9× less HBM reads/writes, 2–4× wall-clock speedup on A100. Same mathematical result.
This is the practical lesson: understanding the tensor engine and memory hierarchy lets you rethink the algorithm, not just the hardware. The most impactful optimization is often upstream of the silicon.
Summary: the hierarchy of concerns
┌─────────────────────────────────────────────────────┐
│ Model architecture (transformer, SSM, MoE, ...) │ ← which operations dominate?
├─────────────────────────────────────────────────────┤
│ Numerical precision (FP32 / BF16 / INT8 / FP8) │ ← throughput vs. accuracy
├─────────────────────────────────────────────────────┤
│ Algorithm (FlashAttention, PagedAttention, etc.) │ ← HBM traffic minimization
├─────────────────────────────────────────────────────┤
│ Tiling and scheduling (tile size, double buffer) │ ← SRAM utilization
├─────────────────────────────────────────────────────┤
│ Tensor engine (systolic array, MMA units) │ ← peak TOPS
├─────────────────────────────────────────────────────┤
│ Memory subsystem (HBM bandwidth, SRAM capacity) │ ← practical throughput
└─────────────────────────────────────────────────────┘
Every layer constrains the one above it. A perfect tensor engine sitting behind a narrow memory bus delivers mediocre real-world performance. A hardware-aware algorithm running on a modest tensor engine can outperform brute-force compute. Both matter.
Next steps
- [ ] Profile a real attention kernel to see where cycles are spent
- [ ] Study FlashAttention-2 tiling strategy in detail
- [ ] Understand how quantization (INT8/FP8) affects the PE datapath
- [ ] Explore fused kernel design: combining layernorm + GEMM in one dispatch
- [ ] Read: Jouppi et al. (2017) — "In-datacenter performance analysis of a TPU"
- [ ] Read: Dao et al. (2022) — "FlashAttention: fast and memory-efficient exact attention"