AI implementations are often limited by data movement rather than arithmetic throughput. This topic focuses on the memory hierarchy, bandwidth constraints, locality, allocation behavior, and the serving-time structures that determine whether hardware execution units remain supplied with data.
Core Ideas to Master
1. HBM (High Bandwidth Memory) & The Memory Wall
- Concept: The use of stacked DRAM and wide memory interfaces near an accelerator package to provide high off-chip bandwidth.
- Why it's important: Many inference and training kernels cannot sustain peak compute throughput unless operands and intermediate values are delivered at sufficient bandwidth.
- Implementation Link: Explains why optimizations such as tiling, fusion, quantization, and FlashAttention are evaluated by bytes moved as well as operations performed.
- Deep dive: HBM And The Memory Wall
2. Unified Memory Architectures (UMA)
- Concept: CPU and accelerator access through a shared address space or shared physical memory, depending on the system design.
- Why it's important: Unified memory can reduce explicit copy management, but latency, bandwidth, coherence, and page placement still determine performance.
- Implementation Link: Useful for reasoning about local inference, edge deployment, preprocessing pipelines, page migration, and CPU-GPU bandwidth contention.
- Deep dive: Unified Memory Architectures
3. PagedAttention & Virtual Memory for GPUs
- Concept: Applying paging-style indirection to manage the key-value (KV) cache used during autoregressive decoding.
- Why it's important: The KV cache grows dynamically and request lengths vary, so contiguous allocation can waste memory or increase fragmentation.
- Implementation Link: Provides the memory-management mechanism used by serving engines such as vLLM to increase effective cache capacity under dynamic workloads.
- Deep dive: PagedAttention And GPU Virtual Memory
4. Memory Locality & Tiling
- Concept: Decomposing operations into tiles that fit in registers, cache, shared memory, or another faster memory level.
- Why it's important: Tiling increases reuse of data after it has been moved from a slower memory level to a faster one.
- Implementation Link: Central to CUDA, Triton, and compiler-generated kernels for GEMM, attention, normalization, and reductions.
- Deep dive: Memory Locality And Tiling
5. Long-Lived Intermediaries
- Concept: Tensors produced early in training or inference may remain live across many later operations, such as backward propagation or subsequent decoding steps.
- Why it's important: Activation tensors and KV-cache entries can dominate peak memory use even when the model parameters fit in device memory.
- Implementation Link: The direct reason for Activation Checkpointing and KV Cache Optimization.
- Deep dive: Long-Lived Intermediaries
6. Masking the "Shuffle"
- Concept: Scheduling data transfers, communication, and computation so independent work overlaps when hardware resources permit.
- Why it's important: Exposed data movement increases latency and reduces accelerator utilization when computation depends on data that has not arrived.
- Implementation Link: Techniques such as double buffering, asynchronous memory copies, stream scheduling, and communication overlap.
- Deep dive: Masking The Shuffle
Recommended Resources
- vLLM Paper (PagedAttention): "Efficient Memory Management for Large Language Model Serving with PagedAttention."
- FlashAttention Paper: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness."
- NVIDIA Blog on HBM3: Understanding the technical specs and physical constraints of AI memory.