Architectures
Hardware-Aware LLM Design β Bridging Theory and Silicon
Model architecture and hardware are inseparable. Understanding GPU memory hierarchy, tensor cores, and kernel optimization enables designs that are 10-100Γ faster in practice, even with equivalent theoretical complexity.
- Memory Hierarchy β Registers, shared memory, L1/L2 cache, HBM
- Tensor Cores β Matrix multiplication units optimized for specific data types
- Kernel Fusion β Combining operations to minimize memory transfers
- Architecture Design β Shaping models to match hardware capabilities
The fastest algorithm is the one that matches the hardware you have.
Hardware-Aware LLM Design
The theoretical complexity of an algorithm tells only part of the story. Real-world performance depends on how well the computation maps to the underlying hardware. For LLMs, this means understanding GPU architecture and designing models that exploit it.
GPU Memory Hierarchy
The Memory Pyramid
Memory Bandwidth
Arithmetic Intensity
Tensor Cores and Matrix Operations
Tensor Core Architecture
Tensor Core Operations
Data Type Considerations
| Data Type | Size | Tensor Core Support | Throughput (A100) |
|---|---|---|---|
| FP32 | 4 bytes | No | 19.5 TFLOPS |
| TF32 | 4 bytes | Yes | 156 TFLOPS |
| FP16 | 2 bytes | Yes | 312 TFLOPS |
| BF16 | 2 bytes | Yes | 312 TFLOPS |
| INT8 | 1 byte | Yes | 624 TOPS |
| INT4 | 0.5 bytes | Yes | 1,248 TOPS |
Kernel Fusion
The Memory Transfer Problem
# WITHOUT fusion: 3 memory round-trips
def unfused_attention(q, k, v):
scores = torch.matmul(q, k.transpose(-2, -1)) # Load q, k; store scores
weights = F.softmax(scores, dim=-1) # Load scores; store weights
output = torch.matmul(weights, v) # Load weights, v; store output
return output
# WITH fusion: 1 memory round-trip (flash attention)
def fused_attention(q, k, v, block_size=256):
"""Fused attention using tiling."""
B, H, L, D = q.shape
output = torch.zeros_like(q)
lse = torch.full((B, H, L), float('-inf'), device=q.device)
# Process in blocks
for i in range(0, L, block_size):
q_block = q[:, :, i:i+block_size]
# Load k, v once for this q block
for j in range(0, L, block_size):
k_block = k[:, :, j:j+block_size]
v_block = v[:, :, j:j+block_size]
# Compute block attention (all in registers/shared memory)
scores = torch.matmul(q_block, k_block.transpose(-2, -1))
# Online softmax update
block_max = scores.max(dim=-1, keepdim=True).values
scores = scores - block_max
# Update output and normalization
exp_scores = torch.exp(scores)
output[:, :, i:i+block_size] += torch.matmul(exp_scores, v_block)
lse[:, :, i:i+block_size] = torch.logaddexp(
lse[:, :, i:i+block_size],
block_max.squeeze(-1)
)
# Final normalization
output = output / lse.unsqueeze(-1).exp().unsqueeze(-1)
return output
Flash Attention
Architecture Design for Hardware
Optimal Hidden Dimensions
Layer Normalization Placement
class PreNormBlock(nn.Module):
"""Pre-norm transformer block (hardware efficient)."""
def __init__(self, d_model, n_heads, d_ffn):
super().__init__()
self.norm1 = nn.LayerNorm(d_model)
self.attn = MultiHeadAttention(d_model, n_heads)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model, d_ffn)
def forward(self, x):
# Pre-norm: more stable training, better hardware utilization
x = x + self.attn(self.norm1(x))
x = x + self.ffn(self.norm2(x))
return x
Activation Function Selection
| Activation | Compute | Memory | Hardware Efficiency |
|---|---|---|---|
| ReLU | Minimal | Low | Excellent |
| GELU | Moderate | Low | Good |
| SiLU/Swish | Moderate | Low | Good |
| GeGLU | Higher | Higher | Moderate |
Quantization and Hardware
Hardware-Specific Quantization
class HardwareAwareQuantization:
"""Quantization strategies for different hardware."""
@staticmethod
def get_optimal_quantization(hardware_type):
if hardware_type == "A100":
return {
"weight": "INT8", # Tensor core INT8 support
"activation": "FP16", # Keep activations in FP16
"kv_cache": "INT8", # Save memory on KV cache
}
elif hardware_type == "RTX_4090":
return {
"weight": "INT4", # Maximize memory savings
"activation": "FP16", # FP16 activations
"kv_cache": "FP16", # Limited INT4 support
}
elif hardware_type == "CPU":
return {
"weight": "INT4", # GGUF format
"activation": "FP32", # CPU prefers FP32
"kv_cache": "FP32", # No special support
}
Tensor Core Utilization
Memory Optimization Techniques
Weight Sharing
Gradient Checkpointing
class CheckpointedTransformer(nn.Module):
"""Transformer with gradient checkpointing."""
def __init__(self, d_model, n_layers, n_heads):
super().__init__()
self.layers = nn.ModuleList([
TransformerBlock(d_model, n_heads)
for _ in range(n_layers)
])
def forward(self, x):
for layer in self.layers:
# Checkpoint: recompute forward during backward
x = torch.utils.checkpoint.checkpoint(
layer, x, use_reentrant=False
)
return x
Benchmarking Hardware Efficiency
Measuring Real Performance
def benchmark_model(model, input_ids, n_warmup=10, n_iter=100):
"""Benchmark model inference performance."""
import time
# Warmup
for _ in range(n_warmup):
with torch.no_grad():
model(input_ids)
# Benchmark
torch.cuda.synchronize()
start = time.time()
for _ in range(n_iter):
with torch.no_grad():
model(input_ids)
torch.cuda.synchronize()
end = time.time()
# Calculate metrics
tokens_per_second = input_ids.shape[1] * n_iter / (end - start)
memory_used = torch.cuda.max_memory_allocated() / 1024**3
return {
"tokens_per_second": tokens_per_second,
"memory_gb": memory_used,
"latency_ms": (end - start) / n_iter * 1000
}
Practice Exercises
-
Conceptual: Explain why autoregressive generation is memory-bandwidth bound while prefill is compute-bound. How does this affect optimization strategies?
-
Mathematical: Calculate the arithmetic intensity of matrix multiplication for two 4096Γ4096 matrices in FP16. Is this operation compute-bound or memory-bound on an A100?
-
Practical: Benchmark the same transformer model with different hidden dimensions (2048, 4096, 8192) and measure how tensor core alignment affects throughput.
-
Research: Investigate how mixed-precision training (FP16/BF16) affects both training speed and final model quality. What is the optimal precision strategy?
What to Learn Next
-> Flash Attention and Memory Efficiency IO-aware attention optimization for modern GPUs.
-> Quantization Techniques Deep Dive GPTQ, AWQ, GGUF, and hardware-specific quantization.
-> Model Parallelism and Tensor Parallelism Distributing models across multiple GPUs.
-> KV Cache Optimization Optimizing transformer inference memory.
-> LLM Inference Optimization Speeding up model inference for production.
-> Distributed Training for LLMs Training large models across multiple GPUs.