Lecture 14: SCE Accelerator Architecture
Part IV: Applications | SCE Futures
!!! abstract "Learning Objectives" - Trace data movement and partial sums through a systolic array. - Tile transformer matrix multiplications for a finite processing-element array. - Quantify reuse across registers, on-chip buffers, and warm memory. - Explain how AQFP timing and state costs shape a weight-stationary accelerator. - State device and system energy results with their measurement or model basis.
1. From Transformers to Hardware¶
In the previous lecture, we established that transformers are fundamentally composed of:
- Matrix-matrix multiplications (attention projections, FFN layers)
- Matrix-vector multiplications (inference with batch=1)
- Element-wise operations (activations, layer norm)
- Softmax (attention weights)
The dominant operations are GEMMs (General Matrix-Matrix Multiplies), which account for >90% of compute.
The Hardware Design Question¶
How do we build hardware that efficiently executes millions of these matrix operations?
Options:
| Architecture | Strengths | Weaknesses |
|---|---|---|
| CPU | Flexible, large cache | Low parallelism, poor FLOP/W |
| GPU | Massive parallelism | Complex memory hierarchy, power-hungry |
| Systolic Array | High efficiency, simple design | Fixed dataflow, limited flexibility |
| FPGA | Reconfigurable | Lower performance, complex tooling |
Systolic arrays have emerged as the architecture of choice for inference accelerators, used in:
- Google TPU (Tensor Processing Unit)
- Apple Neural Engine
- Many inference ASICs
In [1]:
Copied!
# Setup: Import libraries for visualizations
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyArrowPatch
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import numpy as np
COLORS = {
'primary': '#2196F3', # Blue
'secondary': '#FF9800', # Orange
'success': '#4CAF50', # Green
'danger': '#f44336', # Red
'dark': '#1a1a2e', # Dark navy
'light': '#f5f5f5', # Light gray
'purple': '#9C27B0', # Purple
'cyan': '#00BCD4', # Cyan
'weight': '#E91E63', # Pink for weights
'activation': '#3F51B5', # Indigo for activations
'accumulator': '#009688', # Teal for accumulators
'pe': '#FFC107', # Amber for processing elements
}
plt.rcParams['figure.facecolor'] = 'white'
plt.rcParams['axes.facecolor'] = 'white'
plt.rcParams['font.size'] = 11
print("Setup complete.")
# Setup: Import libraries for visualizations
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyArrowPatch
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import numpy as np
COLORS = {
'primary': '#2196F3', # Blue
'secondary': '#FF9800', # Orange
'success': '#4CAF50', # Green
'danger': '#f44336', # Red
'dark': '#1a1a2e', # Dark navy
'light': '#f5f5f5', # Light gray
'purple': '#9C27B0', # Purple
'cyan': '#00BCD4', # Cyan
'weight': '#E91E63', # Pink for weights
'activation': '#3F51B5', # Indigo for activations
'accumulator': '#009688', # Teal for accumulators
'pe': '#FFC107', # Amber for processing elements
}
plt.rcParams['figure.facecolor'] = 'white'
plt.rcParams['axes.facecolor'] = 'white'
plt.rcParams['font.size'] = 11
print("Setup complete.")
Setup complete.
2. What is a Systolic Array?¶
A systolic array is a network of processing elements (PEs) arranged in a regular grid, where data flows rhythmically through the array like blood through the circulatory system (hence "systolic").
Key Characteristics¶
- Regular structure: 2D grid of identical PEs
- Local communication: Each PE only talks to immediate neighbors
- Pipelined execution: Data flows continuously through the array
- High data reuse: Each value is used multiple times as it flows
Processing Element (PE)¶
Each PE performs a simple operation:
accumulator += input_A × input_B
pass input_A to right neighbor
pass input_B to bottom neighbor
This is the multiply-accumulate (MAC) operation, the fundamental building block of matrix multiplication.
Why Systolic Arrays are Efficient¶
- Minimal memory access: Data is fetched once, used many times
- No global interconnect: Only local wires between neighbors
- Simple control: All PEs execute the same instruction
- High utilization: Once filled, every PE is computing every cycle
In [2]:
Copied!
# Visualize: Basic systolic array structure
fig, ax = plt.subplots(figsize=(12, 10))
n = 4 # 4x4 array
pe_size = 1.0
spacing = 1.5
# Draw PEs
for i in range(n):
for j in range(n):
x = j * spacing + 2
y = (n - 1 - i) * spacing + 1
# PE box
pe = patches.FancyBboxPatch((x - pe_size/2, y - pe_size/2), pe_size, pe_size,
boxstyle="round,pad=0.02",
facecolor=COLORS['pe'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(pe)
ax.text(x, y, f'PE\n({i},{j})', ha='center', va='center', fontsize=8, fontweight='bold')
# Horizontal arrows (right) - for A matrix
if j < n - 1:
ax.annotate('', xy=(x + spacing - pe_size/2 - 0.1, y),
xytext=(x + pe_size/2 + 0.1, y),
arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=2))
# Vertical arrows (down) - for B matrix
if i < n - 1:
ax.annotate('', xy=(x, y - spacing + pe_size/2 + 0.1),
xytext=(x, y - pe_size/2 - 0.1),
arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=2))
# Input labels - A matrix (left side)
for i in range(n):
y = (n - 1 - i) * spacing + 1
ax.annotate('', xy=(2 - pe_size/2 - 0.1, y),
xytext=(0.5, y),
arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=3))
ax.text(0.3, y, f'A[{i},:]', ha='right', va='center', fontsize=10, color=COLORS['activation'])
# Input labels - B matrix (top)
for j in range(n):
x = j * spacing + 2
ax.annotate('', xy=(x, (n-1) * spacing + 1 + pe_size/2 + 0.1),
xytext=(x, (n-1) * spacing + 2.5),
arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=3))
ax.text(x, (n-1) * spacing + 2.7, f'B[:,{j}]', ha='center', va='bottom', fontsize=10, color=COLORS['weight'])
# Legend
ax.plot([], [], 's', markersize=15, color=COLORS['pe'], label='Processing Element (PE)')
ax.plot([], [], '->', markersize=10, color=COLORS['activation'], label='Activation flow (A matrix)', linewidth=2)
ax.plot([], [], '->', markersize=10, color=COLORS['weight'], label='Weight flow (B matrix)', linewidth=2)
ax.legend(loc='lower right', fontsize=10)
# Title and labels
ax.set_xlim(-0.5, 9)
ax.set_ylim(-0.5, 8)
ax.set_aspect('equal')
ax.axis('off')
ax.set_title(f'{n}×{n} Systolic Array Architecture', fontsize=14, fontweight='bold')
# Annotations
ax.text(8, 6, 'Each PE performs:\nMAC operation\n\nacc += a × b\npass a right\npass b down',
fontsize=10, va='top', ha='left',
bbox=dict(boxstyle='round', facecolor='white', edgecolor='gray'))
plt.tight_layout()
plt.show()
print(f"A {n}×{n} systolic array has {n*n} = {n**2} PEs")
print(f"Each cycle, it performs {n**2} MAC operations = {2*n**2} FLOPs")
print(f"TPU v4 has 128×128 = 16,384 PEs!")
# Visualize: Basic systolic array structure
fig, ax = plt.subplots(figsize=(12, 10))
n = 4 # 4x4 array
pe_size = 1.0
spacing = 1.5
# Draw PEs
for i in range(n):
for j in range(n):
x = j * spacing + 2
y = (n - 1 - i) * spacing + 1
# PE box
pe = patches.FancyBboxPatch((x - pe_size/2, y - pe_size/2), pe_size, pe_size,
boxstyle="round,pad=0.02",
facecolor=COLORS['pe'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(pe)
ax.text(x, y, f'PE\n({i},{j})', ha='center', va='center', fontsize=8, fontweight='bold')
# Horizontal arrows (right) - for A matrix
if j < n - 1:
ax.annotate('', xy=(x + spacing - pe_size/2 - 0.1, y),
xytext=(x + pe_size/2 + 0.1, y),
arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=2))
# Vertical arrows (down) - for B matrix
if i < n - 1:
ax.annotate('', xy=(x, y - spacing + pe_size/2 + 0.1),
xytext=(x, y - pe_size/2 - 0.1),
arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=2))
# Input labels - A matrix (left side)
for i in range(n):
y = (n - 1 - i) * spacing + 1
ax.annotate('', xy=(2 - pe_size/2 - 0.1, y),
xytext=(0.5, y),
arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=3))
ax.text(0.3, y, f'A[{i},:]', ha='right', va='center', fontsize=10, color=COLORS['activation'])
# Input labels - B matrix (top)
for j in range(n):
x = j * spacing + 2
ax.annotate('', xy=(x, (n-1) * spacing + 1 + pe_size/2 + 0.1),
xytext=(x, (n-1) * spacing + 2.5),
arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=3))
ax.text(x, (n-1) * spacing + 2.7, f'B[:,{j}]', ha='center', va='bottom', fontsize=10, color=COLORS['weight'])
# Legend
ax.plot([], [], 's', markersize=15, color=COLORS['pe'], label='Processing Element (PE)')
ax.plot([], [], '->', markersize=10, color=COLORS['activation'], label='Activation flow (A matrix)', linewidth=2)
ax.plot([], [], '->', markersize=10, color=COLORS['weight'], label='Weight flow (B matrix)', linewidth=2)
ax.legend(loc='lower right', fontsize=10)
# Title and labels
ax.set_xlim(-0.5, 9)
ax.set_ylim(-0.5, 8)
ax.set_aspect('equal')
ax.axis('off')
ax.set_title(f'{n}×{n} Systolic Array Architecture', fontsize=14, fontweight='bold')
# Annotations
ax.text(8, 6, 'Each PE performs:\nMAC operation\n\nacc += a × b\npass a right\npass b down',
fontsize=10, va='top', ha='left',
bbox=dict(boxstyle='round', facecolor='white', edgecolor='gray'))
plt.tight_layout()
plt.show()
print(f"A {n}×{n} systolic array has {n*n} = {n**2} PEs")
print(f"Each cycle, it performs {n**2} MAC operations = {2*n**2} FLOPs")
print(f"TPU v4 has 128×128 = 16,384 PEs!")
A 4×4 systolic array has 16 = 16 PEs Each cycle, it performs 16 MAC operations = 32 FLOPs TPU v4 has 128×128 = 16,384 PEs!
3. Matrix Multiplication on Systolic Arrays¶
Consider multiplying matrices: $C = A \times B$
Where:
- $A$ is $M \times K$
- $B$ is $K \times N$
- $C$ is $M \times N$
The Algorithm¶
For a $P \times P$ systolic array computing a $P \times P$ output:
- Setup: Each PE initializes its accumulator to 0
- Stream A rows from the left, staggered by one cycle per row
- Stream B columns from the top, staggered by one cycle per column
- Each cycle: Every PE multiplies its inputs and adds to its accumulator
- After K+2P-2 cycles: All results are computed
Data reuse¶
Each element of A is used P times (once per column of B). Each element of B is used P times (once per row of A).
Total data read: $P \times K$ (for A) + $K \times P$ (for B) = $2PK$ elements
Total compute: $2P^2K$ FLOPs (P² outputs, each requiring K multiply-adds)
Arithmetic intensity = $\frac{2P^2K}{2PK \times \text{bytes/elem}} = \frac{P}{\text{bytes/elem}}$
For a 128×128 array with FP16: intensity = 128/2 = 64 FLOPs/byte
In [3]:
Copied!
# Visualize: Step-by-step matrix multiplication on a 3x3 systolic array
# We'll show how data flows through to compute C = A × B
def visualize_systolic_step(ax, step, n=3, k=3):
"""Visualize one step of systolic array execution."""
ax.clear()
pe_size = 0.8
spacing = 1.2
# Example matrices
A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
B = np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]])
# Track which data is at each PE at this step
# Data arrives staggered: row i arrives at time i, col j arrives at time j
for i in range(n):
for j in range(n):
x = j * spacing + 1.5
y = (n - 1 - i) * spacing + 1
# Calculate what's at this PE
# A element from row i arrives at column j at time (i + j)
# The k-th element of row i arrives at time (i + k)
# At column j, we see the (step - i - j)th element
k_idx = step - i - j
if 0 <= k_idx < k:
a_val = A[i, k_idx]
b_val = B[k_idx, j]
active = True
else:
a_val = None
b_val = None
active = False
# PE box
color = COLORS['pe'] if active else COLORS['light']
pe = patches.FancyBboxPatch((x - pe_size/2, y - pe_size/2), pe_size, pe_size,
boxstyle="round,pad=0.02",
facecolor=color, alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(pe)
# Show values
if active:
ax.text(x, y + 0.15, f'{a_val}×{b_val}', ha='center', va='center', fontsize=9, fontweight='bold')
ax.text(x, y - 0.2, f'={a_val*b_val}', ha='center', va='center', fontsize=8, color=COLORS['success'])
else:
ax.text(x, y, '-', ha='center', va='center', fontsize=10, color='gray')
# Title
ax.set_title(f'Cycle {step}: Systolic Array Execution', fontsize=12, fontweight='bold')
ax.set_xlim(0, 5)
ax.set_ylim(-0.5, 4)
ax.set_aspect('equal')
ax.axis('off')
# Create figure showing multiple time steps
fig, axes = plt.subplots(2, 3, figsize=(14, 8))
for idx, step in enumerate([0, 1, 2, 3, 4, 5]):
ax = axes[idx // 3, idx % 3]
visualize_systolic_step(ax, step)
plt.suptitle('Matrix Multiply C = A×B on 3×3 Systolic Array\n(showing which products are computed at each PE)',
fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("Cycle 0: Only PE(0,0) is active")
print("Cycle 1: Diagonal PE(0,1) and PE(1,0) become active")
print("Cycle 2: All PEs on main anti-diagonal active")
print("...")
print("After cycle 4: All partial products computed, accumulators contain final results")
# Visualize: Step-by-step matrix multiplication on a 3x3 systolic array
# We'll show how data flows through to compute C = A × B
def visualize_systolic_step(ax, step, n=3, k=3):
"""Visualize one step of systolic array execution."""
ax.clear()
pe_size = 0.8
spacing = 1.2
# Example matrices
A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
B = np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]])
# Track which data is at each PE at this step
# Data arrives staggered: row i arrives at time i, col j arrives at time j
for i in range(n):
for j in range(n):
x = j * spacing + 1.5
y = (n - 1 - i) * spacing + 1
# Calculate what's at this PE
# A element from row i arrives at column j at time (i + j)
# The k-th element of row i arrives at time (i + k)
# At column j, we see the (step - i - j)th element
k_idx = step - i - j
if 0 <= k_idx < k:
a_val = A[i, k_idx]
b_val = B[k_idx, j]
active = True
else:
a_val = None
b_val = None
active = False
# PE box
color = COLORS['pe'] if active else COLORS['light']
pe = patches.FancyBboxPatch((x - pe_size/2, y - pe_size/2), pe_size, pe_size,
boxstyle="round,pad=0.02",
facecolor=color, alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(pe)
# Show values
if active:
ax.text(x, y + 0.15, f'{a_val}×{b_val}', ha='center', va='center', fontsize=9, fontweight='bold')
ax.text(x, y - 0.2, f'={a_val*b_val}', ha='center', va='center', fontsize=8, color=COLORS['success'])
else:
ax.text(x, y, '-', ha='center', va='center', fontsize=10, color='gray')
# Title
ax.set_title(f'Cycle {step}: Systolic Array Execution', fontsize=12, fontweight='bold')
ax.set_xlim(0, 5)
ax.set_ylim(-0.5, 4)
ax.set_aspect('equal')
ax.axis('off')
# Create figure showing multiple time steps
fig, axes = plt.subplots(2, 3, figsize=(14, 8))
for idx, step in enumerate([0, 1, 2, 3, 4, 5]):
ax = axes[idx // 3, idx % 3]
visualize_systolic_step(ax, step)
plt.suptitle('Matrix Multiply C = A×B on 3×3 Systolic Array\n(showing which products are computed at each PE)',
fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("Cycle 0: Only PE(0,0) is active")
print("Cycle 1: Diagonal PE(0,1) and PE(1,0) become active")
print("Cycle 2: All PEs on main anti-diagonal active")
print("...")
print("After cycle 4: All partial products computed, accumulators contain final results")
Cycle 0: Only PE(0,0) is active Cycle 1: Diagonal PE(0,1) and PE(1,0) become active Cycle 2: All PEs on main anti-diagonal active ... After cycle 4: All partial products computed, accumulators contain final results
In [4]:
Copied!
# Calculate: Timing and efficiency of systolic array execution
def systolic_analysis(P, K, clock_freq_ghz=1.0, dtype_bytes=2):
"""Analyze systolic array performance for P×P array computing P×K × K×P matmul."""
# Timing
fill_cycles = P - 1 # Time to fill the array
compute_cycles = K # Time to stream all K elements
drain_cycles = P - 1 # Time to drain results
total_cycles = fill_cycles + compute_cycles + drain_cycles
# Alternative: overlapped execution (steady state)
steady_state_cycles = K # In steady state, just K cycles
# Compute
total_macs = P * P * K
total_flops = 2 * total_macs
# Data movement (assuming each element read once from external memory)
a_bytes = P * K * dtype_bytes # A matrix
b_bytes = K * P * dtype_bytes # B matrix
c_bytes = P * P * dtype_bytes # C matrix (output)
total_bytes_in = a_bytes + b_bytes
total_bytes_out = c_bytes
# Arithmetic intensity
arithmetic_intensity = total_flops / (total_bytes_in + total_bytes_out)
# Utilization
# Maximum possible: P² MACs per cycle
peak_macs = P * P * total_cycles
utilization = total_macs / peak_macs
# Steady state utilization (after warmup)
steady_utilization = total_macs / (P * P * K)
# Performance
time_ns = total_cycles / clock_freq_ghz
tflops = total_flops / (time_ns * 1e-9) / 1e12
return {
'total_cycles': total_cycles,
'steady_cycles': steady_state_cycles,
'total_flops': total_flops,
'bytes_in': total_bytes_in,
'bytes_out': total_bytes_out,
'arithmetic_intensity': arithmetic_intensity,
'utilization': utilization,
'steady_utilization': steady_utilization,
'time_ns': time_ns,
'tflops': tflops,
}
print("="*70)
print("SYSTOLIC ARRAY PERFORMANCE ANALYSIS")
print("="*70)
configs = [
(4, 4, "Tiny (4×4)"),
(16, 64, "Small (16×16)"),
(128, 128, "TPU v2/v3 (128×128)"),
(128, 4096, "TPU with d_model=4096"),
]
for P, K, name in configs:
stats = systolic_analysis(P, K)
print(f"\n{name}: {P}×{P} array, K={K}")
print(f" Cycles: {stats['total_cycles']} total ({P-1} fill + {K} compute + {P-1} drain)")
print(f" FLOPs: {stats['total_flops']:,} ({stats['total_flops']/1e6:.2f}M)")
print(f" Data: {stats['bytes_in']/1024:.1f} KB in, {stats['bytes_out']/1024:.1f} KB out")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.1f} FLOPs/byte")
print(f" Utilization: {stats['utilization']*100:.1f}% (amortized) / {stats['steady_utilization']*100:.1f}% (steady)")
# Calculate: Timing and efficiency of systolic array execution
def systolic_analysis(P, K, clock_freq_ghz=1.0, dtype_bytes=2):
"""Analyze systolic array performance for P×P array computing P×K × K×P matmul."""
# Timing
fill_cycles = P - 1 # Time to fill the array
compute_cycles = K # Time to stream all K elements
drain_cycles = P - 1 # Time to drain results
total_cycles = fill_cycles + compute_cycles + drain_cycles
# Alternative: overlapped execution (steady state)
steady_state_cycles = K # In steady state, just K cycles
# Compute
total_macs = P * P * K
total_flops = 2 * total_macs
# Data movement (assuming each element read once from external memory)
a_bytes = P * K * dtype_bytes # A matrix
b_bytes = K * P * dtype_bytes # B matrix
c_bytes = P * P * dtype_bytes # C matrix (output)
total_bytes_in = a_bytes + b_bytes
total_bytes_out = c_bytes
# Arithmetic intensity
arithmetic_intensity = total_flops / (total_bytes_in + total_bytes_out)
# Utilization
# Maximum possible: P² MACs per cycle
peak_macs = P * P * total_cycles
utilization = total_macs / peak_macs
# Steady state utilization (after warmup)
steady_utilization = total_macs / (P * P * K)
# Performance
time_ns = total_cycles / clock_freq_ghz
tflops = total_flops / (time_ns * 1e-9) / 1e12
return {
'total_cycles': total_cycles,
'steady_cycles': steady_state_cycles,
'total_flops': total_flops,
'bytes_in': total_bytes_in,
'bytes_out': total_bytes_out,
'arithmetic_intensity': arithmetic_intensity,
'utilization': utilization,
'steady_utilization': steady_utilization,
'time_ns': time_ns,
'tflops': tflops,
}
print("="*70)
print("SYSTOLIC ARRAY PERFORMANCE ANALYSIS")
print("="*70)
configs = [
(4, 4, "Tiny (4×4)"),
(16, 64, "Small (16×16)"),
(128, 128, "TPU v2/v3 (128×128)"),
(128, 4096, "TPU with d_model=4096"),
]
for P, K, name in configs:
stats = systolic_analysis(P, K)
print(f"\n{name}: {P}×{P} array, K={K}")
print(f" Cycles: {stats['total_cycles']} total ({P-1} fill + {K} compute + {P-1} drain)")
print(f" FLOPs: {stats['total_flops']:,} ({stats['total_flops']/1e6:.2f}M)")
print(f" Data: {stats['bytes_in']/1024:.1f} KB in, {stats['bytes_out']/1024:.1f} KB out")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.1f} FLOPs/byte")
print(f" Utilization: {stats['utilization']*100:.1f}% (amortized) / {stats['steady_utilization']*100:.1f}% (steady)")
====================================================================== SYSTOLIC ARRAY PERFORMANCE ANALYSIS ====================================================================== Tiny (4×4): 4×4 array, K=4 Cycles: 10 total (3 fill + 4 compute + 3 drain) FLOPs: 128 (0.00M) Data: 0.1 KB in, 0.0 KB out Arithmetic intensity: 1.3 FLOPs/byte Utilization: 40.0% (amortized) / 100.0% (steady) Small (16×16): 16×16 array, K=64 Cycles: 94 total (15 fill + 64 compute + 15 drain) FLOPs: 32,768 (0.03M) Data: 4.0 KB in, 0.5 KB out Arithmetic intensity: 7.1 FLOPs/byte Utilization: 68.1% (amortized) / 100.0% (steady) TPU v2/v3 (128×128): 128×128 array, K=128 Cycles: 382 total (127 fill + 128 compute + 127 drain) FLOPs: 4,194,304 (4.19M) Data: 64.0 KB in, 32.0 KB out Arithmetic intensity: 42.7 FLOPs/byte Utilization: 33.5% (amortized) / 100.0% (steady) TPU with d_model=4096: 128×128 array, K=4096 Cycles: 4350 total (127 fill + 4096 compute + 127 drain) FLOPs: 134,217,728 (134.22M) Data: 2048.0 KB in, 32.0 KB out Arithmetic intensity: 63.0 FLOPs/byte Utilization: 94.2% (amortized) / 100.0% (steady)
4. Google TPU Architecture¶
The Tensor Processing Unit (TPU) is Google's custom ASIC for machine learning, built around a large systolic array.
TPU v1 (2016) - Inference Only¶
| Component | Specification |
|---|---|
| Systolic array | 256 × 256 = 65,536 MACs |
| Clock speed | 700 MHz |
| Peak performance | 92 TOPS (INT8) |
| On-chip memory | 28 MB |
| Memory bandwidth | 34 GB/s (DDR3) |
| Power | 40W |
TPU v4 (2021) - Training and Inference¶
| Component | Specification |
|---|---|
| Systolic arrays | 2 × (128 × 128) = 32,768 MACs |
| Clock speed | ~1.1 GHz |
| Peak performance | 275 TFLOPS (BF16) |
| HBM capacity | 32 GB |
| Memory bandwidth | 1.2 TB/s |
| Power | ~200W |
Architecture Diagram¶
In [5]:
Copied!
# Visualize: TPU architecture block diagram
fig, ax = plt.subplots(figsize=(14, 10))
def draw_component(ax, x, y, w, h, label, color, sublabel=''):
rect = patches.FancyBboxPatch((x, y), w, h,
boxstyle="round,pad=0.02",
facecolor=color, alpha=0.85,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
if sublabel:
ax.text(x + w/2, y + h/2 + 0.15, label, ha='center', va='center',
fontsize=10, fontweight='bold')
ax.text(x + w/2, y + h/2 - 0.15, sublabel, ha='center', va='center', fontsize=8)
else:
ax.text(x + w/2, y + h/2, label, ha='center', va='center',
fontsize=10, fontweight='bold')
# Host interface
draw_component(ax, 0.5, 8, 2, 0.8, 'Host Interface', COLORS['light'], 'PCIe')
# Control
draw_component(ax, 0.5, 6.5, 2, 1, 'Control Unit', COLORS['secondary'], 'Instruction\ndecode')
# Weight FIFO
draw_component(ax, 3.5, 7.5, 2, 1.3, 'Weight FIFO', COLORS['weight'], '256×256×8b')
# Matrix Multiply Unit (main systolic array)
draw_component(ax, 3, 4, 5, 3, 'Matrix Multiply Unit', COLORS['pe'], '256×256 Systolic Array\n65,536 MACs\n92 TOPS (INT8)')
# Unified Buffer
draw_component(ax, 9, 4, 2.5, 3, 'Unified Buffer', COLORS['activation'], '24 MB SRAM\n(Activations)')
# Accumulators
draw_component(ax, 3, 2, 5, 1.2, 'Accumulators', COLORS['accumulator'], '4 MB (32-bit)')
# Activation Unit
draw_component(ax, 3, 0.3, 3, 1.2, 'Activation Unit', COLORS['success'], 'ReLU, Sigmoid, etc.')
# Normalize/Pool
draw_component(ax, 6.5, 0.3, 3, 1.2, 'Normalize/Pool', COLORS['cyan'], 'Softmax, LayerNorm')
# DDR Memory
draw_component(ax, 0.5, 0.3, 2, 1.2, 'DDR3 Memory', COLORS['light'], '8 GB, 34 GB/s')
# Arrows showing data flow
# Host to Control
ax.annotate('', xy=(1.5, 7.5), xytext=(1.5, 8), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Weight FIFO to MMU
ax.annotate('', xy=(4.5, 7), xytext=(4.5, 7.5), arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=3))
# Unified Buffer to MMU
ax.annotate('', xy=(8, 5.5), xytext=(9, 5.5), arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=3))
# MMU to Accumulators
ax.annotate('', xy=(5.5, 3.2), xytext=(5.5, 4), arrowprops=dict(arrowstyle='->', color=COLORS['accumulator'], lw=3))
# Accumulators to Activation
ax.annotate('', xy=(4.5, 1.5), xytext=(4.5, 2), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Activation to Normalize
ax.annotate('', xy=(6.5, 0.9), xytext=(6, 0.9), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Normalize back to Unified Buffer
ax.annotate('', xy=(10.25, 4), xytext=(8, 1.5),
arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=-0.3', color=COLORS['activation'], lw=2))
# DDR to Weight FIFO
ax.annotate('', xy=(3.5, 8.1), xytext=(2.5, 0.9),
arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.5', color=COLORS['weight'], lw=2))
ax.set_xlim(-0.5, 12.5)
ax.set_ylim(-0.5, 9.5)
ax.axis('off')
ax.set_title('TPU v1 Architecture Block Diagram', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()
print("Key insight: The systolic array (MMU) is the heart of the TPU.")
print("All other components exist to feed it data and handle results.")
print("TPU design philosophy: Maximize utilization of the matrix multiply unit.")
# Visualize: TPU architecture block diagram
fig, ax = plt.subplots(figsize=(14, 10))
def draw_component(ax, x, y, w, h, label, color, sublabel=''):
rect = patches.FancyBboxPatch((x, y), w, h,
boxstyle="round,pad=0.02",
facecolor=color, alpha=0.85,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
if sublabel:
ax.text(x + w/2, y + h/2 + 0.15, label, ha='center', va='center',
fontsize=10, fontweight='bold')
ax.text(x + w/2, y + h/2 - 0.15, sublabel, ha='center', va='center', fontsize=8)
else:
ax.text(x + w/2, y + h/2, label, ha='center', va='center',
fontsize=10, fontweight='bold')
# Host interface
draw_component(ax, 0.5, 8, 2, 0.8, 'Host Interface', COLORS['light'], 'PCIe')
# Control
draw_component(ax, 0.5, 6.5, 2, 1, 'Control Unit', COLORS['secondary'], 'Instruction\ndecode')
# Weight FIFO
draw_component(ax, 3.5, 7.5, 2, 1.3, 'Weight FIFO', COLORS['weight'], '256×256×8b')
# Matrix Multiply Unit (main systolic array)
draw_component(ax, 3, 4, 5, 3, 'Matrix Multiply Unit', COLORS['pe'], '256×256 Systolic Array\n65,536 MACs\n92 TOPS (INT8)')
# Unified Buffer
draw_component(ax, 9, 4, 2.5, 3, 'Unified Buffer', COLORS['activation'], '24 MB SRAM\n(Activations)')
# Accumulators
draw_component(ax, 3, 2, 5, 1.2, 'Accumulators', COLORS['accumulator'], '4 MB (32-bit)')
# Activation Unit
draw_component(ax, 3, 0.3, 3, 1.2, 'Activation Unit', COLORS['success'], 'ReLU, Sigmoid, etc.')
# Normalize/Pool
draw_component(ax, 6.5, 0.3, 3, 1.2, 'Normalize/Pool', COLORS['cyan'], 'Softmax, LayerNorm')
# DDR Memory
draw_component(ax, 0.5, 0.3, 2, 1.2, 'DDR3 Memory', COLORS['light'], '8 GB, 34 GB/s')
# Arrows showing data flow
# Host to Control
ax.annotate('', xy=(1.5, 7.5), xytext=(1.5, 8), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Weight FIFO to MMU
ax.annotate('', xy=(4.5, 7), xytext=(4.5, 7.5), arrowprops=dict(arrowstyle='->', color=COLORS['weight'], lw=3))
# Unified Buffer to MMU
ax.annotate('', xy=(8, 5.5), xytext=(9, 5.5), arrowprops=dict(arrowstyle='->', color=COLORS['activation'], lw=3))
# MMU to Accumulators
ax.annotate('', xy=(5.5, 3.2), xytext=(5.5, 4), arrowprops=dict(arrowstyle='->', color=COLORS['accumulator'], lw=3))
# Accumulators to Activation
ax.annotate('', xy=(4.5, 1.5), xytext=(4.5, 2), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Activation to Normalize
ax.annotate('', xy=(6.5, 0.9), xytext=(6, 0.9), arrowprops=dict(arrowstyle='->', color='black', lw=2))
# Normalize back to Unified Buffer
ax.annotate('', xy=(10.25, 4), xytext=(8, 1.5),
arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=-0.3', color=COLORS['activation'], lw=2))
# DDR to Weight FIFO
ax.annotate('', xy=(3.5, 8.1), xytext=(2.5, 0.9),
arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.5', color=COLORS['weight'], lw=2))
ax.set_xlim(-0.5, 12.5)
ax.set_ylim(-0.5, 9.5)
ax.axis('off')
ax.set_title('TPU v1 Architecture Block Diagram', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()
print("Key insight: The systolic array (MMU) is the heart of the TPU.")
print("All other components exist to feed it data and handle results.")
print("TPU design philosophy: Maximize utilization of the matrix multiply unit.")
Key insight: The systolic array (MMU) is the heart of the TPU. All other components exist to feed it data and handle results. TPU design philosophy: Maximize utilization of the matrix multiply unit.
5. Tiling Large Matrices¶
Real transformer matrices are much larger than any systolic array:
| Matrix | Typical Size | vs 128×128 Array |
|---|---|---|
| $W_Q$ in Llama 7B | 4096 × 4096 | 32 × 32 = 1024 tiles |
| FFN up in Llama 7B | 4096 × 11008 | 32 × 86 = 2752 tiles |
| Attention scores (8K ctx) | 8192 × 8192 | 64 × 64 = 4096 tiles |
Tiling Strategy¶
To compute $C = A \times B$ where all matrices exceed array size:
- Partition A into row-blocks of size P
- Partition B into column-blocks of size P
- Partition the K dimension into blocks of size K_tile
- Accumulate partial results across K-tiles
$$C[i:i+P, j:j+P] = \sum_{k=0}^{K/K_{tile}-1} A[i:i+P, k \cdot K_{tile}:(k+1) \cdot K_{tile}] \times B[k \cdot K_{tile}:(k+1) \cdot K_{tile}, j:j+P]$$
Tiling Trade-offs¶
| Factor | Smaller Tiles | Larger Tiles |
|---|---|---|
| On-chip memory | Less required | More required |
| Data reuse | Lower | Higher |
| Memory bandwidth | Higher pressure | Lower pressure |
| Fill/drain overhead | Higher | Lower |
In [6]:
Copied!
# Visualize: Tiling a large matrix multiply
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# Example: 12×12 matrices tiled for 4×4 systolic array
M, N, K = 12, 12, 12
P = 4 # Tile size (array size)
def draw_tiled_matrix(ax, rows, cols, tile_size, title, highlight_tile=None):
"""Draw a matrix with tile boundaries."""
ax.set_xlim(-0.5, cols + 0.5)
ax.set_ylim(-0.5, rows + 0.5)
# Draw tiles
n_row_tiles = rows // tile_size
n_col_tiles = cols // tile_size
for i in range(n_row_tiles):
for j in range(n_col_tiles):
x = j * tile_size
y = rows - (i + 1) * tile_size
is_highlighted = (highlight_tile is not None and
highlight_tile[0] == i and highlight_tile[1] == j)
color = COLORS['primary'] if is_highlighted else COLORS['light']
alpha = 0.8 if is_highlighted else 0.3
rect = patches.Rectangle((x, y), tile_size, tile_size,
facecolor=color, alpha=alpha,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
# Tile label
ax.text(x + tile_size/2, y + tile_size/2, f'({i},{j})',
ha='center', va='center', fontsize=9)
ax.set_title(title, fontsize=12, fontweight='bold')
ax.set_aspect('equal')
ax.set_xlabel(f'{cols}', fontsize=10)
ax.set_ylabel(f'{rows}', fontsize=10)
ax.set_xticks([])
ax.set_yticks([])
# Show A, B, C with highlighted tiles for one output tile computation
# To compute C[1,2], we need A[1,:] × B[:,2]
output_tile = (1, 2)
# Matrix A - highlight row block
draw_tiled_matrix(axes[0], M, K, P, f'Matrix A ({M}×{K})', highlight_tile=(1, None))
# Highlight entire row
for j in range(K // P):
rect = patches.Rectangle((j * P, M - 2*P), P, P,
facecolor=COLORS['primary'], alpha=0.5,
edgecolor=COLORS['danger'], linewidth=3, linestyle='--')
axes[0].add_patch(rect)
axes[0].text(K/2, M - P - P/2, 'Row tiles for C[1,:]', ha='center', va='center', fontsize=10,
bbox=dict(facecolor='white', edgecolor=COLORS['danger'], boxstyle='round'))
# Matrix B - highlight column block
draw_tiled_matrix(axes[1], K, N, P, f'Matrix B ({K}×{N})', highlight_tile=(None, 2))
for i in range(K // P):
rect = patches.Rectangle((2 * P, K - (i+1)*P), P, P,
facecolor=COLORS['secondary'], alpha=0.5,
edgecolor=COLORS['danger'], linewidth=3, linestyle='--')
axes[1].add_patch(rect)
axes[1].text(2*P + P/2, K/2, 'Col tiles\nfor C[:,2]', ha='center', va='center', fontsize=10,
bbox=dict(facecolor='white', edgecolor=COLORS['danger'], boxstyle='round'))
# Matrix C - highlight output tile
draw_tiled_matrix(axes[2], M, N, P, f'Matrix C ({M}×{N})', highlight_tile=output_tile)
plt.suptitle(f'Tiling {M}×{K} × {K}×{N} for {P}×{P} Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
n_tiles_total = (M // P) * (N // P) * (K // P)
print(f"\nTo compute full C matrix:")
print(f" Output tiles: {M//P} × {N//P} = {(M//P)*(N//P)} tiles")
print(f" K-dimension tiles: {K//P}")
print(f" Total tile operations: {n_tiles_total}")
print(f" Each tile: {P}×{P}×{P} = {P**3} MACs")
print(f" Total MACs: {n_tiles_total * P**3:,} = {M*N*K:,} (matches M×N×K)")
# Visualize: Tiling a large matrix multiply
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# Example: 12×12 matrices tiled for 4×4 systolic array
M, N, K = 12, 12, 12
P = 4 # Tile size (array size)
def draw_tiled_matrix(ax, rows, cols, tile_size, title, highlight_tile=None):
"""Draw a matrix with tile boundaries."""
ax.set_xlim(-0.5, cols + 0.5)
ax.set_ylim(-0.5, rows + 0.5)
# Draw tiles
n_row_tiles = rows // tile_size
n_col_tiles = cols // tile_size
for i in range(n_row_tiles):
for j in range(n_col_tiles):
x = j * tile_size
y = rows - (i + 1) * tile_size
is_highlighted = (highlight_tile is not None and
highlight_tile[0] == i and highlight_tile[1] == j)
color = COLORS['primary'] if is_highlighted else COLORS['light']
alpha = 0.8 if is_highlighted else 0.3
rect = patches.Rectangle((x, y), tile_size, tile_size,
facecolor=color, alpha=alpha,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
# Tile label
ax.text(x + tile_size/2, y + tile_size/2, f'({i},{j})',
ha='center', va='center', fontsize=9)
ax.set_title(title, fontsize=12, fontweight='bold')
ax.set_aspect('equal')
ax.set_xlabel(f'{cols}', fontsize=10)
ax.set_ylabel(f'{rows}', fontsize=10)
ax.set_xticks([])
ax.set_yticks([])
# Show A, B, C with highlighted tiles for one output tile computation
# To compute C[1,2], we need A[1,:] × B[:,2]
output_tile = (1, 2)
# Matrix A - highlight row block
draw_tiled_matrix(axes[0], M, K, P, f'Matrix A ({M}×{K})', highlight_tile=(1, None))
# Highlight entire row
for j in range(K // P):
rect = patches.Rectangle((j * P, M - 2*P), P, P,
facecolor=COLORS['primary'], alpha=0.5,
edgecolor=COLORS['danger'], linewidth=3, linestyle='--')
axes[0].add_patch(rect)
axes[0].text(K/2, M - P - P/2, 'Row tiles for C[1,:]', ha='center', va='center', fontsize=10,
bbox=dict(facecolor='white', edgecolor=COLORS['danger'], boxstyle='round'))
# Matrix B - highlight column block
draw_tiled_matrix(axes[1], K, N, P, f'Matrix B ({K}×{N})', highlight_tile=(None, 2))
for i in range(K // P):
rect = patches.Rectangle((2 * P, K - (i+1)*P), P, P,
facecolor=COLORS['secondary'], alpha=0.5,
edgecolor=COLORS['danger'], linewidth=3, linestyle='--')
axes[1].add_patch(rect)
axes[1].text(2*P + P/2, K/2, 'Col tiles\nfor C[:,2]', ha='center', va='center', fontsize=10,
bbox=dict(facecolor='white', edgecolor=COLORS['danger'], boxstyle='round'))
# Matrix C - highlight output tile
draw_tiled_matrix(axes[2], M, N, P, f'Matrix C ({M}×{N})', highlight_tile=output_tile)
plt.suptitle(f'Tiling {M}×{K} × {K}×{N} for {P}×{P} Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
n_tiles_total = (M // P) * (N // P) * (K // P)
print(f"\nTo compute full C matrix:")
print(f" Output tiles: {M//P} × {N//P} = {(M//P)*(N//P)} tiles")
print(f" K-dimension tiles: {K//P}")
print(f" Total tile operations: {n_tiles_total}")
print(f" Each tile: {P}×{P}×{P} = {P**3} MACs")
print(f" Total MACs: {n_tiles_total * P**3:,} = {M*N*K:,} (matches M×N×K)")
To compute full C matrix: Output tiles: 3 × 3 = 9 tiles K-dimension tiles: 3 Total tile operations: 27 Each tile: 4×4×4 = 64 MACs Total MACs: 1,728 = 1,728 (matches M×N×K)
In [7]:
Copied!
# Calculate: Tiling requirements for real transformer dimensions
def analyze_tiling(M, N, K, P, dtype_bytes=2):
"""Analyze tiling requirements for M×K × K×N with P×P array."""
# Number of tiles
m_tiles = (M + P - 1) // P # Ceiling division
n_tiles = (N + P - 1) // P
k_tiles = (K + P - 1) // P
# Total tile operations
total_tile_ops = m_tiles * n_tiles * k_tiles
# On-chip memory needed per tile operation
a_tile_bytes = P * P * dtype_bytes # One A tile
b_tile_bytes = P * P * dtype_bytes # One B tile
c_tile_bytes = P * P * dtype_bytes # One C tile (accumulator)
# Minimum memory for double buffering
min_memory = 2 * (a_tile_bytes + b_tile_bytes) + c_tile_bytes
# Data movement with perfect reuse
# A tiles: each of m_tiles×k_tiles A-tiles read once, reused for all n_tiles
# B tiles: each of k_tiles×n_tiles B-tiles read once, reused for all m_tiles
a_total_bytes = M * K * dtype_bytes
b_total_bytes = K * N * dtype_bytes
c_total_bytes = M * N * dtype_bytes
# Total compute
total_flops = 2 * M * N * K
# Cycles (assuming perfect pipelining)
cycles_per_tile = P + K // k_tiles + P - 2 # Simplified
total_cycles_naive = total_tile_ops * (P + P + P - 2) # Very conservative
total_cycles_pipelined = m_tiles * n_tiles * k_tiles * P # Approximate
return {
'tile_dims': (m_tiles, n_tiles, k_tiles),
'total_tile_ops': total_tile_ops,
'min_memory_kb': min_memory / 1024,
'a_bytes': a_total_bytes,
'b_bytes': b_total_bytes,
'c_bytes': c_total_bytes,
'total_flops': total_flops,
'arithmetic_intensity': total_flops / (a_total_bytes + b_total_bytes + c_total_bytes),
}
print("="*80)
print("TILING ANALYSIS FOR TRANSFORMER OPERATIONS (128×128 systolic array, BF16)")
print("="*80)
# Llama 2 7B operations
operations = [
('QKV Projection (batch=1)', 1, 12288, 4096), # d_model to 3×d_model
('QKV Projection (batch=32)', 32, 12288, 4096),
('Attention QK^T (8K ctx)', 8192, 8192, 128), # Per head
('FFN Up (batch=1)', 1, 11008, 4096),
('FFN Up (batch=32)', 32, 11008, 4096),
('FFN Down (batch=1)', 1, 4096, 11008),
]
P = 128 # TPU-like array size
for name, M, N, K in operations:
stats = analyze_tiling(M, N, K, P)
print(f"\n{name}: [{M}×{K}] × [{K}×{N}]")
print(f" Tiles: {stats['tile_dims'][0]}×{stats['tile_dims'][1]}×{stats['tile_dims'][2]} = {stats['total_tile_ops']} ops")
print(f" FLOPs: {stats['total_flops']/1e6:.1f}M")
print(f" Data: A={stats['a_bytes']/1024:.0f}KB, B={stats['b_bytes']/1024:.0f}KB, C={stats['c_bytes']/1024:.0f}KB")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.1f} FLOP/byte")
print(f" Min on-chip memory: {stats['min_memory_kb']:.0f} KB")
# Calculate: Tiling requirements for real transformer dimensions
def analyze_tiling(M, N, K, P, dtype_bytes=2):
"""Analyze tiling requirements for M×K × K×N with P×P array."""
# Number of tiles
m_tiles = (M + P - 1) // P # Ceiling division
n_tiles = (N + P - 1) // P
k_tiles = (K + P - 1) // P
# Total tile operations
total_tile_ops = m_tiles * n_tiles * k_tiles
# On-chip memory needed per tile operation
a_tile_bytes = P * P * dtype_bytes # One A tile
b_tile_bytes = P * P * dtype_bytes # One B tile
c_tile_bytes = P * P * dtype_bytes # One C tile (accumulator)
# Minimum memory for double buffering
min_memory = 2 * (a_tile_bytes + b_tile_bytes) + c_tile_bytes
# Data movement with perfect reuse
# A tiles: each of m_tiles×k_tiles A-tiles read once, reused for all n_tiles
# B tiles: each of k_tiles×n_tiles B-tiles read once, reused for all m_tiles
a_total_bytes = M * K * dtype_bytes
b_total_bytes = K * N * dtype_bytes
c_total_bytes = M * N * dtype_bytes
# Total compute
total_flops = 2 * M * N * K
# Cycles (assuming perfect pipelining)
cycles_per_tile = P + K // k_tiles + P - 2 # Simplified
total_cycles_naive = total_tile_ops * (P + P + P - 2) # Very conservative
total_cycles_pipelined = m_tiles * n_tiles * k_tiles * P # Approximate
return {
'tile_dims': (m_tiles, n_tiles, k_tiles),
'total_tile_ops': total_tile_ops,
'min_memory_kb': min_memory / 1024,
'a_bytes': a_total_bytes,
'b_bytes': b_total_bytes,
'c_bytes': c_total_bytes,
'total_flops': total_flops,
'arithmetic_intensity': total_flops / (a_total_bytes + b_total_bytes + c_total_bytes),
}
print("="*80)
print("TILING ANALYSIS FOR TRANSFORMER OPERATIONS (128×128 systolic array, BF16)")
print("="*80)
# Llama 2 7B operations
operations = [
('QKV Projection (batch=1)', 1, 12288, 4096), # d_model to 3×d_model
('QKV Projection (batch=32)', 32, 12288, 4096),
('Attention QK^T (8K ctx)', 8192, 8192, 128), # Per head
('FFN Up (batch=1)', 1, 11008, 4096),
('FFN Up (batch=32)', 32, 11008, 4096),
('FFN Down (batch=1)', 1, 4096, 11008),
]
P = 128 # TPU-like array size
for name, M, N, K in operations:
stats = analyze_tiling(M, N, K, P)
print(f"\n{name}: [{M}×{K}] × [{K}×{N}]")
print(f" Tiles: {stats['tile_dims'][0]}×{stats['tile_dims'][1]}×{stats['tile_dims'][2]} = {stats['total_tile_ops']} ops")
print(f" FLOPs: {stats['total_flops']/1e6:.1f}M")
print(f" Data: A={stats['a_bytes']/1024:.0f}KB, B={stats['b_bytes']/1024:.0f}KB, C={stats['c_bytes']/1024:.0f}KB")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.1f} FLOP/byte")
print(f" Min on-chip memory: {stats['min_memory_kb']:.0f} KB")
================================================================================ TILING ANALYSIS FOR TRANSFORMER OPERATIONS (128×128 systolic array, BF16) ================================================================================ QKV Projection (batch=1): [1×4096] × [4096×12288] Tiles: 1×96×32 = 3072 ops FLOPs: 100.7M Data: A=8KB, B=98304KB, C=24KB Arithmetic intensity: 1.0 FLOP/byte Min on-chip memory: 160 KB QKV Projection (batch=32): [32×4096] × [4096×12288] Tiles: 1×96×32 = 3072 ops FLOPs: 3221.2M Data: A=256KB, B=98304KB, C=768KB Arithmetic intensity: 31.7 FLOP/byte Min on-chip memory: 160 KB Attention QK^T (8K ctx): [8192×128] × [128×8192] Tiles: 64×64×1 = 4096 ops FLOPs: 17179.9M Data: A=2048KB, B=2048KB, C=131072KB Arithmetic intensity: 124.1 FLOP/byte Min on-chip memory: 160 KB FFN Up (batch=1): [1×4096] × [4096×11008] Tiles: 1×86×32 = 2752 ops FLOPs: 90.2M Data: A=8KB, B=88064KB, C=22KB Arithmetic intensity: 1.0 FLOP/byte Min on-chip memory: 160 KB FFN Up (batch=32): [32×4096] × [4096×11008] Tiles: 1×86×32 = 2752 ops FLOPs: 2885.7M Data: A=256KB, B=88064KB, C=688KB Arithmetic intensity: 31.7 FLOP/byte Min on-chip memory: 160 KB FFN Down (batch=1): [1×11008] × [11008×4096] Tiles: 1×32×86 = 2752 ops FLOPs: 90.2M Data: A=22KB, B=88064KB, C=8KB Arithmetic intensity: 1.0 FLOP/byte Min on-chip memory: 160 KB
6. Data Movement Analysis¶
The key question for hardware design: How much data must move for each computation?
Three Levels of Data Movement¶
Off-chip → On-chip (DRAM/HBM to SRAM)
- Slowest, highest energy (~200× register access)
- This is the memory bandwidth bottleneck
On-chip SRAM → Array (Buffer to PE)
- ~10× register energy
- Must stream data to keep array fed
PE to PE (within systolic array)
- ~1× register energy (local wire)
- This is where systolic arrays shine
Data Movement for Matrix Multiply¶
For $C_{M \times N} = A_{M \times K} \times B_{K \times N}$:
| Data | Bytes | When Read | Reuse |
|---|---|---|---|
| A | M × K × dtype | Once (ideally) | Each row used N times |
| B | K × N × dtype | Once (ideally) | Each column used M times |
| C | M × N × dtype | Written once | - |
Total I/O: $(MK + KN + MN) \times$ dtype_bytes
Total Compute: $2MNK$ FLOPs
Arithmetic Intensity: $\frac{2MNK}{(MK + KN + MN) \times \text{bytes}}$
In [8]:
Copied!
# Visualize: Data movement in matrix multiplication
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
# Left: Conceptual data flow
ax = axes[0]
# DRAM
dram = patches.FancyBboxPatch((0.5, 5), 4, 1,
boxstyle="round,pad=0.02",
facecolor=COLORS['light'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(dram)
ax.text(2.5, 5.5, 'Off-chip Memory (HBM/DDR)\nWeights + KV Cache', ha='center', va='center', fontsize=10)
# On-chip SRAM
sram = patches.FancyBboxPatch((0.5, 3), 4, 1.2,
boxstyle="round,pad=0.02",
facecolor=COLORS['cyan'], alpha=0.7,
edgecolor='black', linewidth=2)
ax.add_patch(sram)
ax.text(2.5, 3.6, 'On-chip SRAM (Unified Buffer)\nTiles of A, B, Partial C', ha='center', va='center', fontsize=10)
# Systolic Array
array = patches.FancyBboxPatch((1, 0.5), 3, 1.8,
boxstyle="round,pad=0.02",
facecolor=COLORS['pe'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(array)
ax.text(2.5, 1.4, 'Systolic Array\n128×128 PEs', ha='center', va='center', fontsize=11, fontweight='bold')
# Arrows with bandwidth labels
ax.annotate('', xy=(2.5, 4.2), xytext=(2.5, 5),
arrowprops=dict(arrowstyle='<->', color=COLORS['danger'], lw=3))
ax.text(4.8, 4.6, 'Memory BW\n~3 TB/s (H100)', fontsize=9, color=COLORS['danger'])
ax.annotate('', xy=(2.5, 2.3), xytext=(2.5, 3),
arrowprops=dict(arrowstyle='<->', color=COLORS['success'], lw=3))
ax.text(4.8, 2.65, 'SRAM BW\n~30 TB/s', fontsize=9, color=COLORS['success'])
# Add energy costs
ax.text(-0.5, 5.5, '~200 pJ/byte', fontsize=9, color=COLORS['danger'])
ax.text(-0.5, 3.6, '~10 pJ/byte', fontsize=9, color=COLORS['success'])
ax.text(-0.5, 1.4, '~1 pJ/op', fontsize=9, color=COLORS['primary'])
ax.set_xlim(-1, 6)
ax.set_ylim(0, 6.5)
ax.axis('off')
ax.set_title('Memory Hierarchy & Data Movement', fontsize=12, fontweight='bold')
# Right: Bytes in vs out for different operations
ax = axes[1]
# Different scenarios
scenarios = {
'Training\n(large batch)': {'bytes_in': 1.0, 'bytes_out': 0.5, 'compute': 50},
'Inference\nbatch=32': {'bytes_in': 1.0, 'bytes_out': 0.05, 'compute': 2},
'Inference\nbatch=1': {'bytes_in': 1.0, 'bytes_out': 0.002, 'compute': 0.06},
}
x = np.arange(len(scenarios))
width = 0.35
bytes_in = [s['bytes_in'] for s in scenarios.values()]
bytes_out = [s['bytes_out'] for s in scenarios.values()]
bars1 = ax.bar(x - width/2, bytes_in, width, label='Bytes In (weights)', color=COLORS['weight'])
bars2 = ax.bar(x + width/2, bytes_out, width, label='Bytes Out (output)', color=COLORS['success'])
ax.set_ylabel('Relative Data Volume', fontsize=11)
ax.set_xticks(x)
ax.set_xticklabels(scenarios.keys(), fontsize=10)
ax.legend(fontsize=10)
ax.set_title('Input vs Output Data Volume', fontsize=12, fontweight='bold')
ax.set_yscale('log')
ax.grid(True, alpha=0.3, axis='y')
# Add compute/byte ratio
for i, (name, s) in enumerate(scenarios.items()):
ax.text(i, max(bytes_in[i], bytes_out[i]) * 2,
f"{s['compute']:.1f}\nFLOP/B", ha='center', fontsize=9,
bbox=dict(facecolor='white', edgecolor='gray', boxstyle='round'))
plt.tight_layout()
plt.show()
print("\nKey insight: For LLM inference with batch=1:")
print(" - We read ~70B parameters (140GB in FP16)")
print(" - We produce ~4KB output per token (one hidden state)")
print(" - Ratio: ~35,000,000:1 input to output!")
print(" - Almost all bandwidth is consumed reading weights, not writing results")
# Visualize: Data movement in matrix multiplication
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
# Left: Conceptual data flow
ax = axes[0]
# DRAM
dram = patches.FancyBboxPatch((0.5, 5), 4, 1,
boxstyle="round,pad=0.02",
facecolor=COLORS['light'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(dram)
ax.text(2.5, 5.5, 'Off-chip Memory (HBM/DDR)\nWeights + KV Cache', ha='center', va='center', fontsize=10)
# On-chip SRAM
sram = patches.FancyBboxPatch((0.5, 3), 4, 1.2,
boxstyle="round,pad=0.02",
facecolor=COLORS['cyan'], alpha=0.7,
edgecolor='black', linewidth=2)
ax.add_patch(sram)
ax.text(2.5, 3.6, 'On-chip SRAM (Unified Buffer)\nTiles of A, B, Partial C', ha='center', va='center', fontsize=10)
# Systolic Array
array = patches.FancyBboxPatch((1, 0.5), 3, 1.8,
boxstyle="round,pad=0.02",
facecolor=COLORS['pe'], alpha=0.9,
edgecolor='black', linewidth=2)
ax.add_patch(array)
ax.text(2.5, 1.4, 'Systolic Array\n128×128 PEs', ha='center', va='center', fontsize=11, fontweight='bold')
# Arrows with bandwidth labels
ax.annotate('', xy=(2.5, 4.2), xytext=(2.5, 5),
arrowprops=dict(arrowstyle='<->', color=COLORS['danger'], lw=3))
ax.text(4.8, 4.6, 'Memory BW\n~3 TB/s (H100)', fontsize=9, color=COLORS['danger'])
ax.annotate('', xy=(2.5, 2.3), xytext=(2.5, 3),
arrowprops=dict(arrowstyle='<->', color=COLORS['success'], lw=3))
ax.text(4.8, 2.65, 'SRAM BW\n~30 TB/s', fontsize=9, color=COLORS['success'])
# Add energy costs
ax.text(-0.5, 5.5, '~200 pJ/byte', fontsize=9, color=COLORS['danger'])
ax.text(-0.5, 3.6, '~10 pJ/byte', fontsize=9, color=COLORS['success'])
ax.text(-0.5, 1.4, '~1 pJ/op', fontsize=9, color=COLORS['primary'])
ax.set_xlim(-1, 6)
ax.set_ylim(0, 6.5)
ax.axis('off')
ax.set_title('Memory Hierarchy & Data Movement', fontsize=12, fontweight='bold')
# Right: Bytes in vs out for different operations
ax = axes[1]
# Different scenarios
scenarios = {
'Training\n(large batch)': {'bytes_in': 1.0, 'bytes_out': 0.5, 'compute': 50},
'Inference\nbatch=32': {'bytes_in': 1.0, 'bytes_out': 0.05, 'compute': 2},
'Inference\nbatch=1': {'bytes_in': 1.0, 'bytes_out': 0.002, 'compute': 0.06},
}
x = np.arange(len(scenarios))
width = 0.35
bytes_in = [s['bytes_in'] for s in scenarios.values()]
bytes_out = [s['bytes_out'] for s in scenarios.values()]
bars1 = ax.bar(x - width/2, bytes_in, width, label='Bytes In (weights)', color=COLORS['weight'])
bars2 = ax.bar(x + width/2, bytes_out, width, label='Bytes Out (output)', color=COLORS['success'])
ax.set_ylabel('Relative Data Volume', fontsize=11)
ax.set_xticks(x)
ax.set_xticklabels(scenarios.keys(), fontsize=10)
ax.legend(fontsize=10)
ax.set_title('Input vs Output Data Volume', fontsize=12, fontweight='bold')
ax.set_yscale('log')
ax.grid(True, alpha=0.3, axis='y')
# Add compute/byte ratio
for i, (name, s) in enumerate(scenarios.items()):
ax.text(i, max(bytes_in[i], bytes_out[i]) * 2,
f"{s['compute']:.1f}\nFLOP/B", ha='center', fontsize=9,
bbox=dict(facecolor='white', edgecolor='gray', boxstyle='round'))
plt.tight_layout()
plt.show()
print("\nKey insight: For LLM inference with batch=1:")
print(" - We read ~70B parameters (140GB in FP16)")
print(" - We produce ~4KB output per token (one hidden state)")
print(" - Ratio: ~35,000,000:1 input to output!")
print(" - Almost all bandwidth is consumed reading weights, not writing results")
Key insight: For LLM inference with batch=1: - We read ~70B parameters (140GB in FP16) - We produce ~4KB output per token (one hidden state) - Ratio: ~35,000,000:1 input to output! - Almost all bandwidth is consumed reading weights, not writing results
In [9]:
Copied!
# Calculate: Detailed data movement for one transformer layer
def layer_data_movement(d_model, d_ff, n_heads, seq_len, batch_size, dtype_bytes=2):
"""Calculate data movement for one transformer layer."""
d_head = d_model // n_heads
# Input activations
input_bytes = batch_size * seq_len * d_model * dtype_bytes
# === Attention ===
# QKV projection weights
qkv_weight_bytes = 3 * d_model * d_model * dtype_bytes
# QKV output (intermediate)
qkv_output_bytes = 3 * batch_size * seq_len * d_model * dtype_bytes
# Attention scores (not materialized with Flash Attention)
# Without Flash: batch * n_heads * seq_len * seq_len * dtype_bytes
# With Flash: only O(seq_len) per block
attn_scores_flash = batch_size * n_heads * seq_len * 128 * dtype_bytes # Block size 128
# Output projection weights
out_proj_weight_bytes = d_model * d_model * dtype_bytes
# === FFN ===
# Up/gate projection weights (SwiGLU has 2)
ffn_up_weight_bytes = 2 * d_model * d_ff * dtype_bytes
# Down projection weights
ffn_down_weight_bytes = d_ff * d_model * dtype_bytes
# FFN intermediate
ffn_intermediate_bytes = batch_size * seq_len * d_ff * dtype_bytes
# === Totals ===
weight_bytes = qkv_weight_bytes + out_proj_weight_bytes + ffn_up_weight_bytes + ffn_down_weight_bytes
activation_bytes = input_bytes + qkv_output_bytes + attn_scores_flash + ffn_intermediate_bytes
# FLOPs
# QKV projection: 3 × batch × seq × d_model × d_model × 2
qkv_flops = 3 * 2 * batch_size * seq_len * d_model * d_model
# Attention: batch × n_heads × (seq × seq × d_head × 2 for QK + seq × seq × d_head × 2 for AV)
attn_flops = 2 * batch_size * n_heads * 2 * seq_len * seq_len * d_head
# Out projection
out_flops = 2 * batch_size * seq_len * d_model * d_model
# FFN
ffn_flops = 2 * batch_size * seq_len * d_model * d_ff * 3 # up, gate, down
total_flops = qkv_flops + attn_flops + out_flops + ffn_flops
return {
'weight_bytes': weight_bytes,
'activation_bytes': activation_bytes,
'total_bytes': weight_bytes + activation_bytes,
'total_flops': total_flops,
'intensity': total_flops / (weight_bytes + activation_bytes),
'weight_fraction': weight_bytes / (weight_bytes + activation_bytes),
}
print("="*80)
print("DATA MOVEMENT PER TRANSFORMER LAYER (Llama 2 7B config)")
print("="*80)
# Llama 2 7B config
d_model = 4096
d_ff = 11008
n_heads = 32
configs = [
(1, 1, "Inference batch=1, 1 token"),
(1, 2048, "Inference batch=1, 2K context"),
(32, 1, "Inference batch=32, 1 token"),
(32, 2048, "Training batch=32, 2K seq"),
]
for batch, seq, name in configs:
stats = layer_data_movement(d_model, d_ff, n_heads, seq, batch)
print(f"\n{name}:")
print(f" Weight data: {stats['weight_bytes']/1e6:.1f} MB ({stats['weight_fraction']*100:.0f}% of total)")
print(f" Activation data: {stats['activation_bytes']/1e6:.1f} MB")
print(f" Total data: {stats['total_bytes']/1e6:.1f} MB")
print(f" Total FLOPs: {stats['total_flops']/1e9:.2f} GFLOPs")
print(f" Arithmetic intensity: {stats['intensity']:.1f} FLOP/byte")
# Time estimate on H100
h100_bw = 3.35e12 # bytes/s
h100_compute = 990e12 # FLOPS
time_bw_limited = stats['total_bytes'] / h100_bw * 1000 # ms
time_compute_limited = stats['total_flops'] / h100_compute * 1000 # ms
print(f" H100 time: {max(time_bw_limited, time_compute_limited):.3f} ms "
f"({'BW' if time_bw_limited > time_compute_limited else 'Compute'}-limited)")
# Calculate: Detailed data movement for one transformer layer
def layer_data_movement(d_model, d_ff, n_heads, seq_len, batch_size, dtype_bytes=2):
"""Calculate data movement for one transformer layer."""
d_head = d_model // n_heads
# Input activations
input_bytes = batch_size * seq_len * d_model * dtype_bytes
# === Attention ===
# QKV projection weights
qkv_weight_bytes = 3 * d_model * d_model * dtype_bytes
# QKV output (intermediate)
qkv_output_bytes = 3 * batch_size * seq_len * d_model * dtype_bytes
# Attention scores (not materialized with Flash Attention)
# Without Flash: batch * n_heads * seq_len * seq_len * dtype_bytes
# With Flash: only O(seq_len) per block
attn_scores_flash = batch_size * n_heads * seq_len * 128 * dtype_bytes # Block size 128
# Output projection weights
out_proj_weight_bytes = d_model * d_model * dtype_bytes
# === FFN ===
# Up/gate projection weights (SwiGLU has 2)
ffn_up_weight_bytes = 2 * d_model * d_ff * dtype_bytes
# Down projection weights
ffn_down_weight_bytes = d_ff * d_model * dtype_bytes
# FFN intermediate
ffn_intermediate_bytes = batch_size * seq_len * d_ff * dtype_bytes
# === Totals ===
weight_bytes = qkv_weight_bytes + out_proj_weight_bytes + ffn_up_weight_bytes + ffn_down_weight_bytes
activation_bytes = input_bytes + qkv_output_bytes + attn_scores_flash + ffn_intermediate_bytes
# FLOPs
# QKV projection: 3 × batch × seq × d_model × d_model × 2
qkv_flops = 3 * 2 * batch_size * seq_len * d_model * d_model
# Attention: batch × n_heads × (seq × seq × d_head × 2 for QK + seq × seq × d_head × 2 for AV)
attn_flops = 2 * batch_size * n_heads * 2 * seq_len * seq_len * d_head
# Out projection
out_flops = 2 * batch_size * seq_len * d_model * d_model
# FFN
ffn_flops = 2 * batch_size * seq_len * d_model * d_ff * 3 # up, gate, down
total_flops = qkv_flops + attn_flops + out_flops + ffn_flops
return {
'weight_bytes': weight_bytes,
'activation_bytes': activation_bytes,
'total_bytes': weight_bytes + activation_bytes,
'total_flops': total_flops,
'intensity': total_flops / (weight_bytes + activation_bytes),
'weight_fraction': weight_bytes / (weight_bytes + activation_bytes),
}
print("="*80)
print("DATA MOVEMENT PER TRANSFORMER LAYER (Llama 2 7B config)")
print("="*80)
# Llama 2 7B config
d_model = 4096
d_ff = 11008
n_heads = 32
configs = [
(1, 1, "Inference batch=1, 1 token"),
(1, 2048, "Inference batch=1, 2K context"),
(32, 1, "Inference batch=32, 1 token"),
(32, 2048, "Training batch=32, 2K seq"),
]
for batch, seq, name in configs:
stats = layer_data_movement(d_model, d_ff, n_heads, seq, batch)
print(f"\n{name}:")
print(f" Weight data: {stats['weight_bytes']/1e6:.1f} MB ({stats['weight_fraction']*100:.0f}% of total)")
print(f" Activation data: {stats['activation_bytes']/1e6:.1f} MB")
print(f" Total data: {stats['total_bytes']/1e6:.1f} MB")
print(f" Total FLOPs: {stats['total_flops']/1e9:.2f} GFLOPs")
print(f" Arithmetic intensity: {stats['intensity']:.1f} FLOP/byte")
# Time estimate on H100
h100_bw = 3.35e12 # bytes/s
h100_compute = 990e12 # FLOPS
time_bw_limited = stats['total_bytes'] / h100_bw * 1000 # ms
time_compute_limited = stats['total_flops'] / h100_compute * 1000 # ms
print(f" H100 time: {max(time_bw_limited, time_compute_limited):.3f} ms "
f"({'BW' if time_bw_limited > time_compute_limited else 'Compute'}-limited)")
================================================================================ DATA MOVEMENT PER TRANSFORMER LAYER (Llama 2 7B config) ================================================================================ Inference batch=1, 1 token: Weight data: 404.8 MB (100% of total) Activation data: 0.1 MB Total data: 404.8 MB Total FLOPs: 0.40 GFLOPs Arithmetic intensity: 1.0 FLOP/byte H100 time: 0.121 ms (BW-limited) Inference batch=1, 2K context: Weight data: 404.8 MB (76% of total) Activation data: 129.0 MB Total data: 533.7 MB Total FLOPs: 897.65 GFLOPs Arithmetic intensity: 1681.9 FLOP/byte H100 time: 0.907 ms (Compute-limited) Inference batch=32, 1 token: Weight data: 404.8 MB (100% of total) Activation data: 2.0 MB Total data: 406.8 MB Total FLOPs: 12.95 GFLOPs Arithmetic intensity: 31.8 FLOP/byte H100 time: 0.121 ms (BW-limited) Training batch=32, 2K seq: Weight data: 404.8 MB (9% of total) Activation data: 4127.2 MB Total data: 4531.9 MB Total FLOPs: 28724.74 GFLOPs Arithmetic intensity: 6338.3 FLOP/byte H100 time: 29.015 ms (Compute-limited)
7. Memory Hierarchy and Data Reuse¶
Efficient accelerators maximize data reuse to minimize expensive off-chip memory access.
Types of Data Reuse in Matrix Multiply¶
| Reuse Type | Description | Systolic Array Benefit |
|---|---|---|
| Temporal | Same data used multiple times | Accumulators stay in PE |
| Spatial-A | A elements shared across columns | A flows right through array |
| Spatial-B | B elements shared across rows | B flows down through array |
The On-Chip Memory Budget¶
For a 128×128 systolic array:
- Minimum: 2 tiles × (128×128) × 2 bytes = 64 KB (just A and B)
- Practical: Need input FIFOs, output buffers, double buffering
- TPU v4: ~32 MB on-chip SRAM
Double Buffering¶
To hide memory latency:
- While computing tile (i, j), load tile (i, j+1)
- Requires 2× tile storage
- Achieves near-100% compute utilization
In [10]:
Copied!
# Visualize: Data reuse patterns in systolic array
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
def draw_reuse_pattern(ax, pattern_type):
"""Draw data reuse pattern."""
n = 4
for i in range(n):
for j in range(n):
x, y = j, n - 1 - i
if pattern_type == 'A':
# Same row = same A element
color = plt.cm.Blues((i + 1) / (n + 1))
elif pattern_type == 'B':
# Same column = same B element
color = plt.cm.Oranges((j + 1) / (n + 1))
else:
# Accumulator - each PE unique
color = plt.cm.Greens(((i * n + j) + 1) / (n * n + 1))
rect = patches.Rectangle((x, y), 0.9, 0.9,
facecolor=color, alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
if pattern_type == 'A':
ax.text(x + 0.45, y + 0.45, f'a{i}', ha='center', va='center', fontsize=10)
elif pattern_type == 'B':
ax.text(x + 0.45, y + 0.45, f'b{j}', ha='center', va='center', fontsize=10)
else:
ax.text(x + 0.45, y + 0.45, f'c{i}{j}', ha='center', va='center', fontsize=9)
ax.set_xlim(-0.5, n + 0.5)
ax.set_ylim(-0.5, n + 0.5)
ax.set_aspect('equal')
ax.axis('off')
# A reuse (spatial across columns)
draw_reuse_pattern(axes[0], 'A')
axes[0].set_title('A Element Reuse\n(shared across row)', fontsize=12, fontweight='bold')
axes[0].annotate('', xy=(4.5, 2), xytext=(0, 2),
arrowprops=dict(arrowstyle='->', color=COLORS['primary'], lw=3))
# B reuse (spatial across rows)
draw_reuse_pattern(axes[1], 'B')
axes[1].set_title('B Element Reuse\n(shared across column)', fontsize=12, fontweight='bold')
axes[1].annotate('', xy=(2, -0.3), xytext=(2, 4.2),
arrowprops=dict(arrowstyle='->', color=COLORS['secondary'], lw=3))
# C accumulation (temporal reuse)
draw_reuse_pattern(axes[2], 'C')
axes[2].set_title('C Accumulation\n(temporal reuse in PE)', fontsize=12, fontweight='bold')
plt.suptitle('Data Reuse Patterns in Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("For P×P array computing P×K × K×P:")
print(f" Each A element reused P times (across columns)")
print(f" Each B element reused P times (across rows)")
print(f" Each C accumulator updated K times (temporal reuse)")
print(f"\nThis is why systolic arrays achieve high efficiency!")
# Visualize: Data reuse patterns in systolic array
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
def draw_reuse_pattern(ax, pattern_type):
"""Draw data reuse pattern."""
n = 4
for i in range(n):
for j in range(n):
x, y = j, n - 1 - i
if pattern_type == 'A':
# Same row = same A element
color = plt.cm.Blues((i + 1) / (n + 1))
elif pattern_type == 'B':
# Same column = same B element
color = plt.cm.Oranges((j + 1) / (n + 1))
else:
# Accumulator - each PE unique
color = plt.cm.Greens(((i * n + j) + 1) / (n * n + 1))
rect = patches.Rectangle((x, y), 0.9, 0.9,
facecolor=color, alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
if pattern_type == 'A':
ax.text(x + 0.45, y + 0.45, f'a{i}', ha='center', va='center', fontsize=10)
elif pattern_type == 'B':
ax.text(x + 0.45, y + 0.45, f'b{j}', ha='center', va='center', fontsize=10)
else:
ax.text(x + 0.45, y + 0.45, f'c{i}{j}', ha='center', va='center', fontsize=9)
ax.set_xlim(-0.5, n + 0.5)
ax.set_ylim(-0.5, n + 0.5)
ax.set_aspect('equal')
ax.axis('off')
# A reuse (spatial across columns)
draw_reuse_pattern(axes[0], 'A')
axes[0].set_title('A Element Reuse\n(shared across row)', fontsize=12, fontweight='bold')
axes[0].annotate('', xy=(4.5, 2), xytext=(0, 2),
arrowprops=dict(arrowstyle='->', color=COLORS['primary'], lw=3))
# B reuse (spatial across rows)
draw_reuse_pattern(axes[1], 'B')
axes[1].set_title('B Element Reuse\n(shared across column)', fontsize=12, fontweight='bold')
axes[1].annotate('', xy=(2, -0.3), xytext=(2, 4.2),
arrowprops=dict(arrowstyle='->', color=COLORS['secondary'], lw=3))
# C accumulation (temporal reuse)
draw_reuse_pattern(axes[2], 'C')
axes[2].set_title('C Accumulation\n(temporal reuse in PE)', fontsize=12, fontweight='bold')
plt.suptitle('Data Reuse Patterns in Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("For P×P array computing P×K × K×P:")
print(f" Each A element reused P times (across columns)")
print(f" Each B element reused P times (across rows)")
print(f" Each C accumulator updated K times (temporal reuse)")
print(f"\nThis is why systolic arrays achieve high efficiency!")
For P×P array computing P×K × K×P: Each A element reused P times (across columns) Each B element reused P times (across rows) Each C accumulator updated K times (temporal reuse) This is why systolic arrays achieve high efficiency!
In [11]:
Copied!
# Visualize: Double buffering concept
fig, ax = plt.subplots(figsize=(14, 6))
# Timeline
cycles = 12
tile_time = 3
# Memory load (green)
for i in range(4):
start = i * tile_time
rect = patches.Rectangle((start, 2), tile_time - 0.1, 0.8,
facecolor=COLORS['success'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 2.4, f'Load\nTile {i}', ha='center', va='center', fontsize=9)
# Compute (blue) - starts after first load
for i in range(4):
start = tile_time + i * tile_time
rect = patches.Rectangle((start, 1), tile_time - 0.1, 0.8,
facecolor=COLORS['primary'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 1.4, f'Compute\nTile {i}', ha='center', va='center', fontsize=9)
# Store (orange) - starts after first compute
for i in range(4):
start = 2 * tile_time + i * tile_time
if start < cycles + 3:
rect = patches.Rectangle((start, 0), tile_time - 0.1, 0.8,
facecolor=COLORS['secondary'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 0.4, f'Store\nTile {i}', ha='center', va='center', fontsize=9)
# Labels
ax.text(-1.5, 2.4, 'Memory\nLoad', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(-1.5, 1.4, 'Array\nCompute', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(-1.5, 0.4, 'Memory\nStore', ha='center', va='center', fontsize=10, fontweight='bold')
# Highlight overlap
rect = patches.Rectangle((tile_time, -0.3), 2*tile_time, 3.4,
facecolor='none', edgecolor=COLORS['danger'],
linewidth=3, linestyle='--')
ax.add_patch(rect)
ax.text(tile_time + tile_time, 3.3, 'Steady state: all units active',
ha='center', va='bottom', fontsize=10, color=COLORS['danger'])
ax.set_xlim(-2.5, 15)
ax.set_ylim(-0.5, 4)
ax.set_xlabel('Time (cycles)', fontsize=11)
ax.set_title('Double Buffering: Overlapping Memory and Compute', fontsize=14, fontweight='bold')
ax.grid(True, alpha=0.3, axis='x')
ax.set_yticks([])
plt.tight_layout()
plt.show()
print("Without double buffering: Load → Compute → Store (serial)")
print("With double buffering: Load/Compute/Store overlap (pipelined)")
print("\nRequirement: 2× buffer memory to hold current and next tile")
# Visualize: Double buffering concept
fig, ax = plt.subplots(figsize=(14, 6))
# Timeline
cycles = 12
tile_time = 3
# Memory load (green)
for i in range(4):
start = i * tile_time
rect = patches.Rectangle((start, 2), tile_time - 0.1, 0.8,
facecolor=COLORS['success'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 2.4, f'Load\nTile {i}', ha='center', va='center', fontsize=9)
# Compute (blue) - starts after first load
for i in range(4):
start = tile_time + i * tile_time
rect = patches.Rectangle((start, 1), tile_time - 0.1, 0.8,
facecolor=COLORS['primary'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 1.4, f'Compute\nTile {i}', ha='center', va='center', fontsize=9)
# Store (orange) - starts after first compute
for i in range(4):
start = 2 * tile_time + i * tile_time
if start < cycles + 3:
rect = patches.Rectangle((start, 0), tile_time - 0.1, 0.8,
facecolor=COLORS['secondary'], alpha=0.8,
edgecolor='black', linewidth=2)
ax.add_patch(rect)
ax.text(start + tile_time/2, 0.4, f'Store\nTile {i}', ha='center', va='center', fontsize=9)
# Labels
ax.text(-1.5, 2.4, 'Memory\nLoad', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(-1.5, 1.4, 'Array\nCompute', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(-1.5, 0.4, 'Memory\nStore', ha='center', va='center', fontsize=10, fontweight='bold')
# Highlight overlap
rect = patches.Rectangle((tile_time, -0.3), 2*tile_time, 3.4,
facecolor='none', edgecolor=COLORS['danger'],
linewidth=3, linestyle='--')
ax.add_patch(rect)
ax.text(tile_time + tile_time, 3.3, 'Steady state: all units active',
ha='center', va='bottom', fontsize=10, color=COLORS['danger'])
ax.set_xlim(-2.5, 15)
ax.set_ylim(-0.5, 4)
ax.set_xlabel('Time (cycles)', fontsize=11)
ax.set_title('Double Buffering: Overlapping Memory and Compute', fontsize=14, fontweight='bold')
ax.grid(True, alpha=0.3, axis='x')
ax.set_yticks([])
plt.tight_layout()
plt.show()
print("Without double buffering: Load → Compute → Store (serial)")
print("With double buffering: Load/Compute/Store overlap (pipelined)")
print("\nRequirement: 2× buffer memory to hold current and next tile")
Without double buffering: Load → Compute → Store (serial) With double buffering: Load/Compute/Store overlap (pipelined) Requirement: 2× buffer memory to hold current and next tile
8. Mapping Transformers to Systolic Arrays¶
The next step is to see how actual transformer operations map to systolic array execution.
Operation Mapping¶
| Transformer Op | Matrix Shape | Systolic Mapping | Notes |
|---|---|---|---|
| QKV Projection | (B×S, d) × (d, 3d) | GEMM | Fused as one op |
| Attention QK^T | (B×H, S, d_h) × (B×H, d_h, S) | Batched GEMM | Per head |
| Attention × V | (B×H, S, S) × (B×H, S, d_h) | Batched GEMM | Per head |
| Out Projection | (B×S, d) × (d, d) | GEMM | |
| FFN Up | (B×S, d) × (d, 4d) | GEMM | |
| FFN Down | (B×S, 4d) × (4d, d) | GEMM |
Batch-one inference¶
During autoregressive generation:
- Only processing one new token at a time
- Input is (1, d) not (B×S, d)
- GEMM becomes GEMV (matrix-vector)
This fundamentally changes the arithmetic intensity:
| Operation | Training/Prefill | Inference (B=1) |
|---|---|---|
| Shape | (B×S, d) × (d, k) | (1, d) × (d, k) |
| FLOPs | 2 × B × S × d × k | 2 × d × k |
| Bytes | ≈ d × k × 2 (weights) | d × k × 2 (weights) |
| Intensity | ~B×S (high) | ~1 (very low) |
In [12]:
Copied!
# Visualize: Mapping transformer operations to systolic array
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
def draw_gemm_on_array(ax, M, N, K, array_size, title):
"""Draw how a GEMM maps to systolic array."""
P = array_size
# Calculate tiles
m_tiles = (M + P - 1) // P
n_tiles = (N + P - 1) // P
k_tiles = (K + P - 1) // P
# Draw array utilization grid
max_dim = max(m_tiles, n_tiles)
cell_size = 4.0 / max(max_dim, 1)
for i in range(m_tiles):
for j in range(n_tiles):
x = j * cell_size
y = (m_tiles - 1 - i) * cell_size
rect = patches.Rectangle((x, y), cell_size * 0.95, cell_size * 0.95,
facecolor=COLORS['pe'], alpha=0.8,
edgecolor='black', linewidth=1)
ax.add_patch(rect)
# Stats
total_tiles = m_tiles * n_tiles * k_tiles
# Handle edge case where dimensions are smaller than array
if M < P or N < P:
utilization = (M * N) / (P * P) * 100
else:
utilization = 100 # Full tiles
ax.set_xlim(-0.5, 5)
ax.set_ylim(-0.5, 5)
ax.set_aspect('equal')
ax.axis('off')
ax.set_title(title, fontsize=11, fontweight='bold')
# Add info text
info = f"Shape: [{M}×{K}] × [{K}×{N}]\n"
info += f"Tiles: {m_tiles}×{n_tiles}×{k_tiles} = {total_tiles}\n"
info += f"Utilization: {utilization:.0f}%"
ax.text(4.5, 0.5, info, fontsize=9, va='bottom', ha='left',
bbox=dict(facecolor='white', edgecolor='gray', boxstyle='round'))
# Llama 2 7B ops with 128×128 array
P = 128
# QKV projection (training, batch=32, seq=2048)
draw_gemm_on_array(axes[0, 0], 32*2048, 12288, 4096, P,
'QKV Projection (Training)\nbatch=32, seq=2048')
# QKV projection (inference, batch=1, 1 token)
draw_gemm_on_array(axes[0, 1], 1, 12288, 4096, P,
'QKV Projection (Inference)\nbatch=1, 1 token')
# FFN (training)
draw_gemm_on_array(axes[1, 0], 32*2048, 11008, 4096, P,
'FFN Up (Training)\nbatch=32, seq=2048')
# FFN (inference)
draw_gemm_on_array(axes[1, 1], 1, 11008, 4096, P,
'FFN Up (Inference)\nbatch=1, 1 token')
plt.suptitle('Mapping Transformer Ops to 128×128 Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("\nCritical observation:")
print(" Training: Large batch fills the array efficiently")
print(" Inference: Only 1 row of the array is used!")
print(" → 128×128 = 16,384 PEs, but only 128 active = 0.8% utilization")
print("\nThis is why inference is memory-bound, not compute-bound.")
# Visualize: Mapping transformer operations to systolic array
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
def draw_gemm_on_array(ax, M, N, K, array_size, title):
"""Draw how a GEMM maps to systolic array."""
P = array_size
# Calculate tiles
m_tiles = (M + P - 1) // P
n_tiles = (N + P - 1) // P
k_tiles = (K + P - 1) // P
# Draw array utilization grid
max_dim = max(m_tiles, n_tiles)
cell_size = 4.0 / max(max_dim, 1)
for i in range(m_tiles):
for j in range(n_tiles):
x = j * cell_size
y = (m_tiles - 1 - i) * cell_size
rect = patches.Rectangle((x, y), cell_size * 0.95, cell_size * 0.95,
facecolor=COLORS['pe'], alpha=0.8,
edgecolor='black', linewidth=1)
ax.add_patch(rect)
# Stats
total_tiles = m_tiles * n_tiles * k_tiles
# Handle edge case where dimensions are smaller than array
if M < P or N < P:
utilization = (M * N) / (P * P) * 100
else:
utilization = 100 # Full tiles
ax.set_xlim(-0.5, 5)
ax.set_ylim(-0.5, 5)
ax.set_aspect('equal')
ax.axis('off')
ax.set_title(title, fontsize=11, fontweight='bold')
# Add info text
info = f"Shape: [{M}×{K}] × [{K}×{N}]\n"
info += f"Tiles: {m_tiles}×{n_tiles}×{k_tiles} = {total_tiles}\n"
info += f"Utilization: {utilization:.0f}%"
ax.text(4.5, 0.5, info, fontsize=9, va='bottom', ha='left',
bbox=dict(facecolor='white', edgecolor='gray', boxstyle='round'))
# Llama 2 7B ops with 128×128 array
P = 128
# QKV projection (training, batch=32, seq=2048)
draw_gemm_on_array(axes[0, 0], 32*2048, 12288, 4096, P,
'QKV Projection (Training)\nbatch=32, seq=2048')
# QKV projection (inference, batch=1, 1 token)
draw_gemm_on_array(axes[0, 1], 1, 12288, 4096, P,
'QKV Projection (Inference)\nbatch=1, 1 token')
# FFN (training)
draw_gemm_on_array(axes[1, 0], 32*2048, 11008, 4096, P,
'FFN Up (Training)\nbatch=32, seq=2048')
# FFN (inference)
draw_gemm_on_array(axes[1, 1], 1, 11008, 4096, P,
'FFN Up (Inference)\nbatch=1, 1 token')
plt.suptitle('Mapping Transformer Ops to 128×128 Systolic Array', fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.show()
print("\nCritical observation:")
print(" Training: Large batch fills the array efficiently")
print(" Inference: Only 1 row of the array is used!")
print(" → 128×128 = 16,384 PEs, but only 128 active = 0.8% utilization")
print("\nThis is why inference is memory-bound, not compute-bound.")
Critical observation: Training: Large batch fills the array efficiently Inference: Only 1 row of the array is used! → 128×128 = 16,384 PEs, but only 128 active = 0.8% utilization This is why inference is memory-bound, not compute-bound.
In [13]:
Copied!
# Calculate: Complete inference pass data movement
def full_inference_analysis(d_model, n_layers, d_ff, n_heads, vocab_size,
seq_len, dtype_bytes=2):
"""Analyze complete inference pass for one token."""
d_head = d_model // n_heads
# === Weights (read from memory) ===
# Per-layer weights
qkv_weights = 3 * d_model * d_model * dtype_bytes
out_proj_weights = d_model * d_model * dtype_bytes
ffn_up_weights = 2 * d_model * d_ff * dtype_bytes # gate + up
ffn_down_weights = d_ff * d_model * dtype_bytes
layer_weights = qkv_weights + out_proj_weights + ffn_up_weights + ffn_down_weights
total_layer_weights = layer_weights * n_layers
# Embedding + output projection
embed_weights = vocab_size * d_model * dtype_bytes
output_weights = d_model * vocab_size * dtype_bytes # Often tied
total_weights = total_layer_weights + embed_weights
# === KV Cache (read existing, write new) ===
# For each layer, read seq_len KV pairs, write 1 new pair
kv_read = n_layers * 2 * seq_len * d_model * dtype_bytes
kv_write = n_layers * 2 * 1 * d_model * dtype_bytes
# === Activations (mostly on-chip) ===
# Input embedding: 1 × d_model
# Between layers: 1 × d_model
activation_per_layer = d_model * dtype_bytes # Very small!
# === FLOPs ===
# Per layer:
qkv_flops = 3 * 2 * 1 * d_model * d_model
attn_qk_flops = 2 * n_heads * 1 * seq_len * d_head # New Q against all K
attn_v_flops = 2 * n_heads * 1 * seq_len * d_head # Weights × V
out_flops = 2 * 1 * d_model * d_model
ffn_flops = 3 * 2 * 1 * d_model * d_ff # gate + up + down
layer_flops = qkv_flops + attn_qk_flops + attn_v_flops + out_flops + ffn_flops
total_flops = layer_flops * n_layers + 2 * d_model * vocab_size # + output projection
# === Summary ===
total_read = total_weights + kv_read
total_write = kv_write + d_model * dtype_bytes # output hidden state
return {
'total_weights_gb': total_weights / 1e9,
'kv_read_gb': kv_read / 1e9,
'kv_write_bytes': kv_write,
'total_read_gb': total_read / 1e9,
'total_write_mb': total_write / 1e6,
'total_flops_gflops': total_flops / 1e9,
'arithmetic_intensity': total_flops / (total_read + total_write),
'read_write_ratio': total_read / total_write,
}
print("="*80)
print("COMPLETE INFERENCE ANALYSIS: One Token Generation")
print("="*80)
# Llama 2 7B
config_7b = {
'd_model': 4096, 'n_layers': 32, 'd_ff': 11008,
'n_heads': 32, 'vocab_size': 32000
}
for seq_len in [1, 512, 2048, 8192]:
stats = full_inference_analysis(**config_7b, seq_len=seq_len)
print(f"\nLlama 2 7B, context length = {seq_len}:")
print(f" Weights read: {stats['total_weights_gb']:.2f} GB")
print(f" KV cache read: {stats['kv_read_gb']:.4f} GB ({stats['kv_read_gb']/stats['total_read_gb']*100:.1f}% of reads)")
print(f" Total read: {stats['total_read_gb']:.2f} GB")
print(f" Total write: {stats['total_write_mb']:.3f} MB")
print(f" Read:Write ratio: {stats['read_write_ratio']:.0f}:1")
print(f" FLOPs: {stats['total_flops_gflops']:.1f} GFLOPs")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.2f} FLOP/byte")
# Time estimate
h100_bw = 3.35e12 # bytes/s
time_ms = stats['total_read_gb'] * 1e9 / h100_bw * 1000
tokens_per_sec = 1000 / time_ms
print(f" H100 time: {time_ms:.2f} ms → {tokens_per_sec:.0f} tokens/s")
# Calculate: Complete inference pass data movement
def full_inference_analysis(d_model, n_layers, d_ff, n_heads, vocab_size,
seq_len, dtype_bytes=2):
"""Analyze complete inference pass for one token."""
d_head = d_model // n_heads
# === Weights (read from memory) ===
# Per-layer weights
qkv_weights = 3 * d_model * d_model * dtype_bytes
out_proj_weights = d_model * d_model * dtype_bytes
ffn_up_weights = 2 * d_model * d_ff * dtype_bytes # gate + up
ffn_down_weights = d_ff * d_model * dtype_bytes
layer_weights = qkv_weights + out_proj_weights + ffn_up_weights + ffn_down_weights
total_layer_weights = layer_weights * n_layers
# Embedding + output projection
embed_weights = vocab_size * d_model * dtype_bytes
output_weights = d_model * vocab_size * dtype_bytes # Often tied
total_weights = total_layer_weights + embed_weights
# === KV Cache (read existing, write new) ===
# For each layer, read seq_len KV pairs, write 1 new pair
kv_read = n_layers * 2 * seq_len * d_model * dtype_bytes
kv_write = n_layers * 2 * 1 * d_model * dtype_bytes
# === Activations (mostly on-chip) ===
# Input embedding: 1 × d_model
# Between layers: 1 × d_model
activation_per_layer = d_model * dtype_bytes # Very small!
# === FLOPs ===
# Per layer:
qkv_flops = 3 * 2 * 1 * d_model * d_model
attn_qk_flops = 2 * n_heads * 1 * seq_len * d_head # New Q against all K
attn_v_flops = 2 * n_heads * 1 * seq_len * d_head # Weights × V
out_flops = 2 * 1 * d_model * d_model
ffn_flops = 3 * 2 * 1 * d_model * d_ff # gate + up + down
layer_flops = qkv_flops + attn_qk_flops + attn_v_flops + out_flops + ffn_flops
total_flops = layer_flops * n_layers + 2 * d_model * vocab_size # + output projection
# === Summary ===
total_read = total_weights + kv_read
total_write = kv_write + d_model * dtype_bytes # output hidden state
return {
'total_weights_gb': total_weights / 1e9,
'kv_read_gb': kv_read / 1e9,
'kv_write_bytes': kv_write,
'total_read_gb': total_read / 1e9,
'total_write_mb': total_write / 1e6,
'total_flops_gflops': total_flops / 1e9,
'arithmetic_intensity': total_flops / (total_read + total_write),
'read_write_ratio': total_read / total_write,
}
print("="*80)
print("COMPLETE INFERENCE ANALYSIS: One Token Generation")
print("="*80)
# Llama 2 7B
config_7b = {
'd_model': 4096, 'n_layers': 32, 'd_ff': 11008,
'n_heads': 32, 'vocab_size': 32000
}
for seq_len in [1, 512, 2048, 8192]:
stats = full_inference_analysis(**config_7b, seq_len=seq_len)
print(f"\nLlama 2 7B, context length = {seq_len}:")
print(f" Weights read: {stats['total_weights_gb']:.2f} GB")
print(f" KV cache read: {stats['kv_read_gb']:.4f} GB ({stats['kv_read_gb']/stats['total_read_gb']*100:.1f}% of reads)")
print(f" Total read: {stats['total_read_gb']:.2f} GB")
print(f" Total write: {stats['total_write_mb']:.3f} MB")
print(f" Read:Write ratio: {stats['read_write_ratio']:.0f}:1")
print(f" FLOPs: {stats['total_flops_gflops']:.1f} GFLOPs")
print(f" Arithmetic intensity: {stats['arithmetic_intensity']:.2f} FLOP/byte")
# Time estimate
h100_bw = 3.35e12 # bytes/s
time_ms = stats['total_read_gb'] * 1e9 / h100_bw * 1000
tokens_per_sec = 1000 / time_ms
print(f" H100 time: {time_ms:.2f} ms → {tokens_per_sec:.0f} tokens/s")
================================================================================ COMPLETE INFERENCE ANALYSIS: One Token Generation ================================================================================ Llama 2 7B, context length = 1: Weights read: 13.21 GB KV cache read: 0.0005 GB (0.0% of reads) Total read: 13.21 GB Total write: 0.532 MB Read:Write ratio: 24817:1 FLOPs: 13.2 GFLOPs Arithmetic intensity: 1.00 FLOP/byte H100 time: 3.94 ms → 254 tokens/s Llama 2 7B, context length = 512: Weights read: 13.21 GB KV cache read: 0.2684 GB (2.0% of reads) Total read: 13.48 GB Total write: 0.532 MB Read:Write ratio: 25320:1 FLOPs: 13.5 GFLOPs Arithmetic intensity: 1.00 FLOP/byte H100 time: 4.02 ms → 248 tokens/s Llama 2 7B, context length = 2048: Weights read: 13.21 GB KV cache read: 1.0737 GB (7.5% of reads) Total read: 14.29 GB Total write: 0.532 MB Read:Write ratio: 26833:1 FLOPs: 14.3 GFLOPs Arithmetic intensity: 1.00 FLOP/byte H100 time: 4.27 ms → 234 tokens/s Llama 2 7B, context length = 8192: Weights read: 13.21 GB KV cache read: 4.2950 GB (24.5% of reads) Total read: 17.51 GB Total write: 0.532 MB Read:Write ratio: 32882:1 FLOPs: 17.5 GFLOPs Arithmetic intensity: 1.00 FLOP/byte H100 time: 5.23 ms → 191 tokens/s
9. Why Superconducting Logic?¶
Reaching the compute-bound regime¶
Batching and reuse can move transformer inference from a bandwidth-limited regime toward a compute-limited regime. Exact crossover points depend on model dimensions, precision, attention implementation, memory system, and delivered accelerator utilization; batch-size thresholds are therefore workload results rather than universal constants.
Both CMOS and SCE benefit from the same reuse. The technology comparison begins only after workload, numerical format, utilization, memory traffic, and system boundary have been matched.
Device and system evidence¶
| Level | Result | Boundary |
|---|---|---|
| AQFP device | 1.4 zJ/JJ-cycle | Published 8-bit-adder measurement at 4.2 K and 5 GHz |
| Thermodynamic floor | 0.040 zJ/bit erase | Landauer value at 4.2 K |
| Production system model | 11.8–14.8× less wall energy | 5 K, 400 W/W, full routed clock fabric, warm memory and I/O, host-inclusive B200 FP4 comparator |
| Packaged downside model | 6.4–8.5× less wall energy | 5 K, 980 W/W, packaged-machine class, same accounting boundary |
| Plant upside model | 15.7–19.5× less wall energy | 5 K, 207 W/W, central-helium-plant class, same accounting boundary |
The device result provides headroom. A systolic, weight-stationary architecture attempts to preserve it through local reuse, deterministic pipelines, and bounded traffic across the warm-cold boundary. The complete result remains model-derived until integrated hardware runs the workload.
Sources: Takeuchi et al. (2019); current AM static-inclusive system model, as of 2026-08-13.
In [14]:
Copied!
# Visualize: The path from conventional to superconducting accelerators
fig, ax = plt.subplots(figsize=(14, 8))
# Current state
current = patches.FancyBboxPatch((0.5, 5.5), 4, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['danger'], alpha=0.3,
edgecolor='black', linewidth=2)
ax.add_patch(current)
ax.text(2.5, 6.5, 'Current State:\nGPU/TPU + HBM', ha='center', va='center', fontsize=12, fontweight='bold')
ax.text(2.5, 5.9, '• Memory-bound inference\n• Workload-dependent utilization\n• High-power accelerator package',
ha='center', va='center', fontsize=9)
# Near-term improvements
near = patches.FancyBboxPatch((6, 5.5), 4, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['secondary'], alpha=0.3,
edgecolor='black', linewidth=2)
ax.add_patch(near)
ax.text(8, 6.5, 'Near-Term:\nMore HBM + Quantization', ha='center', va='center', fontsize=12, fontweight='bold')
ax.text(8, 5.9, '• HBM3E-class memory\n• INT4/INT8 models\n• Workload-dependent gains',
ha='center', va='center', fontsize=9)
# Superconducting vision
sce = patches.FancyBboxPatch((3, 1.5), 5, 2.5,
boxstyle="round,pad=0.05",
facecolor=COLORS['success'], alpha=0.3,
edgecolor='black', linewidth=3)
ax.add_patch(sce)
ax.text(5.5, 3.2, 'Superconducting Vision:\nCompute-Bound Operation', ha='center', va='center',
fontsize=12, fontweight='bold')
ax.text(5.5, 2.3, '• Target compute-bound workloads:\n prefill, large batches, CNNs, diffusion\n• Internal bandwidth sized from the workload\n• 5 GHz production clock target\n• System power is model-dependent',
ha='center', va='center', fontsize=10)
# Arrows
ax.annotate('', xy=(6, 6.5), xytext=(4.5, 6.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2))
ax.annotate('', xy=(5.5, 4), xytext=(4, 5.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2, connectionstyle='arc3,rad=-0.2'))
ax.annotate('', xy=(5.5, 4), xytext=(7, 5.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2, connectionstyle='arc3,rad=0.2'))
# Challenge callout
challenge = patches.FancyBboxPatch((9.5, 2), 3, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['light'], alpha=0.9,
edgecolor='gray', linewidth=2)
ax.add_patch(challenge)
ax.text(11, 3, 'Key Challenges:', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(11, 2.4, '• 4–5 K operating regime\n• Fabrication and test\n• Warm-cold I/O\n• Software stack',
ha='center', va='center', fontsize=9)
ax.set_xlim(0, 13)
ax.set_ylim(1, 8)
ax.axis('off')
ax.set_title('Evolution Path: From Memory-Bound to Compute-Bound AI', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()
print("The fundamental insight:")
print(" SCE targets workloads that are already compute-bound:")
print(" prefill, large-batch inference, CNNs, diffusion models.")
print(" The model evaluates clock speed and wall-energy efficiency,")
print(" not trying to put 100 GB of memory on chip.")
print('\nAt the cold stage, architecture must integrate:')
print(" • High-speed logic (5 GHz production target)")
print(" • Published device reference: 1.4 zJ/JJ-cycle; system energy requires the full model")
print(" • Internal bandwidth sized from the workload and physical hierarchy")
# Visualize: The path from conventional to superconducting accelerators
fig, ax = plt.subplots(figsize=(14, 8))
# Current state
current = patches.FancyBboxPatch((0.5, 5.5), 4, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['danger'], alpha=0.3,
edgecolor='black', linewidth=2)
ax.add_patch(current)
ax.text(2.5, 6.5, 'Current State:\nGPU/TPU + HBM', ha='center', va='center', fontsize=12, fontweight='bold')
ax.text(2.5, 5.9, '• Memory-bound inference\n• Workload-dependent utilization\n• High-power accelerator package',
ha='center', va='center', fontsize=9)
# Near-term improvements
near = patches.FancyBboxPatch((6, 5.5), 4, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['secondary'], alpha=0.3,
edgecolor='black', linewidth=2)
ax.add_patch(near)
ax.text(8, 6.5, 'Near-Term:\nMore HBM + Quantization', ha='center', va='center', fontsize=12, fontweight='bold')
ax.text(8, 5.9, '• HBM3E-class memory\n• INT4/INT8 models\n• Workload-dependent gains',
ha='center', va='center', fontsize=9)
# Superconducting vision
sce = patches.FancyBboxPatch((3, 1.5), 5, 2.5,
boxstyle="round,pad=0.05",
facecolor=COLORS['success'], alpha=0.3,
edgecolor='black', linewidth=3)
ax.add_patch(sce)
ax.text(5.5, 3.2, 'Superconducting Vision:\nCompute-Bound Operation', ha='center', va='center',
fontsize=12, fontweight='bold')
ax.text(5.5, 2.3, '• Target compute-bound workloads:\n prefill, large batches, CNNs, diffusion\n• Internal bandwidth sized from the workload\n• 5 GHz production clock target\n• System power is model-dependent',
ha='center', va='center', fontsize=10)
# Arrows
ax.annotate('', xy=(6, 6.5), xytext=(4.5, 6.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2))
ax.annotate('', xy=(5.5, 4), xytext=(4, 5.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2, connectionstyle='arc3,rad=-0.2'))
ax.annotate('', xy=(5.5, 4), xytext=(7, 5.5),
arrowprops=dict(arrowstyle='->', color='black', lw=2, connectionstyle='arc3,rad=0.2'))
# Challenge callout
challenge = patches.FancyBboxPatch((9.5, 2), 3, 2,
boxstyle="round,pad=0.05",
facecolor=COLORS['light'], alpha=0.9,
edgecolor='gray', linewidth=2)
ax.add_patch(challenge)
ax.text(11, 3, 'Key Challenges:', ha='center', va='center', fontsize=10, fontweight='bold')
ax.text(11, 2.4, '• 4–5 K operating regime\n• Fabrication and test\n• Warm-cold I/O\n• Software stack',
ha='center', va='center', fontsize=9)
ax.set_xlim(0, 13)
ax.set_ylim(1, 8)
ax.axis('off')
ax.set_title('Evolution Path: From Memory-Bound to Compute-Bound AI', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()
print("The fundamental insight:")
print(" SCE targets workloads that are already compute-bound:")
print(" prefill, large-batch inference, CNNs, diffusion models.")
print(" The model evaluates clock speed and wall-energy efficiency,")
print(" not trying to put 100 GB of memory on chip.")
print('\nAt the cold stage, architecture must integrate:')
print(" • High-speed logic (5 GHz production target)")
print(" • Published device reference: 1.4 zJ/JJ-cycle; system energy requires the full model")
print(" • Internal bandwidth sized from the workload and physical hierarchy")
The fundamental insight: SCE targets workloads that are already compute-bound: prefill, large-batch inference, CNNs, diffusion models. The model evaluates clock speed and wall-energy efficiency, not trying to put 100 GB of memory on chip. At the cold stage, architecture must integrate: • High-speed logic (5 GHz production target) • Published device reference: 1.4 zJ/JJ-cycle; system energy requires the full model • Internal bandwidth sized from the workload and physical hierarchy
In [15]:
Copied!
# Visualize: Current modeled system corners
labels = ['Packaged downside\n5 K / 980 W/W', 'Production\n5 K / 400 W/W', 'Plant upside\n5 K / 207 W/W']
low = np.array([6.4, 11.8, 15.7])
high = np.array([8.5, 14.8, 19.5])
mid = (low + high) / 2
err = np.vstack([mid - low, high - mid])
fig, ax = plt.subplots(figsize=(11, 6))
bars = ax.bar(labels, mid, color=[COLORS['secondary'], COLORS['primary'], COLORS['success']], alpha=0.85)
ax.errorbar(np.arange(len(labels)), mid, yerr=err, fmt='none', color='black', capsize=6, linewidth=1.5)
ax.axhline(1, color=COLORS['dark'], linestyle='--', linewidth=1, alpha=0.6)
ax.set_ylabel('Modeled wall-energy advantage (×)', fontsize=11)
ax.set_title('Static-Inclusive System Model vs Host-Inclusive NVIDIA B200, Dense FP4', fontsize=13, fontweight='bold')
ax.grid(True, alpha=0.25, axis='y')
ax.set_ylim(0, 22)
for bar, lo, hi in zip(bars, low, high):
ax.text(bar.get_x() + bar.get_width()/2, hi + 0.6, f'{lo:.1f}–{hi:.1f}×', ha='center', fontsize=10)
plt.tight_layout()
plt.show()
print('Current modeled wall-energy ranges vs host-inclusive NVIDIA B200, dense FP4:')
print(' • Packaged downside: 6.4–8.5× at 5 K / 980 W/W')
print(' • Production: 11.8–14.8× at 5 K / 400 W/W')
print(' • Plant upside: 15.7–19.5× at 5 K / 207 W/W')
print('These are modeled system results, not measurements of an integrated accelerator.')
# Visualize: Current modeled system corners
labels = ['Packaged downside\n5 K / 980 W/W', 'Production\n5 K / 400 W/W', 'Plant upside\n5 K / 207 W/W']
low = np.array([6.4, 11.8, 15.7])
high = np.array([8.5, 14.8, 19.5])
mid = (low + high) / 2
err = np.vstack([mid - low, high - mid])
fig, ax = plt.subplots(figsize=(11, 6))
bars = ax.bar(labels, mid, color=[COLORS['secondary'], COLORS['primary'], COLORS['success']], alpha=0.85)
ax.errorbar(np.arange(len(labels)), mid, yerr=err, fmt='none', color='black', capsize=6, linewidth=1.5)
ax.axhline(1, color=COLORS['dark'], linestyle='--', linewidth=1, alpha=0.6)
ax.set_ylabel('Modeled wall-energy advantage (×)', fontsize=11)
ax.set_title('Static-Inclusive System Model vs Host-Inclusive NVIDIA B200, Dense FP4', fontsize=13, fontweight='bold')
ax.grid(True, alpha=0.25, axis='y')
ax.set_ylim(0, 22)
for bar, lo, hi in zip(bars, low, high):
ax.text(bar.get_x() + bar.get_width()/2, hi + 0.6, f'{lo:.1f}–{hi:.1f}×', ha='center', fontsize=10)
plt.tight_layout()
plt.show()
print('Current modeled wall-energy ranges vs host-inclusive NVIDIA B200, dense FP4:')
print(' • Packaged downside: 6.4–8.5× at 5 K / 980 W/W')
print(' • Production: 11.8–14.8× at 5 K / 400 W/W')
print(' • Plant upside: 15.7–19.5× at 5 K / 207 W/W')
print('These are modeled system results, not measurements of an integrated accelerator.')
Current modeled wall-energy ranges vs host-inclusive NVIDIA B200, dense FP4: • Packaged downside: 6.4–8.5× at 5 K / 980 W/W • Production: 11.8–14.8× at 5 K / 400 W/W • Plant upside: 15.7–19.5× at 5 K / 207 W/W These are modeled system results, not measurements of an integrated accelerator.
Summary¶
Key Concepts¶
Systolic Arrays:
- Regular grid of processing elements performing MAC operations
- Data flows through array, maximizing reuse
- Each element of A reused P times, each element of B reused P times
- Arithmetic intensity = P / bytes_per_element
Tiling Large Matrices:
- Split M×K × K×N into tiles that fit the P×P array
- Accumulate partial results across K-dimension tiles
- Double buffering hides memory latency
Data Movement Analysis:
- Inference (batch=1): Read entire model, produce tiny output
- Read:Write ratio can exceed 10,000,000:1
- Memory bandwidth is the bottleneck, not compute
The Superconducting Opportunity:
- Large batch sizes increase arithmetic intensity
- Registers over SRAM: every AQFP gate is a flip-flop
- Transform memory-bound → compute-bound
- Use a 5 GHz production target while measuring the clock, I/O, memory, static, and refrigerator terms separately
Connection to Course¶
This lecture explains why superconducting electronics matter for AI:
- The memory wall is a fundamental physics problem
- Conventional scaling (Moore's Law) can't solve it
- The current static-inclusive production model uses 11.8–14.8× less wall energy than its host-inclusive B200 FP4 comparator
The previous lectures (AQFP, majority gates, fabrication) show how to build it.
Key Numbers to Remember¶
| Metric | Value |
|---|---|
| TPU v4 array size | 128 × 128 = 16,384 MACs |
| Inference arithmetic intensity (batch=1) | ~2 FLOP/byte |
| H100 ridge point | ~300 FLOP/byte |
| H100 inference utilization | ~1% |
| Read:Write ratio (LLM inference) | >1,000,000:1 |
| Llama 70B tokens/s (H100, batch=1) | ~24 |
Check Your Understanding¶
- For a $P\times P$ array, derive the cycle count for one full tile including fill and drain latency.
- Compare activation-, weight-, and output-stationary mappings for batch-one inference and batched inference.
- Why does the measured 1.4 zJ/JJ-cycle reference not by itself determine wall energy per transformer token? List the additional model terms.
Lecture 14 complete
Architecture converts device-level headroom into useful work by keeping data local and hardware occupied.
Architecture converts device-level headroom into useful work by keeping data local and hardware occupied.
[← Previous: 13. AI Workloads and Dataflow](../13_ai_workloads_dataflow/) | [Course home](../../)