ML Systems & Inference
Attention internals, the KV cache, inference performance, quantization, parallelism, GPU/TPU kernels, and JAX.
-
Training vs Inference
A trained model is a fixed function; inference is the forward-only generation loop. Where the compute goes.
-
NumPy & JAX
Hands-on array programming: broadcasting, indexing, vectorization, einsum, and the JAX numpy API (functional updates, vmap/grad/jit in practice, PRNG keys, pytrees).
-
Attention Internals
The full softmax(QK^T / sqrt(d)) V derivation, why the sqrt(d) scaling, RoPE and relative position, causal masking.
-
Inference Performance
The roofline model, memory- vs compute-bound, arithmetic intensity, static vs continuous batching.
-
KV Cache & Decoding
What the KV cache stores, prefill vs decode, why memory grows with context, sampling strategies.
-
Quantization & Low Precision
bfloat16 range vs precision, FP8 and INT8, mixed precision, where dropping precision is safe.
-
GPU / TPU Architecture
The memory hierarchy (HBM/SRAM, VMEM/SMEM), warps and SIMT, tensor cores, occupancy, coalescing.
-
Parallelism & Collectives
Data, tensor, and pipeline parallelism, sharding, all-reduce and all-to-all, NVLink and InfiniBand.
-
Kernels (CUDA / Triton / Pallas)
What a kernel is, CUDA vs Triton vs Pallas, and the FlashAttention idea: tiling and online softmax.
-
JAX & XLA Mental Model
Pure functions, jit/grad/vmap, pytrees, NNX split/merge, and JAX to HLO to device.
-
Mixture-of-Experts
Top-k routing, all-to-all dispatch and combine, load balancing and capacity factors.
-
Frontier & Recent Research
Recent inference and architecture research: FlashAttention, DeepSeek (MLA, aux-loss-free MoE, FP8, multi-token prediction), speculative decoding, PagedAttention, linear/state-space models, and newer attention/softmax variants.
Up next
105 of 105 questions here are still unsolved.