TPU and GPU Architecture: Deep-Dive into AI Accelerators

A Chatbot That Doesn't Get Faster

Picture a team building a Hindi-English customer support bot that answers over WhatsApp for millions of users — the kind of vernacular assistant now common across Indian fintech and e-commerce apps. It runs a 7-billion-parameter transformer on an NVIDIA A100 GPU, a card whose spec sheet advertises 312 TFLOPS of BF16 tensor-core throughput. The team profiles two phases of a single request: reading the user's message and conversation history (the "prefill" pass, where the model processes the whole prompt at once), and then generating the reply one word at a time (the "decode" pass, where each new token depends on the one before it). Prefill is fast and the GPU seems to be earning its price tag. Decode — the part that actually determines how quickly the reply appears, word by word — crawls, and profiling shows the tensor cores sitting at barely half a percent of their rated peak. Swapping in a GPU with twice the FLOPS barely moves the decode speed at all.

This is not a bug. It is the single most important fact about how GPUs and TPUs actually behave in production, and it has nothing to do with how many multiply-accumulate units are etched onto the die. It is about how many bytes have to travel from memory for every arithmetic operation performed — and that number is fixed by the shape of the computation, not by the accelerator. This chapter derives that relationship precisely, for both a GPU's memory hierarchy and a TPU's systolic array, and shows the exact arithmetic that predicts when an accelerator will and won't deliver its advertised number.

Bytes Moved, Not FLOPs Issued: The Roofline Model

Every accelerator has two hard ceilings on its speed: how many floating-point operations it can issue per second (its peak compute, in FLOPS), and how many bytes it can move between off-chip memory (HBM) and on-chip compute units per second (its memory bandwidth, in bytes/s). A computation cannot exceed either ceiling. Williams, Waterman, and Patterson formalized this as the roofline model in their 2009 Communications of the ACM paper, "Roofline: An Insightful Visual Performance Model for Multicore Architectures." The key quantity is arithmetic intensity (AI): the number of floating-point operations performed per byte moved from memory,

AI = total FLOPs / total bytes moved from HBM

Given AI, the roofline model bounds attainable performance as

attainable FLOPS = min(peak_compute, AI × bandwidth)

At low AI, a kernel is memory-bound: it is waiting on bytes, and doubling compute units does nothing, because the compute units simply idle longer between arrivals of data. At high AI, a kernel is compute-bound: bytes arrive faster than the arithmetic units can consume them, and bandwidth is irrelevant. The two regimes meet at the ridge point, AI* = peak_compute / bandwidth — the arithmetic intensity at which a chip's compute and memory ceilings are exactly saturated together. Above AI*, more bandwidth wouldn't help. Below it, more FLOPS wouldn't help. This single number tells you, before running anything, which upgrade is worth paying for.

Worked Example: The A100's Ridge Point, and Why Decode Lives Far Below It

The A100 80GB SXM4 card has a published BF16 tensor-core peak of 312 TFLOPS and HBM2e bandwidth of 2039 GB/s. Its ridge point:

peak = 312e12       # FLOPS, BF16 tensor core, no sparsity
bw   = 2039e9       # bytes/s, HBM2e
ridge = peak / bw
print(round(ridge, 1))
# 153.0

Trace it: 312,000 / 2,039 = 153.02, so any kernel needs at least about 153 FLOPs per byte moved before the A100's compute units, not its memory bus, become the bottleneck.

Now consider one decode step for a 7B-class transformer with hidden dimension d = 4096 (the actual hidden size of models like Llama-2-7B), generating a single token for a single user (batch = 1). Each linear layer performs a matrix-vector multiply: a 1×4096 activation against a 4096×4096 weight matrix stored in BF16.

d = 4096
flops_decode  = 2 * d * d          # multiply + add per weight element
bytes_decode  = 2 * d * d          # BF16 weights, read once, batch=1
ai_decode = flops_decode / bytes_decode
print(ai_decode)
# 1.0

def roofline_attainable(ai, peak_flops, bandwidth):
    return min(peak_flops, ai * bandwidth)

attain = roofline_attainable(ai_decode, peak, bw)
print(round(attain / 1e12, 2), "TFLOPS")
print(round(attain / peak * 100, 2), "% of peak")
# 2.04 TFLOPS
# 0.65 % of peak

Trace the arithmetic: FLOPs = 2 × 4096 × 4096 = 33,554,432. Bytes moved to read that same weight matrix once, at 2 bytes/element = 33,554,432. The two are numerically identical for any square weight matrix at batch = 1 — every weight is fetched from HBM, used in exactly one multiply-add, and discarded. AI = 1 FLOP/byte, independent of how big the matrix is, because doubling d quadruples both the FLOPs and the bytes together. Plugging AI = 1 into the roofline formula: attainable = min(312 TFLOPS, 1 × 2039 GB/s) = 2.04 TFLOPS — 0.65% of the card's peak. This is not a software inefficiency to be tuned away; it is the arithmetic ceiling imposed by the shape of a batch-1 matrix-vector multiply on this specific piece of hardware. The chatbot's decode step was never going to use the tensor cores' rated speed, no matter how well the kernel is written.

Worked Example: Prefill, and How Much Batching It Takes to Cross the Ridge

Prefill processes the whole prompt at once. With a prompt of S = 512 tokens, the same linear layer now performs a matrix-matrix multiply: a 512×4096 activation matrix against the 4096×4096 weight matrix.

S = 512
flops_prefill = 2 * S * d * d
bytes_prefill = (2 * d * d) + (2 * S * d) + (2 * S * d)   # weights + input acts + output acts
ai_prefill = flops_prefill / bytes_prefill
print(ai_prefill)
# 409.6

attain_p = roofline_attainable(ai_prefill, peak, bw)
print(round(attain_p / 1e12, 1), "TFLOPS")
# 312.0

Trace it: FLOPs = 2 × 512 × 4096 × 4096 = 17,179,869,184. Bytes = weight bytes (2 × 4096² = 33,554,432) plus input activations (2 × 512 × 4096 = 4,194,304) plus output activations (same, 4,194,304), totalling 41,943,040 bytes. Dividing: 17,179,869,184 / 41,943,040 = 409.6 exactly. That weight matrix now gets reused across all 512 tokens in the batch before it is evicted, so the fixed cost of fetching it is amortized over far more arithmetic — arithmetic intensity rises by more than 400× simply from batching, with the algorithm and hardware unchanged. Since 409.6 > 153 (the ridge point), prefill is compute-bound: the roofline model's ceiling for it is the full 312 TFLOPS. Real kernels never hit that ceiling exactly — instruction overhead, tiling boundaries, and non-matmul operations (softmax, layer norm) typically claim 30-70% of it — but the ceiling itself is set by compute, not bandwidth, which is the qualitative flip that separates prefill from decode.

The natural next question: how much batching does decode itself need before it also crosses the ridge? For B vectors processed together, FLOPs = 2Bd² and bytes = 2d² (weights, fetched once) + 4Bd (input and output activations), giving AI(B) = Bd / (d + 2B).

def ai_batched(B, d=4096):
    return (B * d) / (d + 2 * B)

for B in [1, 32, 165, 1000]:
    ai = ai_batched(B)
    attain = roofline_attainable(ai, peak, bw)
    print(B, round(ai, 2), round(attain / 1e12, 2))
# 1    1.0     2.04
# 32   31.51   64.24
# 165  152.70  311.35
# 1000 671.92  312.0

At B = 32 (a modest batch of concurrent users' decode steps grouped together — the "continuous batching" technique production LLM servers use), arithmetic intensity rises 31.5× and attainable throughput reaches 64.24 TFLOPS, still memory-bound but a 31× improvement over serving one user at a time. Solving Bd/(d+2B) = 153 for the ridge crossing: 4096B = 153(4096 + 2B) → 3790B = 626,688 → B ≈ 165.3, matching the table's B = 165 row landing just short of the peak. This is the precise, derivable reason production LLM serving systems batch hundreds of concurrent decode requests together rather than serving them one at a time: it is not a convenience, it is the only way to reach the hardware's compute ceiling at all.

The Misconception: "More TFLOPS Means Faster"

The chatbot team's instinct — that a card with a bigger advertised FLOPS number will generate replies faster — is the single most common misreading of accelerator spec sheets, and the roofline math above shows exactly why it fails. For a memory-bound kernel (AI below the ridge), attainable performance is AI × bandwidth, a product that does not contain the peak-FLOPS term at all. Two GPUs with identical HBM bandwidth but wildly different peak FLOPS — say 200 TFLOPS versus 400 TFLOPS — would decode this model's tokens at exactly the same speed, because decode's AI = 1 sits so far below either card's ridge point that neither card's extra compute capacity is ever reached. What would actually speed up decode is more HBM bandwidth, or restructuring the workload (batching, as derived above) to raise AI past the ridge. Reading only the FLOPS row of a spec sheet, without checking a workload's arithmetic intensity against the bandwidth row, routinely leads to purchasing decisions and engineering effort that produce zero measured improvement.

The TPU's Bet: A Weight-Stationary Systolic Array

Everything above concerns the GPU's roofline as a function of HBM bandwidth. Google's Tensor Processing Unit takes a structurally different approach to the same underlying problem — instead of streaming both weights and activations from a cache hierarchy for every operation, its matrix-multiply unit (MXU) is a systolic array: an N×N grid of multiply-accumulate cells, following the architecture H. T. Kung described in his 1982 IEEE Computer paper "Why Systolic Architectures?" A weight matrix is loaded once into the grid and held stationary at each cell; activation vectors then flow in from the left, one value shifting into the next cell every clock cycle, while partial sums flow downward, each cell adding its stationary weight's contribution before passing the sum on. Nothing is refetched from memory between cycles — the weight matrix, once loaded, is reused by every activation vector that streams through, which is a hardware-level version of exactly the amortization the batching example above achieved in software.

This dataflow has its own latency cost, distinct from the bandwidth-driven roofline analysis: the pipeline must "fill" before the first result emerges, and "drain" after the last input enters. Trace this by hand for the smallest non-trivial case, N = 2, multiplying a stationary 2×2 weight matrix W against a stream of row vectors, with each row's activation staggered by one cycle so it meets the correct partial sum:

The first complete output vector (y0 at cycle 2, y1 at cycle 3) takes 3 cycles to emerge from a 2×2 array — matching the general fill/drain latency of 2N−1 cycles for an N×N array. Streaming a batch of T vectors through, each entering one cycle behind the previous one, the last vector's result completes at cycle T + 2N − 2 (verified above: T=1 gives 1+2=3, matching the single-vector trace). Total useful MACs performed is T·N²; total cycles consumed is T + 2N − 2; so steady-state utilization is T / (T + 2N − 2) — for T = 4 vectors through a 4×4 array, that's 4/(4+6) = 40%; for T = 1000, it's 1000/1006 ≈ 99.4%. This is a second, independent source of underutilization beyond the roofline's bandwidth argument: even with data already resident and reused, a spatial systolic array still needs enough vectors streaming through to amortize its own pipeline latency.

The real MXU in Google's first TPU (Jouppi et al., "In-Datacenter Performance Analysis of a Tensor Processing Unit," ISCA 2017) is a 256×256 grid — 65,536 8-bit multiply-accumulate cells, clocked at 700 MHz:

N, clock = 256, 700e6
macs = N * N
peak_ops = 2 * macs * clock   # each MAC = 1 multiply + 1 add
print(round(peak_ops / 1e12, 2), "TOPS")
# 91.75 TOPS

2 × 65,536 × 700,000,000 = 91.75 × 10¹², matching the paper's reported peak of roughly 92 TOPS (int8). For this array, fill/drain overhead is 2(256) − 2 = 510 cycles; serving a batch of only T = 256 vectors reaches just 256/(256+510) ≈ 33.4% utilization, while T in the tens of thousands pushes utilization above 99%. This is precisely why TPUs have historically been positioned for large-batch training and bulk inference workloads rather than single-query, latency-sensitive serving — the hardware's own pipeline geometry, independent of any bandwidth argument, demands large batches to be worth its silicon.

The GPU's Bet: Software-Scheduled Tensor Cores

A GPU's tensor core takes the opposite design bet: instead of one large hardware-scheduled spatial array, each streaming multiprocessor (SM) carries small, fixed-function matrix units invoked repeatedly under compiler control. On Ampere-class GPUs, a single warp (32 threads) cooperatively issues a `mma.sync` instruction operating on a fixed tile shape — for BF16 inputs with FP32 accumulation, commonly a 16×8×16 tile (M×N×K) — with operands staged from shared memory into registers via `ldmatrix` before the instruction executes. Because there is no single large array holding an entire weight matrix stationary, the GPU's compiler and libraries (cuBLAS, CUTLASS) must explicitly tile larger matrix multiplies into many of these small `mma.sync` calls, staging each tile HBM → L2 (40 MB on the A100) → shared memory (up to 164 KB per SM) → registers, and double-buffering those transfers so the next tile is loading from HBM while the current one computes — hiding memory latency behind arithmetic rather than eliminating the fetch, the way a systolic array's stationary weights do.

This also sharpens the earlier misconception check: could the decode bottleneck simply be fixed by caching the whole weight matrix in on-chip SRAM instead of re-reading it from HBM? A 7B-parameter model in BF16 occupies 7×10&sup9; × 2 bytes = 14 GB. The A100's L2 cache is 40 MB — roughly 14,000 MB / 40 MB ≈ 350× too small to hold the model, so every decode step really must stream the weights from HBM again, and the memory-bound analysis stands. SRAM caching helps within a batch (reusing a tile already staged in shared memory across the K-dimension of one matmul) but cannot rescue batch = 1 decode, because the working set simply does not fit.

Active Recall

Attempt each question before checking the worked answer beneath it.

1. A weight matrix is now 8192×8192 instead of 4096×4096, BF16, still batch = 1. What is the new arithmetic intensity?
2. For an 8×8 weight-stationary systolic array, how many cycles does it take to fill and drain the pipeline for a single input vector, and what is the steady-state MACs-per-cycle throughput?
3. Continuous batching raises decode batch size from 1 to 32 for the d = 4096 model above. Recompute arithmetic intensity, state whether the layer is now compute- or memory-bound, and give the attainable TFLOPS.
4. A new accelerator has peak compute of 400 TFLOPS (BF16) and HBM bandwidth of 1600 GB/s. Compute its ridge point. A workload has AI = 100 FLOPs/byte on this chip — which regime is it in, and what is the attainable performance?
5. True or false: a TPU with a higher quoted peak TOPS than a GPU will always finish a given matrix multiply faster. Justify using the roofline model.
6. Why can't simply adding more on-chip SRAM to a GPU fix the batch-1 decode bottleneck for a 7B-parameter model?

Worked answers

1. AI is unchanged: FLOPs = 2d² and bytes = 2d² scale identically with d, so AI = 1 FLOP/byte regardless of matrix size. Doubling the hidden dimension doubles both the numerator and denominator, leaving the ratio fixed — arithmetic intensity for a batch-1 GEMV depends only on batch size, never on matrix dimension.

2. Fill/drain latency = 2N − 1 = 2(8) − 1 = 15 cycles for a single vector. Steady-state throughput = N² = 64 MACs/cycle once the pipeline is full and multiple vectors are streaming.

3. AI(32) = (32·4096)/(4096+64) = 131,072/4,160 = 31.51 FLOPs/byte. Since 31.51 < 153 (the A100's ridge point), the layer is still memory-bound. Attainable = 31.51 × 2039 GB/s ≈ 64.24 TFLOPS — a 31× improvement over batch 1's 2.04 TFLOPS, but still only 20.6% of the 312 TFLOPS peak, because 32 vectors still fall well short of the ≈165 needed to cross the ridge for this layer shape.

4. Ridge = 400e12 / 1600e9 = 250 FLOPs/byte. The workload's AI = 100 is below 250, so it is memory-bound. Attainable = 100 × 1600 GB/s = 160 TFLOPS, only 40% of the 400 TFLOPS peak — bandwidth, not the advertised peak, sets the ceiling here.

5. False. Attainable performance is min(peak, AI × bandwidth). If the workload's AI sits below both chips' ridge points, attainable performance is AI × bandwidth for each — a chip with a bigger peak-TOPS number but the same or lower bandwidth gains nothing, since the peak term is never the binding constraint in that regime. The TPU only wins if either the workload's AI exceeds its ridge point (compute-bound) or its bandwidth is genuinely higher.

6. A 7B-parameter BF16 model occupies 14 GB; even the A100's full 40 MB L2 cache is about 350× too small to hold it, so the weights must always be re-streamed from HBM on every decode step no matter how much (practically buildable) SRAM is added. What actually raises arithmetic intensity is reusing a fetched weight across many activation vectors at once — i.e., batching — not enlarging the cache for a working set that will never fit regardless.

Think About It

Think about this: How would you explain tpu and gpu architecture: deep-dive into ai accelerators 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.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where tpu and gpu architecture: deep-dive into ai accelerators is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting tpu and gpu architecture: deep-dive into ai accelerators to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind tpu and gpu architecture: deep-dive into ai accelerators, 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.