GPU & Accelerator Kernels for LLMs

GPU & Accelerator Kernels for LLMs

The hardware substrate every other model-layer skill sits on. Pretraining, fine-tuning, alignment, inference serving, and compression all ultimately resolve to kernels: GPU programs that move bytes through a memory hierarchy and feed tensor cores. This reference is the implementation view: how a GPU executes work, why LLM attention and decode are bottlenecked by memory bandwidth rather than FLOPs, and how kernels (CUDA, Triton, FlashAttention, paged-KV) and compilers (torch.compile, TensorRT-LLM, XLA) are written to fight that bottleneck.

Where this sits among siblings (read the boundary, then the right file):


When to load this reference

Load when the task is about why GPU code is slow and how to make it fast at the kernel level, not about which model or which parallelism strategy:


Core concepts (MECE)

1. The GPU execution model — SMs, warps, SIMT, occupancy

A GPU is a throughput machine built to hide latency with parallelism, the opposite of a latency-optimized CPU. The unit of compute is the Streaming Multiprocessor (SM) — a B200 has ~148 SMs, an H100 ~132. Each SM contains arithmetic units (FP/INT), tensor cores (matrix-multiply accelerators), a register file, shared memory / L1, and one or more warp schedulers.

2. The memory hierarchy — and why attention is IO-bound

Speed and capacity trade off inversely at every level. Approximate H100/B200 figures:

Level Capacity Bandwidth Latency
Registers (per thread) ~256 KB/SM file ~tens of TB/s ~1 cycle
Shared memory / L1 (SRAM) ~228 KB/SM (H100) ~tens of TB/s ~30 cycles
L2 cache ~50 MB ~10 TB/s ~200 cycles
HBM (global / VRAM) 80–192 GB 3.35–8 TB/s ~400+ cycles

3. Arithmetic intensity and the roofline model

Arithmetic intensity (AI) = FLOPs performed ÷ bytes moved from memory (FLOP/byte). The roofline model plots attainable FLOP/s against AI:

The decisive LLM consequence:

4. Precision and tensor cores — BF16 / FP8 / FP4 / MX / INT8

Tensor cores are dedicated matrix-multiply-accumulate (MMA) units: they multiply small tiles (e.g. 16×16) and accumulate, delivering ~10–30× the FLOP/s of the general FP units. Lower precision = more tensor-core throughput and fewer bytes to move (helps the memory-bound regime), so the precision frontier is the central lever for both training and inference speed.

The algorithm for choosing which weights to quantize and how (GPTQ, AWQ, SmoothQuant) lives in llm-compression. This reference is the tensor-core path those algorithms target.

5. CUDA basics — kernels, grids/blocks, coalescing, shared memory

CUDA is the C++ dialect for writing GPU kernels. The launch hierarchy:

6. Triton — block-level kernels and autotuning

OpenAI Triton is a Python DSL+compiler for GPU kernels at a tile (block) granularity, sitting between hand-CUDA and framework ops. You write what each program instance does to a block of data; the compiler handles intra-block thread scheduling, vectorization, shared-memory allocation, and coalescing automatically — you do not manage individual threads or banks.

7. Kernel fusion

Fusion combines a sequence of operations into a single kernel so intermediates stay in registers/SRAM and are never written to HBM. It removes (a) HBM round-trips of intermediate tensors and (b) per-op kernel-launch overhead.

8. FlashAttention — the kernel case study (implementation view)

FlashAttention is IO-aware exact attention: same result as standard attention, but it never materializes the O(seq²) score matrix in HBM. (The math/derivation and the architecture motivation live in transformer-architecture; here is how the kernel is built.)

9. Paged and quantized KV-cache kernels

The KV cache (cached keys/values for every past token) grows with sequence and batch and dominates decode memory. Two kernel-level techniques:

10. NCCL collectives — the communication primitives (ring vs tree)

When a model spans many GPUs, the parallelism strategy (distributed-training) is implemented on top of NCCL collective primitives. The ones that matter:

Ring vs tree (the central trade-off):

11. Profiling and Model FLOPs Utilization (MFU)

You cannot optimize what you cannot measure; raw “GPU utilization” (percent of time a kernel was resident) is misleading — it can read 100% while tensor cores sit mostly idle. The real efficiency metric is MFU.

12. The hardware landscape (Hopper → Blackwell, MI300X, TPU)

Bandwidth and FLOPs set the roofline; HBM capacity sets how big a model/KV cache fits. Approximate per-accelerator figures (2024–mid-2026):

Accelerator Arch HBM Bandwidth Peak dense tensor Notes
H100 Hopper 80 GB HBM3 ~3.35 TB/s ~990 TF BF16 / ~1979 TF FP8 WGMMA, TMA, FP8; the 2023–24 workhorse
H200 Hopper 141 GB HBM3e ~4.8 TB/s same as H100 bandwidth/capacity bump → faster decode
B200 Blackwell 192 GB HBM3e ~8 TB/s ~4.5 PF BF16 / ~9 PF FP8 5th-gen tensor cores, native MXFP8/NVFP4, FP4
GB200 Blackwell Grace+2×B200 NVLink-C2C ~3–3.4× H100/GPU NVL72 rack = 72 GPUs on one NVLink fabric
AMD MI300X CDNA3 192 GB HBM3 ~5.3 TB/s high BF16/FP8 (ROCm) big HBM; CUDA-moat gap on software/kernels
AMD MI350X CDNA4 288 GB HBM3e ~8 TB/s + FP4/FP6 most HBM capacity; competes with B200
Google TPU v6e TPU HBM high ~0.918 PF BF16 systolic MXU, XLA-only, pod-scale ICI

Takeaways: (1) each generation’s bandwidth jump is what speeds decode; (2) Blackwell’s FP4/MXFP8 is what makes 4-bit inference fast in hardware; (3) NVLink/NVSwitch fabric (GB200 NVL72) makes large collectives intra-fabric; (4) NVIDIA’s lead is partly the kernel/software moat (CUDA, cuDNN, NCCL, TensorRT, FlashAttention) — AMD MI300X has competitive silicon (more HBM) but historically trails on ready kernels; TPU is strong but XLA-only (no CUDA).

13. Compilers — torch.compile / TorchInductor, TensorRT-LLM, XLA, Mojo

Compilers turn a high-level model graph into fused, scheduled kernels so humans don’t hand-write each one.


Practical patterns

Anti-patterns

Troubleshooting

Symptom Likely kernel-level cause Where to look
Decode throughput far below HBM bandwidth ÷ model bytes Tiny per-step kernels, launch overhead, no CUDA graph; KV not paged Nsight Systems timeline (gaps); enable CUDA graphs / paged KV
Low MFU in training but GPU “100% util” Memory-bound or starved tensor cores; small tiles Nsight Compute roofline; bigger tiles, fuse, FP8/BF16
Attention OOMs at long context Materializing O(seq²) scores Use FlashAttention; check it’s actually engaged
Multi-GPU step dominated by comm Collective not overlapped; wrong ring/tree at this size nsys for NCCL; let NCCL tune; overlap (strategy → distributed-training)
FP8/FP4 accuracy collapses Bad scaling / wrong block size / outliers Per-block (MX/NVFP4) scaling; Hadamard/incoherent processing; recalibrate
Slow shared-memory kernel Bank conflicts Nsight Compute “shared bank conflicts”; pad tiles
Quantized weights load but run slow Falling off the tensor-core fast path (format mismatch) Confirm the quant format matches the GPU’s MMA path

Cross-references (reciprocal)

References

  1. NVIDIA — Introducing NVFP4 for Efficient and Accurate Low-Precision Inference (developer.nvidia.com, 2025).
  2. NVIDIA Transformer Engine docs — MXFP8 / Using FP8 and FP4 (OCP MX block formats, UE8M0 scaling), 2025.
  3. OCP — Microscaling (MX) Data Formats for Deep Learning spec / arXiv 2310.10537.
  4. Tri Dao et al. — FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608; tridao.me blog; PyTorch blog), 2024.
  5. OpenAI / Triton — Introducing Triton and the official tutorials (triton-lang.org): fused softmax, autotuned matmul, fused attention, block-scaled matmul.
  6. NVIDIA NCCL — developer docs + Understanding NCCL Tuning and Massively Scale … with NCCL (ring vs tree algorithm selection); PAT algorithm (arXiv 2506.20252).
  7. vLLM docs — PagedAttention design and Quantized KV Cache (FP8); INT8 KV-cache quantization (arXiv 2601.04719).
  8. LLM Inference Unveiled: Survey and Roofline Model Insights (arXiv 2402.16363); A Systematic Characterization of LLM Inference on GPUs (arXiv 2512.01644).
  9. PyTorch — Why Is PyTorch Compile So Fast: Kernel Fusion; Introduction to torch.compile and How It Works with vLLM (vLLM blog, 2025).
  10. NVIDIA TensorRT-LLM docs — Torch Compile & Piecewise CUDA Graph; Collabora torch.compile vs TensorRT (2024).
  11. Trainy — GPU Utilization Is a Misleading Metric; Using Model FLOPs Utilization (MFU); NVIDIA Profiling LLM Training Workflows on Grace Hopper (Nsight Systems/Compute).
  12. Hardware comparisons — Exxact Blackwell vs Hopper; SemiAnalysis MI300X vs H100/H200; Artificial Analysis TPU v6e vs MI300X vs H100/B200 (2025).
  13. Deep Kernel Fusion for Transformers (arXiv 2602.11808) — fusion targets, HBM-traffic reduction.
  14. Siboehm — How to Optimize a CUDA Matmul Kernel (coalescing, tiling, shared-memory, bank conflicts worklog).