Foundation Model Kernel Optimization in 2026: A Field Guide Across Dense, MoE, Multimodal, and Diffusion

Most “kernel optimization” conversations still start with FLOPs. In 2026 that is usually the wrong first question. Foundation-model runtime is dominated by HBM traffic, KV cache, temporary tensors, collectives, dynamic permutation, and launch overhead. FlashAttention, fused linear–cross-entropy, paged KV, MoE grouped GEMM, and DeepEP-style dispatch all share one essence: do not materialize intermediates, or make each byte travel once.

This post is a field guide to the kernel / fused-kernel / collective families that matter on critical paths — Dense Transformers, MoE, LLM train/serve, VLM / audio / video, contrastive multimodal training, Diffusion / DiT, and SSM / long convolution — plus the communication, memory, layer-structure, and model–kernel co-design moves that decide whether a fast op shows up in end-to-end tokens/s.

Scope note. “Every kernel” here means families on the critical path, not every CUDA symbol in a framework. Batch, sequence length, hidden size, GPU generation, and parallel strategy spawn different instances; symbol-level enumeration is unstable and rarely useful.

Related posts on this blog: Roofline, GPU Optimization Playbook, Large MoE Performance, torch.compile, VLM one-GPU recipe.

TL;DR

1. Optimize data movement first, not peak FLOPs. Fusion, paging, packing, and low precision win by cutting bytes and launches.
2. Fusion is layered. Elementwise chains are cheap wins; GEMM epilogue fusion is the safe high-ROI boundary; GEMM↔GEMM and compute↔comm fusion need shared layout/pipeline and can hurt occupancy.
3. Train ≠ prefill ≠ decode. Prefill is large-GEMM heavy; decode is GEMV / small-GEMM + KV bandwidth; training adds backward, optimizer, and cross-rank sync.
4. Attention co-design is hardware-asymmetric. FA1 fixed IO; FA2 fixed work partition; FA3 used Hopper TMA/WGMMA; FA4 rebalances Blackwell Tensor Core / SFU / TMEM / 2-CTA ratios.
5. MoE is a pipeline. Router → top-k → histogram → permute → All-to-All → grouped GEMM → act → combine. Speeding only expert GEMM often amplifies communication.
6. Overlap needs finer dependencies than “another CUDA stream.” Chunk/tile/threadblock events beat whole-tensor collectives.
7. Memory has six levers: dematerialize, compress, pack/page/layout, recompute, shard/offload/prefetch, allocator lifetime. Peak GB alone is a bad KPI.
8. Multimodal’s biggest lever is usually token budget (native resolution, packing, temporal sparsity) — architecture–runtime co-design, not one vision kernel.
9. Diffusion’s unique lever is cross-timestep redundancy (cache, fewer steps, CFG parallel) — fewer full-network calls, not only faster ops.
10. End-to-end KPIs win. Kernel TFLOP/s without graph-break / layout / exposed-comm / p99 / quality checks is storytelling.

1. A practical performance model

For a kernel or path, a lower bound looks like:

[ T \gtrsim \max\left( \frac{F}{P_{\text{effective}}}, \frac{B_{\text{HBM}}}{BW_{\text{HBM}}}, \frac{B_{\text{net}}}{BW_{\text{net}}} \right) +T_{\text{launch}}+T_{\text{sync}}+T_{\text{imbalance}} ]

• (F): useful FLOPs; (P_{\text{effective}}) is limited by Tensor Core use, tile choice, occupancy, dependency chains, SFU throughput.
• (B_{\text{HBM}}): bytes that actually cross HBM — materialization, cold KV, layout copies inflate this beyond logical tensor size.
• (B_{\text{net}}): NVLink / PCIe / IB bytes; topology, algorithm, message granularity, and congestion set effective bandwidth.
• Launch / sync / imbalance dominate small-batch, short-sequence, fine-grained MoE, and ragged workloads.

NVIDIA’s Nsight Compute Profiling Guide splits kernels into memory- vs compute-bound. Foundation models need two more axes: network and launch / load imbalance. For the arithmetic-intensity framing, see Roofline on this blog.

1.1 Symptom → direction

2. A dictionary of techniques

2.1 Different GPU generations want different “good” kernels

Do not copy H100 block sizes, warp counts, and stage counts onto B200 or MI300X. Tensor Core throughput often outruns SFU, register bandwidth, SMEM, and the network — the non-matmul and orchestration paths become the new Amdahl bottleneck.

3. Dense Transformer / LLM / sequence kernels

Notation: B batch, S sequence, H hidden, V vocab, D head dim. Covers train fwd/bwd, prefill, decode, loss, optimizer.

3.1 What “fuse elementwise and GEMM” actually means

The safe high-ROI boundary is usually the GEMM epilogue. For Y = act(XW + b) * gate + residual:

1. Mainloop accumulates XW on Tensor Cores.
2. While the accumulator tile is still in register/TMEM, add bias/scale.
3. Apply activation, gate, residual, dtype convert / quant.
4. Write final Y to HBM once.

That avoids three-to-four large intermediates. On Blackwell, accumulators live in TMEM, so epilogue traffic and SFU scheduling differ from Hopper; CODA goes further and expresses many non-attention Transformer ops as GEMM + epilogue programs.

Do not fuse blindly when:

• Per-thread registers explode into spills;
• Ops want incompatible tiles / parallel axes;
• Reductions need cross-CTA sync;
• One producer fans out to many consumers (fusion duplicates work);
• Dynamic shapes explode JIT variants;
• Backward recompute turns a compute-bound path slower.

3.2 Four generations of FlashAttention

Lesson: when one resource gets faster, the bottleneck moves. FA4 is not “FA3 ported to Blackwell”; it rewrites the schedule around new resource ratios.

4. MoE: optimize the whole token–expert pipeline

With T tokens, E experts, top-k=K, MoE cost is not “K expert MLPs.” It includes T×E router logits, top-k, counts/offsets, permute, cross-rank dispatch/combine, imbalance tails, and the backward inverse path. For the systems framing of memory / communication / compute walls, see Large MoE Performance on this blog.

4.1 The right overlap granularity

Bad: full dispatch → all expert GEMMs → full combine.

Better:

1. Chunk tokens by destination / rank / expert.
2. Enqueue grouped-GEMM tasks as chunks arrive.
3. Push completed output tiles into the combine send queue immediately.
4. Specialize communication vs compute CTAs (or event-driven streams/communicators).
5. Cap communication CTAs so they do not steal all SM/L2/HBM.

FLUX decomposes then re-fuses communication and computation at fine tiles (paper: up to ~96% communication hidden in some settings). Comet uses threadblock specialization to balance resources. The KPI is exposed communication on the critical path, not “did the NCCL kernel overlap.”

4.2 Megatron-Core MoE map (flags change by version)

As of mid-2026, Megatron-Core 0.17 MoE is a combinator of routing, dispatcher, parallel map, precision, fusion, overlap, and memory strategies. Authoritative sources: Megatron-Core 0.17 MoE docs and NVIDIA’s 2026 MoE tech report.

Recommended triage order: model semantics → parallel mapping (keep EP×ETP on NVLink/MNNVL; fold Attn vs MoE) → dispatcher → GroupedGEMM + router/permute fusion + FP8 recipe → shared-expert / batch A2A overlap / CUDA Graph / recompute-offload (ablate — contention is common).

5. Cross-rank communication, parallelism, and scheduling

5.1 Four levels of overlap

Domino splits tensors and batches into independently advancing pieces. Measure with Nsight Systems:

[ T_{\text{exposed-comm}} = T_{\text{step}} - T_{\text{ideal-compute-only}} - T_{\text{other}} ]

Do not sum stream durations. Contended NCCL kernels can look “longer” while the critical path shortens.

6. Multimodal: vision, audio, video, fusion

Three layers matter:

1. Input pipeline — decode, resize, crop, normalize, resample, STFT/mel: GPU/HW decode, async prefetch, fewer layout copies.
2. Modality encoders — mostly Transformer/conv kernels, but token count, window/spatiotemporal sparsity, and ragged packing dominate.
3. Fusion & losses — projectors/resamplers, cross/joint attention, contrastive similarity, generative LM loss: kill padding, avoid B²/S² materialization, reduce global sync.

6.1 Multimodal priority is usually tokens, not patch kernels

Approx. VLM cost with visual tokens (S_v): per-layer linear ~(O((S_t+S_v)H^2)), global attention ~(O((S_t+S_v)^2 H)).

• 2× faster patch embedding touches one shallow op.
• 2× fewer (S_v) cuts linear, attention, activation, KV, and communication on every later layer.
• Removing 40% padding via packing compounds the same way without changing semantics.

That is why NaViT native-resolution packing, Qwen2-VL dynamic resolution, resampler latents, and video spatiotemporal sparsity so often beat an isolated CUDA micro-optimization. Practical one-GPU stacking notes: Optimizing VLM Training on One GPU.

7. Diffusion, DiT, U-Net, generative video

Diffusion’s distinctive structure is “the same denoiser, many times.” Optimize reuse across timesteps, solver step count, CFG branching, and cross-device patch/pipeline — not only per-step kernels.

7.1 Do not mix three kinds of speedup

Report ablations separately. “8 steps instead of 50” is not a kernel win.

8. Memory system: six levers, not only checkpointing

8.1 By object

Parameters & optimizer. BF16 weights + FP32 master + FP32 m/v can dwarf the weights. FSDP/ZeRO is the capacity lever; fused/8-bit optimizers are the bandwidth/state lever. INT4 weight-only becomes decode speed only with prepacked layout + dequant-in-mainloop (Marlin-class designs). FP8/MXFP8/NVFP4 scale granularity must match GEMM tiles; TE distinguishes communication-compact vs GEMM-ready layouts and swizzles between them.

Activations. Dematerialize first, checkpoint second. Prefer selective recompute on “large storage, cheap FLOPs” ops. Reducing Activation Recomputation in Large Transformer Models combines sequence parallel with selective recompute. PyTorch’s activation checkpointing techniques and compile min-cut partitions help — until graph breaks intervene.

KV cache. Capacity ≈ layers × tokens × KV_heads × head_dim × 2 × bytes. GQA/MQA cut heads; MLA uses low-rank latents; quant cuts bytes; paging cuts fragmentation; prefix sharing cuts duplicates. Smaller pages → less fragmentation, more indirection; larger pages reverse that. KV quant must include scale reads and dequant; decode is usually BW-bound, so a little extra arithmetic often still wins.

Workspace & allocator. Faster cuBLASLt/cuDNN/CUTLASS/collective algos may want large workspace — OOM is not always “model tensors.” Preallocate max workspace; avoid per-token cudaMalloc; CUDA Graph wants stable addresses. For dynamic serving, track allocated / reserved / payload / fragmentation / page metadata separately from max_memory_allocated.

8.2 Sequence packing boundaries

Packing kills padding only if attention knows segment edges:

• Describe segments with cu_seqlens / offsets;
• Reset causal masks inside each segment;
• Remap position ids, RoPE/M-RoPE, and loss masks together;
• Persist multimodal special tokens, image grids, audio time indexes with pack metadata;
• Honor boundaries in backward and context parallel.

Concat + vanilla causal attention leaks across samples. See also Why Variable Sequence Length Breaks DDP Throughput.

9. Compilers, DSLs, libraries, megakernels

Further reading: TorchInductor design notes, torch.compile tutorial, Triton persistent matmul, block-scaled matmul, ThunderKittens.

Division of labor that usually works:

1. Compile away peripheral pointwise/reshape/simple reduction/epilogue; clear graph breaks.
2. Call specialized kernels for attention, linear-CE, MoE, low-bit GEMM, decode KV, sampling.
3. Unify layouts at those boundaries so transpose/cast/contiguous copies disappear.
4. AOT/autotune-cache hot shapes; robust fallbacks for the long tail.
5. Only write a megakernel when profiles prove the op boundary itself is top bottleneck and shapes are stable.

10. Model–kernel co-design that actually moves the needle

10.1 Two judgment criteria

A — Does the win compound across layers? Fewer tokens / KV / denoising steps ride through tens or hundreds of layers. A 2× faster projector does not.

B — Does the new structure map to regular tiles? Block sparsity, fixed latents, small top-k experts, and block scaling tend to run well. Random token-level sparsity and arbitrary tiny ragged ops often look great on paper FLOPs and lose to dense on GPUs.

10.2 Parallel decoder / Parallel Transformer Block

Here “parallel” is not tensor or pipeline parallel. It rewrites one decoder block’s dependency graph.

Serial pre-norm:

[ u=x+\operatorname{Attn}(\operatorname{Norm}_a(x)),\qquad y=u+\operatorname{MLP}(\operatorname{Norm}_m(u)) ]

Parallel Transformer Block (PaLM / shared-norm form):

[ z=\operatorname{Norm}(x),\qquad y=x+\operatorname{Attn}(z)+\operatorname{MLP}(z) ]

Falcon-style separate-norm variant:

[ y=x+\operatorname{Attn}(\operatorname{Norm}_a(x))+\operatorname{MLP}(\operatorname{Norm}_m(x)) ]

Because MLP no longer depends on attention output, systems gain scheduling freedom: branch-level parallelism, shared input/norm reads, horizontal QKV vs gate/up scheduling, potential collective merge after local branch combine, and freer backward reordering. It does not magically cut major GEMM FLOPs, and concurrent branches can contend for Tensor Cores, HBM, L2, and SMs.

PaLM uses parallel layers. MegaScale applies PTB in Megatron-LM and pipelines SP/TP AllGather/ReduceScatter with FFN linears; on a 175B / 256-GPU ablation they report MFU 47.7%→52.3% (+4.6 pp), and 53.3% with sliding-window attention — workload-specific. NeMo’s Megatron-Core customization docs list Falcon’s parallel-attention/MLP variant.

PTB is worth considering for full pretraining (or large continued-pretrain budget) when attention and MLP durations are comparable or one branch can cover the other’s collectives, TP/SP layouts allow clean merge, and schedulers can limit resource fights. It is usually not a lossless kernel swap onto an existing Llama-class checkpoint for short fine-tunes.

10.3 Algorithmic speedups to book separately

Book these as algorithmic speedup, then ablate against kernel, communication, and runtime gains — otherwise you cannot tell whether you won by computing less, rewriting the graph, hiding communication, or making one op faster.

11. Priority lists by scenario

Dense LLM training

1. Confirm BF16/FP8 GEMM and FlashAttention are actually running; clear layout/cast fallbacks.
2. Fused RMSNorm/RoPE/SwiGLU and fused linear-CE to kill large intermediates/logits.
3. Sequence packing + selective recompute for padding and activation peaks.
4. Chunk-overlap TP/FSDP AG/RS with GEMM; then pipeline schedule.
5. Fused optimizer / multi-tensor grad norm; only then chase embeddings and tiny pointwise.

LLM prefill / high-throughput serving

1. FlashAttention varlen + continuous batching for the right head dim / mask / precision.
2. Large-GEMM epilogue fusion, FP8, shape buckets, CUDA Graph.
3. Paged KV writes, prefix sharing, chunked prefill to protect decode.
4. TP collective overlap; jointly schedule TTFT, throughput, and decode p99 — not only prefill TFLOP/s.

LLM decode / low-latency serving

1. Cut KV bytes first: MQA/GQA/MLA, KV dtype, context, page/prefix hits.
2. Paged / split-KV decode attention; JIT/autotune batch/head/page.
3. INT4 weight-only / FP8 GEMV–small GEMM; continuous batching; persistent kernels.
4. CUDA Graph, fused sampling, LoRA SGMV/BGMV; shrink CPU scheduler/launch.
5. Multi-GPU: account for per-token TP AllReduce — small batches can be network-latency dominated.

MoE train / serve

1. Plot token→expert/rank histograms and stage times — is it GEMM, permute, A2A, or tail?
2. Grouped/persistent GEMM + fused gate/up/activation/combine.
3. DeepEP-class dispatch/combine; FP8; hierarchical topology; chunk overlap.
4. Routing locality, placement/replication, capacity/dropless policy.
5. Only then micro-optimize router/top-k or bet on a full-path megakernel.

VLM / image / video / audio

1. Measure effective tokens vs padding; dynamic resolution, buckets, packing.
2. Async decode/resize/normalize/resample/STFT into target dtype/layout.
3. Varlen attention, conv fusion, local/spatiotemporal sparsity in encoders.
4. Align projector/resampler with LLM layouts; fixed latents buy backend shape stability.
5. Contrastive: chunk similarity/loss and analyze AllGather; generative VLMs reuse LLM loss tricks.

Diffusion / DiT / video generation

1. Break down one step (conv/attn/GEMM, CFG, scheduler, VAE) × step count.
2. BF16/FP8, FA/SDPA, channels-last, compile, fused QKV/norm/act.
3. Prompt/KV cache, CUDA Graph, CFG batching/parallel, scheduler fusion.
4. DeepCache/TeaCache / few-step solvers — report quality separately.
5. Multi-GPU single-sample: DistriFusion / PipeFusion / xDiT; topology and resolution set break-even.

12. Reproducible benchmarking & profiling

From end-to-end to kernel

1. Freeze the workload: model commit, dtype, batch/token budget, length/resolution distribution, train vs serve mode, hardware, topology, software versions, power/clocks.
2. E2E baseline: train → step time / tokens/s / MFU / peak memory; serve → TTFT, TPOT, tokens/s, p50/p95/p99, KV footprint; diffusion → latency + quality metrics.
3. Nsight Systems / rocprof: holes, CPU launch, copies, exposed collectives, stream overlap, tails.
4. Nsight Compute / ROCm Compute Profiler: roofline, TC/MFMA, HBM/L2, occupancy, spills, SMEM/LDS, warp stalls, branches, wave quantization on top kernels.
5. Microbench real shapes from histograms — not only tidy 4096³. Include warmup, JIT/autotune, workspace, layout conversion.
6. Return to E2E after each swap: new copies, graph breaks, comm contention, allocator behavior, or quality loss can erase kernel wins.

Fields worth reporting per kernel

Common benchmark mistakes

• Mixing theoretical FLOPs (including skipped padding/sparse blocks) with executed FLOPs.
• Giving the new kernel prepacked inputs while charging the baseline for conversion.
• Only testing large steady batches when serving is dynamic small-batch + long-tail context.
• Forward-only numbers for training decisions.
• Comparing CUDA-Graph’d launches against a Python/CPU-scheduled baseline.
• Quant GEMM without calibration / dequant / scale / quality / unsupported-shape fallbacks.
• Summing overlapped stream durations instead of critical path / exposed communication.
• Average expert load instead of p99 expert/rank load.
• Diffusion latency without steps, resolution, CFG, seed, scheduler, and quality.
• Numeric checks on one random tensor — miss extremes, NaN/Inf, long sequences, masks, empty experts, tail pages, unaligned shapes.

For an end-to-end PyTorch profiling workflow, see Profiling a PyTorch Training Job End to End.

13. An executable decision tree

If you only have a month of engineering time

• Week 1: Shape/time/bytes dashboard; clear graph breaks, fallbacks, explicit copies; freeze a correctness harness.
• Week 2: Integrate/upgrade mature specialized kernels (FA / TE / Liger / DeepGEMM / Marlin / FlashInfer, …) and unify layouts.
• Week 3: One or two fused kernels on the largest memory-bound path; chunk overlap on the largest communication path.
• Week 4: E2E tuning, long-tail shape/failure fallbacks, numerics/convergence, p99 / multi-node retest; lock a regression matrix.

Do not start by writing a megakernel. That investment pays when profiles show the op boundary itself is the top bottleneck and shapes/deployments are stable enough to amortize the cost.

14. Primary-source index

Attention, KV, serving

• FlashAttention, FlashAttention-2, FlashAttention-3, FlashAttention-4, FA4 author notes
• FlashInfer, FlashInfer sampling
• PagedAttention / vLLM, vAttention
• MQA, GQA, DeepSeek-V2 / MLA, Native Sparse Attention

GEMM, fusion, precision, training kernels

• CUTLASS, Triton tutorials, Stream-K, DeepGEMM
• Transformer Engine, SmoothQuant, AWQ, Marlin
• Liger Kernel, Cut Cross Entropy, CODA
• Apex FusedAdam, 8-bit optimizers, FlashOptim

MoE, distributed systems, communication

• MegaBlocks, ScatterMoE, SonicMoE, TritonMoE
• DeepEP, FLUX, Comet, DeepSeek-V3
• Megatron-Core MoE 0.17, NVIDIA Megatron-Core MoE report, MoE Parallel Folding
• Megatron-LM, PTD-P, ZeRO, ZeRO++, Domino
• Ulysses, Ring Attention, USP, NCCL docs
• PaLM parallel layers, MegaScale PTB, Blockwise Parallel Transformer

Multimodal, diffusion, other architectures

• NaViT, Qwen2-VL, Qwen2.5-VL, SigLIP, VeOmni
• DiT, Stable Diffusion 3, DistriFusion, PipeFusion, xDiT
• DeepCache, TeaCache, SVDQuant, FPSAttention
• Mamba, Mamba-2 / SSD, FlashFFTConv

Hardware / tooling docs cited inline

• Nsight Compute Profiling Guide
• Hopper Tuning Guide
• CUTLASS tcgen05 programming
• ROCm AI inference workload tuning
• Transformer Engine user guide
• PyTorch activation checkpointing techniques

15. Closing judgment

In 2026, foundation-model kernel work is four-layer co-design:

1. On-chip dataflow — tiles, async copy, warp specialization, TMEM/LDS, epilogues;
2. Operator graph — fusion, recompute, layout, paging, persistent scheduling;
3. Multi-GPU dataflow — collective decomposition, tile overlap, topology, expert/token scheduling;
4. The model itself — GQA/MLA, sparse attention, dynamic visual tokens, MoE routing, SSM, diffusion cache / fewer steps.

The strongest recipes usually move two or three layers at once:

• MLA + paged/quantized KV + split-KV decode kernel
• Dropless MoE + grouped GEMM + DeepEP + tile-level overlap
• Dynamic resolution + multimodal packing + varlen FlashAttention
• DiT + FP8 fused GEMM/attention + temporal cache + patch parallel
• Parallel decoder block + TP/SP chunk overlap + fused branch output/collective
• MoE Parallel Folding + GroupedGEMM + Flex dispatcher + shared-expert/A2A overlap

So the right question is not “which kernel is fastest?” It is: under real shapes, real topologies, and real precision constraints, which model expression and data layout make critical bytes cross HBM/network as rarely as possible — and keep the remaining compute saturating the machine?