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):
- Distributed-parallelism strategy (FSDP / ZeRO / TP / PP / EP, the
3D-placement decision) →
distributed-training. This reference is the kernel/hardware layer beneath it: the NCCL collective primitives (all-reduce / all-gather, ring vs tree) those strategies call. - Serving-engine configuration and policy (vLLM flags, continuous
batching, speculative decoding, autoscaling) →
llm-inference-serving. This reference is the kernel mechanics beneath it: the paged / quantized KV-cache kernels the engine schedules. - The transformer architecture — why GQA/MLA shrink the KV cache, and
FlashAttention’s math (online softmax derivation, IO-aware exactness) →
transformer-architecture. This reference is the kernel/implementation view of FlashAttention: tiling, SRAM reuse, warp specialization, WGMMA/TMA. - Quantization algorithms (GPTQ, AWQ, SmoothQuant — how to choose the
low-precision weights) →
llm-compression. This reference covers the low-precision kernel / tensor-core mechanics: the FP8 / FP4 / MX / INT8 tensor-core paths that make those quantized weights fast.
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:
- “Why is decode memory-bound but prefill compute-bound?” / roofline reasoning.
- “What is occupancy / a warp / an SM / SIMT?” GPU execution-model questions.
- “How does FlashAttention actually work on the hardware?” (tiling, online softmax in SRAM, warp specialization, ping-pong).
- “What’s the difference between FP8, MXFP8, NVFP4, and INT8 on tensor cores?”
- Writing or reading a CUDA or Triton kernel; memory coalescing; shared
memory bank conflicts;
@triton.autotune. - Kernel fusion — when it helps, when it can’t (reductions).
- Paged / quantized KV-cache kernel internals.
- NCCL collective primitives — ring vs tree, why one is bandwidth-optimal and the other latency-optimal.
- Profiling: Nsight Systems vs Nsight Compute, PyTorch profiler, and computing MFU (Model FLOPs Utilization).
- The hardware landscape: Hopper → Blackwell (H100/H200/B200/GB200), AMD MI300X/MI350X, Google TPU — bandwidth/FLOPs/HBM specs and what they imply.
- Compilers:
torch.compile/TorchInductor, TensorRT-LLM, XLA, Mojo.
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.
- SIMT (Single Instruction, Multiple Thread). Threads are grouped into warps of 32. All 32 threads in a warp execute the same instruction each cycle on different data (lockstep). If threads in a warp take different branch paths (warp divergence), the paths execute serially with the inactive lanes masked off, a major performance killer.
- Latency hiding, not latency reduction. A warp scheduler holds many resident warps and, every cycle, issues from whichever warp is ready. When one warp stalls on a ~400-cycle HBM load, the SM switches to another ready warp instead of idling. This is why GPUs need thousands of threads in flight to reach peak throughput.
- Occupancy = (active warps per SM) / (max warps per SM). It is capped by the scarcest per-SM resource: registers per thread, shared memory per block, or the warp/block hardware limit. Higher occupancy gives the scheduler more warps to hide latency with — but it is a means, not a goal: a register-heavy, well-pipelined kernel can hit peak at modest occupancy, and chasing 100% occupancy by shrinking tiles can hurt (less work reuse). The Triton/CUDA tuning loop is: enough occupancy to hide memory latency, large enough tiles to keep tensor cores fed.
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 |
- HBM3 / HBM3e is the off-chip stacked DRAM: H100 HBM3 ≈ 3.35 TB/s, H200 HBM3e ≈ 4.8 TB/s, B200 HBM3e ≈ 8 TB/s @ 192 GB. It is huge but ~100× slower per byte than SRAM.
- SRAM (shared memory) is the on-chip scratchpad a kernel explicitly manages — orders of magnitude faster than HBM but only ~hundreds of KB per SM. The whole game of a fast kernel is: stage a tile into SRAM, do all the math you can on it, then write back — minimizing HBM round-trips.
- Why attention is IO-bound. Naive attention materializes the full
S = QKᵀscore matrix (size O(seq²)) in HBM, runs softmax over it (another HBM read + write), then multiplies by V (another read). The arithmetic is cheap relative to the bytes moved, so the kernel spends most of its time waiting on HBM. FlashAttention exists precisely to keepSin SRAM and never write it to HBM (see §8).
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:
- A sloped line (peak bandwidth × AI) on the left: the memory-bound region.
- A flat line (peak compute) on the right: the compute-bound region.
- The ridge point is where they cross — the AI at which a kernel transitions from bandwidth-limited to compute-limited. For H100 BF16 the ridge is roughly ~300 FLOP/byte; a kernel below it cannot reach peak FLOP/s no matter how fast the math units are.
The decisive LLM consequence:
- Prefill (prompt processing) is compute-bound. It is a big matrix–matrix multiply (GEMM): a long sequence × weight matrices → high AI, reuses each loaded weight across many tokens. Prefill wants high AI and is limited by tensor-core FLOP/s.
- Decode (autoregressive generation) is memory-bandwidth-bound. To emit one token you must stream the entire weight matrix (and the growing KV cache) from HBM, but you only do a matrix–vector multiply (batch=1) — almost no reuse, AI ≈ O(1). Decode latency ≈ (model bytes + KV bytes) ÷ HBM bandwidth. This is why: (a) decode throughput tracks HBM bandwidth, not FLOPs; (b) batching raises AI (reuse weights across many sequences) and is the single biggest decode-throughput lever; (c) quantization (fewer bytes per weight) directly speeds decode; (d) KV-cache size directly costs latency.
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.
- FP16 / BF16. 16-bit. BF16 (8-bit exponent, 7-bit mantissa) has the same dynamic range as FP32 — the default training/inference precision; rarely overflows, so usually no loss scaling.
- TF32. A 19-bit internal tensor-core mode for FP32 inputs (10-bit mantissa); a near-free Ampere+ speedup for FP32 workloads.
- FP8 (E4M3 / E5M2). 8-bit, native on Hopper+ tensor cores. E4M3 (more mantissa) for forward/weights, E5M2 (more range) for gradients. Needs scaling (per-tensor delayed scaling, or finer) to fit values in the narrow range. H100 FP8 ≈ 2× its BF16 FLOP/s.
- MX microscaling formats (OCP standard). Instead of one scale per tensor/row, an MX block of 32 consecutive elements shares one power-of-two scale stored as UE8M0 (8-bit exponent). MXFP8 = FP8 (E4M3/E5M2) elements + 1×32 block scale; MXFP4 = FP4 (E2M1) elements + 1×32 block scale; MXFP6 also exists. Finer-grained scaling than per-tensor → better accuracy at low bit-width. Native on Blackwell (SM 10.0+).
- NVFP4 (NVIDIA’s Blackwell FP4). Compatible E2M1 elements but a smaller block of 16 with two-level scaling: a per-block FP8 E4M3 scale plus a per-tensor FP32 scale. The smaller block + higher-precision scale localizes to the data’s dynamic range better than MXFP4’s 32-block UE8M0 scale, giving lower quantization error. Blackwell’s 5th-gen tensor cores handle FP4 grouping, dynamic scaling, and the 4-bit MMA in hardware. Reported: ~3.5× smaller memory vs FP16, ~1.8× vs FP8, with <1% accuracy degradation on key LM tasks for many models. Block size is hardware-fixed — picking the wrong block produces a checkpoint the tensor cores can’t consume.
- INT8. Integer 8-bit MMA, very high throughput (H100 INT8 ≈ 2× BF16); common for weight/activation quant (W8A8) and KV-cache quant. Needs careful scale/zero-point calibration.
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:
- A kernel is a function run by many threads. You launch it over a grid
of thread blocks; each block holds up to 1024 threads (executed as
warps of 32).
blockIdx,threadIdx,blockDimindex the data each thread owns. A block runs entirely on one SM and shares that SM’s shared memory. - Memory coalescing is the #1 global-memory rule: when the 32 threads of a warp access consecutive, aligned addresses, the hardware merges them into one (or a few) wide HBM transactions. Strided or scattered access splits into many transactions and wastes most of the bandwidth, the dominant cause of a slow memory-bound kernel.
- Shared memory is the SRAM scratchpad a block uses to stage and reuse data. Tiled matmul is the canonical pattern: each block cooperatively loads a tile of A and B from HBM into shared memory (coalesced), does the partial dot products from SRAM, advances to the next tile. This converts repeated HBM reads into one HBM read + many SRAM reads, raising arithmetic intensity.
- Bank conflicts. Shared memory is split into 32 banks. If multiple threads
in a warp hit different addresses in the same bank, the accesses serialize.
The standard fix is padding (e.g. a
[32][33]tile) so consecutive threads land in distinct banks. - CUDA graphs capture a sequence of kernel launches and replay them as one unit, eliminating per-launch CPU overhead, which matters in decode, where each step is many tiny kernels (see §13 for the compiler/CUDA-graph tie-in).
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.
- Programming model:
@triton.jit;pid = tl.program_id(0)identifies the block;BLOCK_SIZEis atl.constexpr;tl.load(ptr + offs, mask=...)/tl.store(...)move tiles with boundary masks;tl.dot(a, b)issues a tensor-core matmul on tiles. - Autotuning: decorate with
@triton.autotune(configs=[triton.Config({...}, num_warps=, num_stages=), ...], key=[...]). On first call for a new shape (thekey), Triton benchmarks every config on the real tensor sizes and memoizes the winner: cuDNN-style autotuning with no C++ build.num_stagescontrols software pipelining (overlapping loads with compute);num_warpssets the block’s warp count. - Why it matters for LLMs: Triton is the backend
torch.compilegenerates fused kernels into (§13), and the language most custom LLM kernels (fused softmax, layernorm, fused-attention, MoE grouped GEMM, quant kernels) are now written in. The official tutorials walk vector-add → fused-softmax → autotuned matmul → fused-attention → block-scaled (MX) matmul.
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.
- Best targets: chains of elementwise and tile-local ops —
bias → activation → dropout,RMSNorm → matmulpreambles, dequant + matmul. LLM forward/backward fire hundreds of tiny ops; fusing them is a large win on launch overhead and memory traffic. - Hard / impossible to fuse: reductions with long-range dependencies (softmax across a long axis, large all-reduce) need cross-tile/cross-SM communication that breaks single-kernel streaming. FlashAttention is the clever exception: it fuses attention by reformulating softmax into an online/streaming recurrence (§8) so no full-row reduction is materialized.
- Memory-bound ops benefit most (they were limited by bytes, and fusion cuts bytes); compute-bound GEMMs benefit less from fusion itself but still gain from fused epilogues (bias/activation folded into the GEMM store).
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.)
- Tiling. Q, K, V are split into blocks. The kernel loops over K/V blocks, loading each Q/K/V tile from HBM into SRAM, computing the partial scores and partial output there, and accumulating, so the score tile lives only in SRAM and is discarded, never written to HBM.
- Online (streaming) softmax. Softmax normally needs the whole row’s max and
sum first. FlashAttention keeps a running max
mand running denominatorℓand rescales the accumulated output as each new K/V block arrives, producing the exact softmax without ever holding the full row. This is what makes attention fusible into one kernel. - Recomputation in the backward pass. Rather than store the huge
intermediate
S, the backward pass recomputes tiles from the saved stats — trading a little extra FLOPs for a large HBM-traffic/memory saving (a selective-recompute idea). - FlashAttention-2 raised tensor-core utilization by reducing non-matmul FLOPs, better work partitioning across warps, and parallelizing over the sequence dimension.
- FlashAttention-3 (Hopper). Exploits Hopper asynchrony: warp specialization (producer warps issue TMA async copies HBM→SRAM while consumer warps run WGMMA tensor-core matmuls), ping-pong scheduling between two warpgroups (one does GEMM while the other does softmax — ~570→620 TFLOPS), and intra-warpgroup pipelining of softmax with GEMM (→~640–660 TFLOPS FP16). It adds FP8 attention with incoherent processing (a random-sign Hadamard transform in O(d log d) to spread outliers), cutting FP8 error ~2.6× vs baseline. Result: ~740 TFLOPS FP16 (~75% of H100 peak, up from ~35%), 1.5–2.0× over FA-2; ~1.2 PFLOPS in FP8.
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:
- PagedAttention kernel. Inspired by OS virtual memory: the KV cache is
stored in fixed-size blocks (pages), not one contiguous per-sequence
buffer. A per-sequence block table maps logical token positions to
physical blocks, so blocks can be allocated on demand and shared across
sequences (e.g. a shared prompt prefix, or beams). The attention kernel
gathers K/V through the block table instead of a flat stride. This cuts KV
fragmentation/waste to <4% and is what lets a server pack many more concurrent
sequences (2–4× throughput). Engine policy (which sequences to batch,
eviction, prefix caching) is
llm-inference-serving; this is the kernel that the policy schedules. - Quantized KV-cache kernels. Storing K/V in FP8 (recommended on Hopper/Blackwell) or INT8 halves/quarters KV bytes, directly relieving the memory-bandwidth-bound decode path and extending context length. The kernel must dequantize on the fly inside the attention compute (or use low-precision MMA paths). vLLM ships FP8 KV-cache; INT8 KV-cache kernels (naive/tiled/ coarsened/vectorized variants) report up to 4× KV memory reduction with small accuracy loss.
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:
- all-reduce — sum (or other op) a tensor across all ranks, every rank gets the result (gradient sync in data parallel; the per-block sum in tensor parallel).
- all-gather — each rank contributes a shard, every rank ends with the full concatenation (FSDP/ZeRO parameter gather).
- reduce-scatter — reduce then partition (the FSDP gradient half; reduce-scatter + all-gather = one ring all-reduce).
- all-to-all — every rank sends a distinct piece to every other rank (MoE expert dispatch/combine).
Ring vs tree (the central trade-off):
- Ring all-reduce arranges ranks in a logical ring and streams shards around it (reduce-scatter phase + all-gather phase). It is bandwidth-optimal — each link is fully utilized and per-rank traffic is independent of rank count — but its latency grows linearly with the number of ranks (~2(N−1) steps), so it is poor for tiny messages at large scale.
- Tree all-reduce reduces up a (double-)binary tree and broadcasts down. Latency is logarithmic in N, so it wins for small, latency-sensitive messages and large clusters.
- NCCL auto-selects per call: it models each algorithm×protocol’s latency and
bandwidth and picks the predicted winner by message size — tree for small,
ring for large — and tunes to the topology (NVLink/NVSwitch intra-node, the
network inter-node). Newer PAT (Parallel Aggregated Trees) gives
logarithmic all-gather/reduce-scatter at scale. Overlapping collectives with
compute (so comm hides behind matmuls) is the strategy-level lever covered in
distributed-training.
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.
- MFU (Model FLOPs Utilization) = (model’s useful FLOP/s, e.g. the
6NDtraining estimate or the inference FLOPs) ÷ (the hardware’s peak FLOP/s at that precision). It is hardware-agnostic and tells you how close you are to the roofline. 40–50% sustained MFU is a good real-world training target; decode is far lower because it is memory-bound (low AI), so MFU is the wrong lens for decode — there, % of peak bandwidth is the metric. (HFU, hardware FLOPs utilization, additionally counts recomputed FLOPs.) - Nsight Systems (
nsys) — system-wide timeline: CPU↔GPU overlap, kernel gaps, stream/launch behavior, NCCL. Low overhead; the first tool — find the top/longest or stalling kernels and the bubbles. - Nsight Compute (
ncu) — single-kernel deep dive: achieved occupancy, memory vs compute bound, the kernel’s roofline, warp-stall reasons, bank conflicts. The second tool, once Nsight Systems names the suspect kernel. - PyTorch profiler (
torch.profiler+ TensorBoard / Holistic Trace Analysis / Chrome trace) — medium overhead, framework-aware: maps kernels back to model ops, with stack traces, shapes, and memory. Best for “which layer is slow” and for correlating Python with kernels. - Typical loop: Nsight Systems → find bubbles / a hot kernel → Nsight Compute → classify (memory- vs compute-bound on the roofline) → fix (coalesce, fuse, raise occupancy/tile size, change precision) → re-measure 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.
torch.compile(PyTorch 2.x). Front end TorchDynamo captures the graph; back end TorchInductor lowers it to fused Triton kernels (GPU) / C++/OpenMP (CPU). Fusion of elementwise chains + reduced launch overhead is the main speedup; pairs with CUDA graphs to kill per-launch cost. It is the default acceleration path and the one vLLM now uses (-O3, piecewise CUDA graphs) for its model code.- TensorRT-LLM. NVIDIA’s inference compiler/runtime: aggressive deep
graph fusion, fused multi-head attention, FP8/FP4 paths, in-flight batching,
piecewise CUDA graphs, and it now uses
torch.compilefor lightweight vertical fusion. Tends to win on large models where deep fusion pays off;torch.compilealone can match or beat it on smaller models. - XLA. Google’s array compiler (JAX, TF, TPU; PyTorch/XLA). The native path
for TPUs and whole-graph fusion via HLO; on NVIDIA GPUs its gains over
torch.compileare usually modest. - Mojo. Modular’s Python-superset systems language aimed at writing portable high-performance kernels (an alternative to CUDA C++/Triton, MLIR-based); emerging, not yet a default in mainstream LLM stacks — watch, don’t depend.
Practical patterns
- Diagnose with the roofline first. Before optimizing, classify the kernel: memory-bound or compute-bound (Nsight Compute draws this). Memory-bound → coalesce, fuse, raise reuse, drop precision. Compute-bound → bigger tiles, use tensor cores, lower precision.
- Decode = bandwidth problem. To speed decode: batch (raise AI),
quantize weights + KV (fewer bytes), shrink the KV cache (GQA/MLA — an
architecture lever in
transformer-architecture), use CUDA graphs to kill launch overhead. Do not expect more FLOPs to help. - Prefer Triton + autotune over hand-CUDA for new custom kernels unless you
need an instruction the DSL can’t express; let the compiler handle banks and
coalescing, and let
@triton.autotunefind tile sizes per shape. - Let
torch.compilefuse first. Reach for hand kernels only where the compiler leaves bandwidth on the table (profile to prove it). - Match the precision to the hardware’s fixed block size. On Blackwell, NVFP4 wants 16-element blocks, MXFP8 wants 32 — quantize to the format the tensor cores actually consume.
- Measure MFU for training, % peak bandwidth for decode. Report the right metric for the regime; a “100% GPU utilization” claim with low MFU means the tensor cores are starved.
Anti-patterns
- Chasing 100% occupancy. Occupancy is a means to hide latency, not a target; shrinking tiles to raise it can lower reuse and hurt throughput.
- Trusting
nvidia-smi“GPU-Util”. It reports time a kernel was resident, not tensor-core efficiency; use MFU / Nsight Compute instead. - Strided / uncoalesced global access in the hot loop — the most common cause of a memory-bound kernel running at a fraction of HBM bandwidth.
- Trying to fuse a long-range reduction (e.g. naive softmax across the whole sequence) into one streaming kernel — it needs cross-tile communication; reformulate (online softmax) or keep it separate.
- Optimizing FLOPs to speed decode. Decode is bandwidth-bound; FLOP-side tuning yields little. Optimize bytes moved.
- Ignoring the fixed MX/NVFP4 block size — a wrong block size produces a checkpoint the tensor cores cannot run.
- Hand-writing CUDA before profiling. Premature kernel hacking before the roofline tells you what’s actually limiting.
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)
transformer-architecture— FlashAttention’s math (online-softmax derivation, IO-aware exactness), GQA/MLA why, the KV cache concept. This file is the kernel implementation of those.llm-inference-serving— serving-engine policy (vLLM batching, prefix caching, speculative decoding, autoscaling) on top of the paged/quantized KV kernels here.distributed-training— parallelism strategy (FSDP/ZeRO/TP/PP/EP, 3D placement, compute–comm overlap) on top of the NCCL collective primitives here.llm-compression— quantization algorithms (GPTQ/AWQ/SmoothQuant) that target the FP8/FP4/MX/INT8 tensor-core paths here.
References
- NVIDIA — Introducing NVFP4 for Efficient and Accurate Low-Precision Inference (developer.nvidia.com, 2025).
- NVIDIA Transformer Engine docs — MXFP8 / Using FP8 and FP4 (OCP MX block formats, UE8M0 scaling), 2025.
- OCP — Microscaling (MX) Data Formats for Deep Learning spec / arXiv 2310.10537.
- Tri Dao et al. — FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608; tridao.me blog; PyTorch blog), 2024.
- OpenAI / Triton — Introducing Triton and the official tutorials (triton-lang.org): fused softmax, autotuned matmul, fused attention, block-scaled matmul.
- NVIDIA NCCL — developer docs + Understanding NCCL Tuning and Massively Scale … with NCCL (ring vs tree algorithm selection); PAT algorithm (arXiv 2506.20252).
- vLLM docs — PagedAttention design and Quantized KV Cache (FP8); INT8 KV-cache quantization (arXiv 2601.04719).
- LLM Inference Unveiled: Survey and Roofline Model Insights (arXiv 2402.16363); A Systematic Characterization of LLM Inference on GPUs (arXiv 2512.01644).
- PyTorch — Why Is PyTorch Compile So Fast: Kernel Fusion; Introduction to torch.compile and How It Works with vLLM (vLLM blog, 2025).
- NVIDIA TensorRT-LLM docs — Torch Compile & Piecewise CUDA Graph; Collabora torch.compile vs TensorRT (2024).
- Trainy — GPU Utilization Is a Misleading Metric; Using Model FLOPs Utilization (MFU); NVIDIA Profiling LLM Training Workflows on Grace Hopper (Nsight Systems/Compute).
- Hardware comparisons — Exxact Blackwell vs Hopper; SemiAnalysis MI300X vs H100/H200; Artificial Analysis TPU v6e vs MI300X vs H100/B200 (2025).
- Deep Kernel Fusion for Transformers (arXiv 2602.11808) — fusion targets, HBM-traffic reduction.
- Siboehm — How to Optimize a CUDA Matmul Kernel (coalescing, tiling, shared-memory, bank conflicts worklog).