Accelerator Architecture Field Guide

An 8-billion-parameter code assistant receives eight repository prompts, each 4,096 tokens long. Its prompt pass is fast on one accelerator, yet token generation crawls. Move the same model graph to another accelerator and the slow operation changes. Nothing about “more peak FLOPS” explains either result.

The CUDA foundation introduced kernels, thread blocks, warps, device memory, and synchronized timing. The model-parallelism lesson then showed that a logical shard count isn't a physical route. Keep both ideas: an accelerator runs work through a concrete execution hierarchy, and every tensor or collective travels through a concrete memory path.

We follow one model projection across five programming stacks. The intent isn't to crown a winner. Instead, preserve the workload's meaning while execution groups, local memory, compiler boundaries, and interconnects change.

Snapshot boundary, checked August 29, 2026: Concrete targets here are NVIDIA H100 (Hopper) and B200-class Blackwell under CUDA, AMD MI300/CDNA 3 and MI350/CDNA 4 under the ROCm 7.14 documentation, Google TPU v6e with Pallas, AWS Trainium2 with Neuron 2.32 documentation, and Apple silicon through PyTorch MPS. Hardware availability, software support, and cloud shapes change. Recheck official compatibility pages before buying capacity or freezing a production image.^{[1]Reference 1NVIDIA Hopper Tuning Guidehttps://docs.nvidia.com/cuda/archive/13.0.0/hopper-tuning-guide/index.html[2]Reference 2Blackwell Tuning Guidehttps://docs.nvidia.com/cuda/archive/13.0.2/pdf/Blackwell_Tuning_Guide.pdf[3]Reference 3ROCm Core SDK 7.14.0 Release Noteshttps://rocm.docs.amd.com/en/docs-7.14.0/about/release-notes.html[4]Reference 4TPU v6ehttps://cloud.google.com/tpu/docs/v6e[5]Reference 5The Trainium Memory Hierarchyhttps://awsdocs-neuron.readthedocs-hosted.com/en/v2.32.0/nki/get-started/about/memory-hierarchy-overview.html[6]Reference 6Accelerated PyTorch training on Mac.https://developer.apple.com/metal/pytorch/}

One projection, two different machines inside it

Take the first feed-forward projection from the running assistant. Each input row has K = 4,096 features, and the projection produces N = 11,008 features. Its simplified matrix multiplication is:

$$Y_{M \times N} = X_{M \times K} W_{K \times N}$$

M is the number of token rows processed together. BF16 stores each matrix element in 2 bytes. The weight matrix contains about 86 MiB, independent of phase.

During prefill, all prompt tokens are processed in parallel. Eight 4,096-token prompts give M = 32,768. During decode, each active sequence contributes one new token, so the same eight requests give M = 8. The weight shape hasn't changed. Reuse has.

For a matrix multiplication, the arithmetic work is approximately:

$$\text{FLOPs} = 2MKN$$

An idealized traffic floor that reads each input and weight once and writes each output once is:

$$\text{bytes} = 2(MK + KN + MN)$$

Dividing work by traffic gives arithmetic intensity, measured in floating-point operations per byte. The script computes both phases as a shape calculation, not a hardware benchmark.

Real kernels reread some data, cache some tiles, fuse operations, and may use quantized weights. Still, the phase split survives: prefill presents large matrix work, while low-concurrency decode repeatedly streams a large weight set for a few rows. Peak matrix throughput predicts prefill better than it predicts decode.

Why can the same projection be compute-heavy during prefill and memory-heavy during decode?

Answer

Prefill applies each loaded weight tile to 32,768 token rows in this workload, while decode applies it to only 8. The large prefill matrix gets far more arithmetic from each byte moved. Decode has much lower arithmetic intensity, so weight traffic and fixed dispatch costs can dominate.

The invariant behind every accelerator

Fast accelerator code keeps reused data close to compute, performs enough work before eviction, and overlaps the next transfer with current arithmetic. Vendor names differ, but four ownership questions don't:

1. Which group executes one instruction together?
2. Which memory can that group share at low latency?
3. Who schedules movement between large memory and local memory?
4. Which physical link carries bytes when work spans devices?

Time to first token (TTFT) covers prompt processing and scheduling. Time per output token (TPOT) exposes the serial decode loop. The running projection turns the ownership questions into a four-stage receipt:

A platform port is incomplete until both clocks and an output-correctness check are attached to the same prompt distribution.

The local-memory vocabulary is different, but the job is recognizable:

The table compares ownership, not speed. After a symptom appears, inspect the platform-specific owner in its row.

NVIDIA Hopper and Blackwell: preserve mapping, retune pressure

CUDA presents grids of thread blocks, blocks of threads, 32-thread warps, and streaming multiprocessors (SMs). Threads in one block share on-chip shared memory and can synchronize. Global device memory holds large tensors. Those boundaries from the CUDA prerequisite remain valid on Hopper and Blackwell.^{[7]Reference 7CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Hopper added the Tensor Memory Accelerator (TMA), which can move multidimensional tensor tiles between global and shared memory while thread blocks continue independent work. Hopper also added thread-block clusters and distributed shared memory across blocks in a cluster.^{[1]Reference 1NVIDIA Hopper Tuning Guidehttps://docs.nvidia.com/cuda/archive/13.0.0/hopper-tuning-guide/index.html} Those features are useful only when a kernel has a tile worth reusing. Copying the decode projection's 86 MiB weight matrix through shared memory without enough rows to reuse each tile adds staging without changing the bandwidth problem.

Blackwell retains and extends the CUDA programming model. NVIDIA's Blackwell tuning guide still starts with coalesced global access, reduced redundant traffic, suitable launch configuration, and limited warp divergence. It also warns that occupancy limits differ across Blackwell compute capabilities, so “Blackwell” isn't one register-and-shared-memory budget.^{[2]Reference 2Blackwell Tuning Guidehttps://docs.nvidia.com/cuda/archive/13.0.2/pdf/Blackwell_Tuning_Guide.pdf}

For the running workload, use the same algorithmic split on both generations:

• Prefill: choose a matrix path whose tile shape keeps Tensor Core work dense, then overlap global-to-shared movement with current computation.
• Decode: batch enough active rows to reuse weights, use a kernel specialized for small M, or reduce bytes through a verified quantization path.
• Port from Hopper to Blackwell: ship native code for the target or forward-compatible PTX, then retune. A binary that launches is only a compatibility result.^{[2]Reference 2Blackwell Tuning Guidehttps://docs.nvidia.com/cuda/archive/13.0.2/pdf/Blackwell_Tuning_Guide.pdf}

A CUDA failure that looks like architecture progress

Suppose a Hopper kernel uses more shared memory to stage a larger weight tile. One block now reuses more bytes, but fewer blocks remain resident on each SM. TPOT gets worse.

The symptom isn't proof that shared-memory tiling failed. Larger tiles improved reuse and reduced concurrency at the same time. Compare achieved occupancy, memory traffic, and eligible warps, then test smaller tiles. If decode still moves almost the same weight bytes per token, more staging can't create missing reuse.

A Hopper decode kernel gets slower after its shared-memory tile doubles. Which two effects must be separated?

Answer

The larger tile may reduce global-memory traffic per block, but its shared-memory and register footprint may lower occupancy. Measure both data movement and resident or eligible warps before deciding whether the tile helped.

AMD CDNA and ROCm: familiar syntax, different lane contract

HIP deliberately resembles CUDA: kernels launch grids of blocks, threads use threadIdx and blockIdx, and a work-group shares Local Data Share (LDS). That similarity helps source portability. It doesn't make wave-level assumptions portable.

On CDNA architectures, AMD's HIP programming-model documentation specifies a 64-thread wavefront. NVIDIA code commonly assumes a 32-thread warp. A reduction that hardcodes masks, lane counts, or “four warps per 128-thread block” can return wrong values or waste half a wave after a mechanical HIP conversion.^{[8]Reference 8Introduction to the HIP Programming Modelhttps://rocm.docs.amd.com/projects/HIP/en/develop/understand/programming_model.html}

Repair the execution model before changing syntax:

• derive lane behavior from warpSize or use library primitives with documented semantics;
• recalculate work-group size as a number of 64-lane waves on CDNA;
• remeasure vector-register and LDS pressure, because both limit resident waves;
• compile for the exact LLVM gfx target in the pinned environment.

As of the snapshot date, ROCm 7.14 documentation lists MI300-series accelerators as CDNA 3 with gfx942 and MI350-series accelerators as CDNA 4 with gfx950.^{[3]Reference 3ROCm Core SDK 7.14.0 Release Noteshttps://rocm.docs.amd.com/en/docs-7.14.0/about/release-notes.html} Keep those identifiers in build and benchmark records. “ROCm passed” without hardware target, ROCm version, and library versions isn't reproducible.

RCCL provides collective communication for ROCm stacks. Its current MI300X guidance describes eight-GPU systems where every accelerator pair has dedicated Infinity Fabric links. Using only part of that topology changes available collective routes, so a tensor-parallel result needs the selected ranks, not only TP=4.^{[9]Reference 9RCCL Documentationhttps://rocm.docs.amd.com/projects/rccl/en/docs-7.2.4/}

For the projection, CDNA's decision remains phase-specific. Dense prefill can keep matrix units busy through large tiles. Decode needs enough concurrent rows, a small-M kernel, or fewer weight bytes. Replacing CUDA API names with HIP API names can't change M = 8.

A 128-thread CUDA reduction is ported to CDNA and keeps a hardcoded 32-lane shuffle mask. Why is this a correctness risk rather than only a tuning issue?

Answer

CDNA uses 64-thread wavefronts. A hardcoded 32-lane algorithm may combine only half of each wave or apply invalid lane assumptions, so output can be wrong. Replace the assumption with documented cross-lane primitives or a wave-size-aware algorithm, then test against a reference.

TPU and Pallas: map tiles, not CUDA threads

TPU v6e exposes a TensorCore with matrix-multiply units (MXUs), a vector unit, and a scalar unit. The MXUs handle dense matrix work, while vector and scalar operations have separate execution resources.^{[4]Reference 4TPU v6ehttps://cloud.google.com/tpu/docs/v6e} A port that treats TPU as a GPU with different thread names misses that split.

Pallas is JAX's custom-kernel layer for GPU and TPU. On TPU, kernel inputs usually reside in high-bandwidth memory (HBM), while kernel-body references point into faster vector memory (VMEM) or scalar memory (SMEM). BlockSpec describes which input and output tile each program sees; the compiler can overlap HBM transfers with computation.^{[10]Reference 10Writing TPU Kernels with Pallashttps://docs.jax.dev/en/latest/pallas/tpu/details.html}

Grid order carries extra meaning. Pallas TPU programs normally advance sequentially in lexicographic grid order. Consecutive programs that use the same input slice can reuse data already in VMEM and skip another HBM transfer.^{[10]Reference 10Writing TPU Kernels with Pallashttps://docs.jax.dev/en/latest/pallas/tpu/details.html} Reordering grid axes can therefore change traffic without changing the mathematical output.

Map the running projection this way:

1. Put the 86 MiB weights in HBM.
2. Use BlockSpec to select weight and activation tiles that fit VMEM.
3. Arrange the prefill grid so consecutive output tiles reuse one weight or input slice where possible.
4. Accumulate matrix work on MXUs and keep softmax, normalization, or element-wise work on suitable vector paths.
5. For M = 8 decode, test padded or batched shapes against real TPOT. A large MXU doesn't guarantee high use for a skinny matrix.

Pallas block shapes have backend constraints, including divisibility requirements on the trailing dimensions for TPU paths.^{[10]Reference 10Writing TPU Kernels with Pallashttps://docs.jax.dev/en/latest/pallas/tpu/details.html} A shape rejection at compile time isn't an availability incident. The tile contract failed. Pad with correctness masking, choose a legal block, or keep the operation in compiled JAX when a custom kernel doesn't earn its maintenance cost.

A TPU failure caused by legal but poor grid order

Two Pallas kernels produce identical outputs. Kernel A varies the output-column tile fastest, so consecutive programs reuse the same input tile. Kernel B varies the input-row tile fastest and reloads that input tile more often. Kernel B reports more HBM traffic and worse prefill time.

The math can't distinguish them. The grid-to-block mapping can. Inspect BlockSpec, grid order, and HBM-to-VMEM transfers before changing model precision.

Trainium and NKI: assign work to engines and memories

Trainium2's NeuronCore-v3 contains tensor, vector, scalar, and general-purpose SIMD (GPSIMD) engines plus software-managed on-chip SRAM.^{[11]Reference 11NeuronCore-v3 Architecturehttps://awsdocs-neuron.readthedocs-hosted.com/en/v2.29.1/about-neuron/arch/neuron-hardware/neuron-core-v3.html} Neural Kernel Interface (NKI) code makes that specialization visible.

NKI's memory path has three named levels for our projection:

• HBM holds kernel inputs and outputs.
• SBUF is the main software-managed on-chip buffer shared by compute engines.
• PSUM holds partial matrix-multiply accumulations near the tensor engine.

An NKI kernel loads HBM tiles into SBUF, performs work from internal memory, accumulates matrix results in PSUM when appropriate, and stores completed outputs back through SBUF to HBM.^{[5]Reference 5The Trainium Memory Hierarchyhttps://awsdocs-neuron.readthedocs-hosted.com/en/v2.32.0/nki/get-started/about/memory-hierarchy-overview.html} If live tiles exceed SBUF or PSUM capacity, the compiler inserts spills and refills. That can preserve correctness while quietly destroying the intended traffic pattern.

The running projection suggests an engine schedule: direct-memory-access engines prefetch the next X and W tiles, the tensor engine performs matrix multiplication, vector or scalar engines handle fused follow-up operations they support, and the current output tile accumulates in PSUM. NKI tile dimensions also distinguish partition and free dimensions, so a valid NumPy shape isn't yet a valid physical layout.

Trainium2 systems connect chips through NeuronLink-v3 in a 4-by-4 two-dimensional torus within a 16-chip instance.^{[12]Reference 12Amazon EC2 Trn2 Architecturehttps://awsdocs-neuron.readthedocs-hosted.com/en/v2.32.0/about-neuron/arch/neuron-hardware/trn2-arch.html} Keep tensor-parallel ranks aligned with the actual topology. A rank count alone hides routing, exactly as it did for CUDA and ROCm.

An NKI failure that keeps answers correct

The prefill kernel grows its fused region to remove one HBM round trip. Live intermediates no longer fit SBUF, so the compiler adds spills. Latency increases and the trace shows extra DMA traffic.

Split the fusion or shrink tiles until the live set stays on chip. “More fusion” isn't a monotonic optimization when software-managed local memory is the limiting resource.

Why can a larger fused NKI kernel run slower even though it launches fewer kernels?

Answer

Fusion lengthens the lifetime of intermediate tiles. If that live set overflows SBUF or PSUM, spills add HBM traffic. Fewer launches can lose to more bytes moved.

Apple MPS: unified memory changes transfer, not capacity

PyTorch's Metal Performance Shaders (MPS) backend maps tensor operations to MPS Graph and tuned MPS kernels on Apple platforms.^{[6]Reference 6Accelerated PyTorch training on Mac.https://developer.apple.com/metal/pytorch/[13]Reference 13MPS backend.https://docs.pytorch.org/docs/stable/notes/mps} MPS Graph represents operations and tensors as a symbolic compute graph that can be compiled into an executable.^{[14]Reference 14Metal Performance Shaders Graphhttps://developer.apple.com/documentation/metalperformanceshadersgraph}

Apple silicon uses a unified physical memory pool. The CPU and GPU don't have the discrete host-RAM-to-VRAM boundary used by a typical data-center GPU. PyTorch still treats cpu and mps as different device targets, and unified memory is still finite.^{[15]Reference 15MLX: An array framework for Apple siliconhttps://github.com/ml-explore/mlx[13]Reference 13MPS backend.https://docs.pytorch.org/docs/stable/notes/mps}

For the 8B BF16 model, weights alone need about 16 GB in decimal units. The machine also needs KV cache, temporary activations, allocator headroom, macOS, and other applications. A Mac advertised with enough total memory can still enter pressure or fail allocation once the full serving ledger is counted.

MPS is strongest in this field guide as a local path: validate graph behavior, test product logic, and measure an on-device latency envelope. Keep three questions separate:

1. Is the operation implemented on MPS?
2. Does the whole model fit within safe system-memory headroom?
3. Does the graph execute fast enough after warmup and synchronization?

When PYTORCH_ENABLE_MPS_FALLBACK=1 is enabled, an unsupported MPS operation can fall back to the CPU and create a latency cliff.^{[16]Reference 16MPS Environment Variables.https://docs.pytorch.org/docs/stable/mps_environment_variables.html} Unified memory makes the detour less visually obvious than a discrete device copy, but it doesn't make CPU and GPU execution equally fast. Profile for CPU operations and synchronization gaps before blaming GPU arithmetic.

Choose programming depth after the bottleneck

Most ports should stop at the highest layer that meets correctness and service-level objectives (SLOs). Each step downward gains control and creates a new maintenance surface:

Portability is a test matrix, not a source-language property. A Triton kernel may support NVIDIA and AMD while using different legal tile sizes. A Pallas kernel may share BlockSpec vocabulary across GPU and TPU while requiring backend-specific memory and shape rules. NKI exposes Trainium's memory hierarchy directly. MPS Graph may compile a complete graph without offering the same custom-kernel surface as CUDA.

For our assistant, start with graph-level BF16 inference on each candidate. Preserve tokenization, weights, prompts, decoding settings, and output checks. Move down only after a trace identifies a stable projection, attention, normalization, or collective hotspot.

Read failures as ownership mistakes

The same symptom can point to a different owner on each platform. Use evidence in this order: output, phase, memory path, execution mapping, then topology.

Don't change precision, batching, kernel code, and topology in one experiment. That destroys attribution. Keep one baseline and change one ownership boundary at a time.

Build an architecture decision receipt

Suppose the assistant must meet p95 TTFT below 450 ms and p95 TPOT below 45 ms for eight 4,096-token prompts. Those thresholds are exercise inputs, not published accelerator results.

For every candidate, record one row with:

• exact accelerator and count;
• host or instance shape and physical topology;
• driver, runtime, framework, compiler, and library versions;
• model revision, dtype, quantization, and kernel path;
• prompt, output, concurrency, arrival, and warmup distributions;
• weight, KV-cache, temporary, and safety-headroom memory;
• p50, p95, and p99 TTFT and TPOT;
• output tolerance plus task-level quality check;
• profiler evidence for the largest compute, memory, and collective regions.

Then make a constrained decision:

1. Reject any point that fails output or task quality.
2. Reject any point that exceeds safe memory headroom.
3. Reject any point that misses either latency SLO.
4. Compare goodput and cost only among remaining points.

Architecture changes which knobs exist. It doesn't change that decision order. An accelerator with a higher unconstrained token rate still loses if its p95 TPOT misses the user's limit.

Two accelerators pass numerical checks. Candidate A has higher total tokens per second but misses the p95 TPOT limit. Candidate B meets TTFT and TPOT with lower total throughput. Which candidate is eligible?

Answer

Candidate B. Correctness, fit, and latency SLOs are gates. Throughput ranks only the candidates that pass all gates.

The field guide leaves one unresolved problem: each diagnosis above depends on trustworthy timing, traces, counters, and output comparisons. The next lesson turns those words into a repeatable benchmark and correctness workflow.