AI Chip Design and GPU/TPU Architectures

In 2024, when India's IndiaAI Mission decided how to give the country's researchers and startups access to large-scale AI compute, it did not commission a new fab to print chips domestically. It leased time on thousands of GPUs from data-centre operators, because designing and manufacturing a competitive AI accelerator is a different order of engineering problem from writing an application, and even a large government cannot shortcut it. Understanding why requires going inside the chip — not at the level of "GPUs are good at parallelism," which is true but empty, but at the level of what a streaming multiprocessor actually does with a matrix multiply, why Google built an entirely different chip called a TPU instead of just buying more GPUs, and how to predict, with arithmetic, whether adding more compute to a chip will make your model train faster at all. This chapter builds that from first principles, ending with a model — the roofline model — that you can use to reason about any accelerator, present or future.

Why a transformer layer cannot run fast on a CPU

A single forward pass through a transformer's feed-forward block is dominated by matrix multiplication: an activation matrix X of shape (batch·sequence, d_model) is multiplied against a weight matrix W of shape (d_model, 4·d_model). For a modest configuration with d_model = 4096 and a batch·sequence of 2048, that single multiply requires roughly 2 × 2048 × 4096 × 16384 ≈ 2.75 × 10¹¹ floating-point operations — and a full transformer runs dozens of these per layer, dozens of layers per forward pass, and a forward and backward pass per training step, repeated millions of times during pretraining.

A CPU core is built to do something almost opposite to this workload well: execute a long, unpredictable sequence of instructions with data-dependent branches, minimizing the *latency* of any single operation. To do that it spends most of its transistor budget on things that have nothing to do with raw arithmetic — branch predictors, out-of-order execution schedulers, large per-core caches, speculative execution — and it commits only a handful of arithmetic-logic units per core. A high-end server CPU might have 64 cores, each capable of a handful of floating-point operations per cycle. Multiply that out at a few GHz and you get on the order of a few teraFLOPS for the entire chip.

Matrix multiplication has the opposite character: every output element is computed by the same instruction (multiply-accumulate) applied to different data, with no branching and a completely predictable, regular pattern of memory access. There is nothing for a branch predictor or an out-of-order scheduler to do — the "decision" of what to compute next is known before the loop even starts. This is precisely the workload that rewards trading control-logic transistors for raw arithmetic-unit transistors: build thousands of simple multiply-accumulate units, and issue the *same* instruction to hundreds of them at once (single-instruction-multiple-thread, or SIMT), rather than one core independently deciding what to do next. That trade is what a GPU (and even more so, a TPU) is built around.

Inside a GPU: streaming multiprocessors, warps, and tensor cores

An NVIDIA A100, a chip that trained a large share of the world's transformer models between 2020 and 2022, is organized as 108 streaming multiprocessors (SMs). Each SM contains its own set of FP32 CUDA cores, a register file, and — since the Volta generation in 2017 — dedicated tensor cores: small hardware units whose entire job is to compute D = A·B + C on small fixed-size tiles of numbers in a single pipelined operation, at reduced precision (FP16 or BF16 multiply, FP32 accumulate) to save on both silicon area and memory traffic. Threads on an SM are grouped into warps of 32, and all 32 threads in a warp execute the same instruction on the same clock cycle, on different data — the hardware realization of the "same instruction, many data" idea above. Memory on the chip is organized as a hierarchy, and this hierarchy matters more for AI workloads than the compute units themselves. Each SM has a small, extremely fast pool of registers and shared memory/L1 cache (on the order of 100–200 KB, on-chip, sub-nanosecond latency). All SMs share a larger L2 cache (40 MB on the A100). Below that sits High Bandwidth Memory (HBM) — off-chip DRAM stacked next to the chip and connected through a very wide bus — which holds the full model weights and activations and is orders of magnitude slower to reach than on-chip memory, even though its bandwidth (over a thousand gigabytes per second) sounds enormous in absolute terms.

The reason this hierarchy is the crux of the whole chapter is a trend that has held for over a decade of accelerator design: peak compute throughput (FLOPS) has grown faster, generation over generation, than memory bandwidth (bytes/second). Every new chip gets relatively better at arithmetic and relatively worse at feeding that arithmetic with data. That single trend is why a huge fraction of real GPU kernels in a transformer — layer normalization, bias addition, activation functions, softmax — are not limited by how fast the chip can multiply numbers, but by how fast it can move numbers from HBM into the compute units. The tool for reasoning about exactly which regime a given operation falls into is the roofline model.

The roofline model: compute-bound or memory-bound?

Every chip has two hard ceilings: a peak compute rate, P_max (FLOPs/second), and a peak memory bandwidth, B_max (bytes/second). For any given computation, define its arithmetic intensity I as the number of floating-point operations it performs divided by the number of bytes it must move to and from memory to do so:

I = FLOPs / bytes moved

The attainable performance of that computation on the chip is capped by whichever ceiling it hits first:

Attainable performance = min(P_max, B_max × I)

The crossover point, called the ridge point, is I* = P_max / B_max. Below the ridge point the operation is memory-bound — it cannot go faster no matter how many more tensor cores you add, because the chip is starved waiting for data. Above the ridge point the operation is compute-bound, limited by the arithmetic units themselves, and more compute genuinely helps.

Work this out with real numbers. The A100 (40 GB PCIe variant) has a published FP16 tensor-core peak of P_max = 312 TFLOPS = 312 × 10¹² FLOPs/s, and an HBM2 bandwidth of B_max = 1555 GB/s = 1555 × 10⁹ bytes/s. The ridge point is:

I* = 312 × 10¹² / 1555 × 10⁹ ≈ 200.6 FLOPs/byte

Now compute the arithmetic intensity of a square matrix multiply C = A·B, where A, B, C are all n×n, stored in FP16 (2 bytes per element). A matmul performs one multiply and one add per output element per reduction step, giving FLOPs = 2n³. If we (idealize and) assume A, B and C are each read or written from HBM exactly once, the bytes moved are 3n² elements × 2 bytes = 6n². So:

I_matmul = 2n³ / 6n² = n / 3

For n = 4096 (roughly the size of the projection matrices inside a large transformer's feed-forward block), I_matmul = 4096/3 ≈ 1365.3 FLOPs/byte — nearly seven times the ridge point of 200.6. This matmul is deep in compute-bound territory: it attains the full 312 TFLOPS peak.

Now do the same for a bias-add — a purely elementwise operation, y = x + b, applied after that matmul, of exactly the kind that appears constantly in a transformer (layer norm, residual add, activation functions all have this shape). Each output element costs one FLOP and requires reading one element of x, one of b (2 bytes each), and writing one element of y (2 bytes) — 6 bytes moved per FLOP performed:

I_bias = 1 / 6 ≈ 0.167 FLOPs/byte

That is roughly 1,200 times below the ridge point. This operation is deeply memory-bound. Let's trace this through code exactly as it would run and print.

def attainable_performance(peak_flops, peak_bw, arithmetic_intensity):
    return min(peak_flops, peak_bw * arithmetic_intensity)

peak_flops = 312e12   # A100 FP16 tensor-core peak, FLOPs/s
peak_bw    = 1555e9   # A100 40GB HBM2 bandwidth, bytes/s

# Case 1: n=4096 square matmul, FP16
n = 4096
flops_matmul = 2 * n**3
bytes_matmul = 6 * n**2
ai_matmul = flops_matmul / bytes_matmul
perf_matmul = attainable_performance(peak_flops, peak_bw, ai_matmul)

# Case 2: elementwise bias-add, FP16
ai_bias = 1 / 6
perf_bias = attainable_performance(peak_flops, peak_bw, ai_bias)

print(f"Matmul   AI: {ai_matmul:.1f} FLOPs/byte, attainable: {perf_matmul/1e12:.1f} TFLOPS")
print(f"Bias-add AI: {ai_bias:.3f} FLOPs/byte, attainable: {perf_bias/1e9:.2f} GFLOPS")

Tracing this line by line: ai_matmul evaluates to 2·4096³ / (6·4096²) = 4096/3 = 1365.333…, and since 1555e9 × 1365.333 ≈ 2.12 × 10¹⁵ is far above peak_flops, attainable_performance returns peak_flops unchanged. The first print statement outputs exactly: Matmul AI: 1365.3 FLOPs/byte, attainable: 312.0 TFLOPS. For the bias-add, 1555e9 × (1/6) ≈ 2.592 × 10¹¹, far below peak_flops, so attainable_performance returns the bandwidth-limited value. The second print statement outputs exactly: Bias-add AI: 0.167 FLOPs/byte, attainable: 259.17 GFLOPS.

Compare those two numbers: 312,000 GFLOPS attained on the matmul versus 259.17 GFLOPS on the bias-add — a gap of roughly 1,200×, purely because the bias-add is starved of data, not starved of arithmetic units.

Correcting a common misconception

The natural assumption, seeing a spec sheet advertise "624 TFLOPS with sparsity" or "petaflops of AI performance," is that a chip with a bigger peak-FLOPS number will always make your model run faster. The roofline analysis above shows this is false for a large share of real transformer operations. If a kernel's arithmetic intensity sits below the ridge point — as the bias-add does — then increasing P_max (more tensor cores, higher clock speed) does nothing to its attainable performance, because the min(P_max, B_max × I) formula is already being capped by the bandwidth term, which never changed. The only ways to speed up a memory-bound kernel are to increase bandwidth (a more expensive memory system) or to raise its arithmetic intensity by doing more work per byte fetched.

That second option — raising arithmetic intensity through fusion — is exactly the engineering idea behind FlashAttention (Dao, Fu, Ermon, Rudra, and Ré, NeurIPS 2022). Standard attention computes the QK^T score matrix, writes the full n×n matrix to HBM, reads it back for the softmax, writes the softmax output back, and reads it again for the final weighted sum with V — several memory-bound round trips through HBM for an operation whose actual arithmetic is comparatively small. FlashAttention restructures the computation to keep intermediate tiles entirely in on-chip SRAM and fuses the score, softmax, and weighted-sum steps into one kernel, so far fewer bytes ever touch HBM per FLOP performed — moving the operation's arithmetic intensity closer to the ridge point rather than trying to buy a chip with more peak FLOPS. This is the standard playbook once you internalize the roofline model: for a memory-bound op, reach for fusion and tiling before reaching for a bigger chip.

TPU: a different dataflow, not just more cores

A GPU's SIMT design is still, fundamentally, a grid of independent processing units that each fetch operands from registers and shared memory for every multiply-accumulate. Google's Tensor Processing Unit, described in Jouppi et al., "In-Datacenter Performance Analysis of a Tensor Processing Unit" (ISCA 2017), takes a structurally different approach for the one operation that dominates deep learning: matrix multiplication. Instead of many independent cores each re-fetching operands, the TPU's matrix unit is a systolic array — a 2-D grid of simple multiply-accumulate cells wired directly to their neighbours, through which data flows like blood through a body (hence "systolic"). The first TPU's matrix unit was a 256×256 grid — 65,536 multiply-accumulate cells — built for 8-bit integer operations.

The key idea is weight-stationary dataflow. Before a matmul begins, the weight matrix is loaded once into the array — one weight value parked permanently in each cell's local register for the duration of the computation. Then activations stream in from one edge of the array (say, the left), one row at a time, and partial sums accumulate as they flow across the array in the perpendicular direction (say, downward). Every cell does exactly one thing every cycle: multiply its resident weight by the activation value passing through it, and add that product to the partial sum arriving from its neighbour, then pass both the activation and the updated partial sum onward to the next cell. Because the weight never leaves its cell once loaded, it is reused — for free, with zero additional memory traffic — by every activation value that streams past it. A single weight fetch from memory ends up contributing to hundreds of multiply-accumulate operations over the course of the streaming pass, which is precisely "raising arithmetic intensity," done in silicon wiring rather than in software fusion.

The diagram below shows a small 4×4 version of this array to make the mechanism concrete.

Because the array must fill before the first result emerges and drain before the last one does, a systolic array pays a fixed startup and shutdown cost measured in cycles proportional to the array's dimension — negligible when streaming a huge matmul through it, but a real overhead when the matrix being multiplied is small relative to the array. That asymmetry is central to why TPUs and GPUs are not simply interchangeable, and it is the subject of one of the practice questions below.

Active recall

Attempt each question before reading the answer beneath it.

Q1. A CPU core and a GPU core both run at roughly similar clock speeds. Explain, in terms of transistor allocation, why the GPU is roughly 100× faster on a matrix multiply but only marginally faster (or even slower) on a single-threaded branch-heavy program like a Python interpreter loop.

Q2. Compute the arithmetic intensity of an FP16 square matmul with n = 8192, and determine whether it is compute-bound or memory-bound on an A100 (peak 312 TFLOPS FP16, 1555 GB/s HBM2 bandwidth). What is the attainable performance?

Q3 (ripple). Suppose a new GPU is released with 4× the A100's peak FLOPS (1248 TFLOPS) but the same HBM bandwidth (1555 GB/s). Recompute the ridge point for this chip. Then determine, with numbers, what happens on this chip to (a) the bias-add example (I = 1/6) and (b) the n = 4096 matmul example (I ≈ 1365.3) from the worked example above.

Q4. Explain, referencing data movement, why a TPU's systolic array can be more energy-efficient per matmul FLOP than a GPU's tensor cores, even when both are doing the same arithmetic.

Q5. A classmate says: "TPUs have no general SM/warp overhead and are pure matmul specialists, so they should always outperform GPUs on AI workloads." Identify the flaw.

Q6. In an idealized systolic-array model, the number of cycles from the moment activations start streaming in to the moment the array has fully drained its last output is approximately 2N cycles for an N×N array (N cycles for the diagonal wavefront to reach the far corner, plus roughly N more cycles to drain all outputs). Estimate this latency for a 128×128 array, and explain in one sentence why this overhead matters less as the matrix being multiplied grows much larger than the array.

A1. A CPU core spends most of its transistor budget on control logic — branch prediction, out-of-order scheduling, large per-core caches — needed to minimize the latency of an unpredictable, data-dependent instruction stream, and it has only a few arithmetic units per core as a result. A GPU core is stripped of nearly all of that control logic and instead packs thousands of simple arithmetic units that all execute the *same* instruction on different data each cycle (SIMT). A matmul has no branches and a perfectly regular, predictable structure, so it is exactly the workload that rewards trading control logic for raw arithmetic units — hence the huge speedup. A branch-heavy Python-style loop has the opposite structure: unpredictable control flow that the GPU's simple cores handle poorly (they stall or diverge across a warp), while the CPU's branch predictor and out-of-order engine are built precisely for this case — hence little or no GPU advantage there.

A2. I = n/3 = 8192/3 ≈ 2730.7 FLOPs/byte. The ridge point is unchanged at ≈200.6 FLOPs/byte (same chip), so 2730.7 ≫ 200.6 means this is compute-bound, just like the n = 4096 case — it attains the full 312 TFLOPS peak. Doubling n only made the operation *more* compute-bound (I scales linearly with n), it did not change which regime it is in.

A3. New ridge point: I* = 1248×10¹² / 1555×10⁹ ≈ 802.6 FLOPs/byte. (a) Bias-add: I = 0.167 is still far below the new ridge point (802.6), so it remains memory-bound, and its attainable performance is unchanged at ≈259.17 GFLOPS — the extra compute the new chip offers is completely wasted on this kernel, because bandwidth, not compute, was always the limiter. This is the ripple a student is likely to miss: quadrupling peak FLOPS does not speed up every kernel. (b) Matmul (I ≈ 1365.3): this is still above the new ridge point of 802.6, so it remains compute-bound and now attains the full new peak of 1248 TFLOPS — a full 4× speedup, because this kernel was genuinely compute-limited and more compute helps it directly. The lesson: whether a hardware upgrade helps a given kernel depends entirely on which side of the ridge point that kernel sits, not on the upgrade itself.

A4. On a GPU, every multiply-accumulate performed by a tensor core requires operands to be fetched from registers or shared memory for that specific operation; a given weight value gets re-fetched (from some level of the memory hierarchy) each time it participates in a new tile computation. On the TPU's systolic array, a weight is fetched from memory exactly once, loaded into its cell, and then reused for every activation that streams past it during the entire pass — potentially hundreds of multiply-accumulates per single weight fetch. Since moving data (especially off-chip) costs far more energy per bit than performing an arithmetic operation on it, eliminating repeated weight fetches by pinning weights in place directly reduces the chip's energy cost per FLOP for this specific workload.

A5. The flaw is treating "specialized" as strictly better rather than as a different point on a flexibility/efficiency tradeoff. A systolic array is optimized for large, dense, regular matmuls where a weight matrix can be loaded once and reused across a long stream of activations. It handles irregular shapes, sparse matrices, small-batch inference, and workloads with dynamic control flow far less gracefully — small matrices leave much of the array's cells never used to full advantage, and the fixed fill/drain latency (as in Q6) becomes a larger fraction of the total runtime when the matrix being multiplied is not much bigger than the array itself. A GPU's independent SMs, by contrast, can each be given a different, smaller, or more irregular chunk of work concurrently. Neither architecture dominates the other across all workload shapes.

A6. Latency ≈ 2N = 2 × 128 = 256 cycles. This overhead matters less as the matrix being multiplied grows much larger than the array because the fixed 2N-cycle fill/drain cost is paid only once per pass, while the number of useful multiply-accumulate cycles the array performs afterward scales with the size of the matrix being streamed through — so the fixed overhead is amortized over an ever-larger amount of useful work and its fraction of total runtime shrinks toward zero.

Think About It

Think about this: How would you explain ai chip design and gpu/tpu architectures to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind ai chip design and gpu/tpu architectures, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.