Fast Attention Kernels and Operator Fusion

A technical post explains how fast attention kernels and operator fusion reduce kernel launches and intermediate memory traffic in transformer inference, covering operator fusion, FlashAttention-style kernels, paged attention kernels, and CUDA graphs with torch.compile(). The post uses a SwiGLU gated MLP example with three matrix multiplications, one SiLU activation and one pointwise multiplication to show where a direct implementation writes intermediates to high-bandwidth memory, and notes the optimizations leave model weights and the mathematical function unchanged, though small numerical differences can arise from operation order and precision. The code depends on vLLM, installable via python -m pip install vllm, and requires a CUDA-enabled PyTorch build matching the driver so torch.compile() can use CUDA graphs.

A transformer layer is built from familiar operations: linear projections, attention, normalization, activation functions, and residual additions. A direct implementation may launch a separate accelerator kernel for every operation and write every intermediate tensor to device memory. The mathematics is correct, but the execution can spend more time moving values and launching kernels than performing useful arithmetic. Fast attention kernels and operator fusion reorganize this work. They keep intermediate values close to the arithmetic units, combine operations, and reuse compiled execution plans. These optimizations normally leave the model weights and mathematical function unchanged, although small numerical differences can arise from operation order and precision. In this post, you will learn about: - How operator fusion reduces launches and intermediate memory traffic - How fast and paged attention kernels change execution and cache access - How CUDA graphs and compilation affect repeated inference steps Let’s get started. Overview This post is divided into four parts; they are: - Operator Fusion - FlashAttention-Style Kernels - Paged Attention Kernels - Fusion and CUDA Graphs with torch.compile Note: Starting with this post, you will run the code that depends on vLLM. You can install vLLM with python -m pip install vllm . This environment should have a CUDA-enabled PyTorch build that matches your driver so that torch.compile can use CUDA graphs. Also note that running the code may download model checkpoints and other data. Be sure you have enough disk space. Operator Fusion The SwiGLU block used by many decoder-only models provides an example: $$ \mathbf{y} = \mathbf{W} {\text{down}}\big \operatorname{SiLU} \mathbf{W} {\text{gate}}\mathbf{x} \odot \mathbf{W} {\text{up}}\mathbf{x}\big $$ A direct implementation launches matrix multiplications for the projections, a kernel for SiLU, another for element-wise multiplication, and a final matrix multiplication. Each boundary may write an intermediate tensor to high-bandwidth memory and read it back. python Run a small gated MLP as an example of operations that can be fused. import torch import torch.nn as nn import torch.nn.functional as F class GatedMLP nn.Module : def init self, hidden size, intermediate size : super . init self.gate proj = nn.Linear hidden size, intermediate size, bias=False self.up proj = nn.Linear hidden size, intermediate size, bias=False self.down proj = nn.Linear intermediate size, hidden size, bias=False def forward self, x : gate = F.silu self.gate proj x up = self.up proj x return self.down proj gate up model = GatedMLP hidden size=16, intermediate size=32 .eval print model torch.randn 2, 4, 16 .shape 1234567891011121314151617181920 import torchimport torch.nn as nnimport torch.nn.functional as F class GatedMLP nn.Module : def init self, hidden size, intermediate size : super . init self.gate proj = nn.Linear hidden size, intermediate size, bias=False self.up proj = nn.Linear hidden size, intermediate size, bias=False self.down proj = nn.Linear intermediate size, hidden size, bias=False def forward self, x : gate = F.silu self.gate proj x up = self.up proj x return self.down proj gate up model = GatedMLP hidden size=16, intermediate size=32 .eval print model torch.randn 2, 4, 16 .shape This code expresses the model clearly and should usually remain the reference implementation. It contains three matrix multiplications, one pointwise SiLU activation, and one pointwise multiplication. An optimized backend may fuse SiLU with multiplication, combine projection work, or choose a layout that avoids extra transposes and kernel launches. Fusion saves time when intermediate memory traffic and launch overhead are significant. It helps less when a large matrix multiplication already dominates execution or when the fused kernel uses an inefficient shape. Kernel choice is empirical: profile the complete workload. Other common transformer fusion opportunities include: - Adding a residual connection and applying RMSNorm or LayerNorm - Applying bias, activation, and element-wise multiplication - Applying rotary position embeddings to queries and keys - Updating the KV cache while transforming key and value tensors - Combining dequantization with matrix multiplication For example, a fused residual and RMS normalization can conceptually perform: python Compute a residual update with RMS normalization. import torch def residual rms norm x, residual, weight, eps=1e-6 : x = x + residual variance = x.float .pow 2 .mean dim=-1, keepdim=True normalized = x torch.rsqrt variance + eps .to x.dtype return normalized weight, x normalized, updated residual = residual rms norm torch.randn 2, 4, 16 , torch.randn 2, 4, 16 , torch.ones 16 print normalized.shape, updated residual.shape 123456789101112 import torch def residual rms norm x, residual, weight, eps=1e-6 : x = x + residual variance = x.float .pow 2 .mean dim=-1, keepdim=True normalized = x torch.rsqrt variance + eps .to x.dtype return normalized weight, x normalized, updated residual = residual rms norm torch.randn 2, 4, 16 , torch.randn 2, 4, 16 , torch.ones 16 print normalized.shape, updated residual.shape The function returns both the normalized value and the updated residual. A fused kernel can read each input once, perform the reduction and scaling, and write the two required outputs. Separate PyTorch operations may create more intermediate traffic. Fusion has boundaries. Operations with incompatible layouts, synchronization requirements, or very large intermediate state may not combine efficiently. A giant fused kernel can also increase register pressure, reduce occupancy, and become slower than several well-tuned kernels. FlashAttention-Style Kernels Ordinary attention with an additive attention mask can be written as: $$ \operatorname{Attention} \mathbf{Q},\mathbf{K},\mathbf{V} = \operatorname{softmax}\big \frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}} + \mathbf{M}\big \mathbf{V} $$ A simple implementation materializes the attention-score matrix $\mathbf{Q}\mathbf{K}^\top$. For a sequence length $T$, that matrix contains $T^2$ elements per head. On a GPU, computing the matrix, writing it to global memory, and reading it back for softmax consumes substantial bandwidth between the streaming multiprocessors SMs and device memory, typically high-bandwidth memory HBM . This memory traffic can become a bottleneck during long prefills. FlashAttention processes queries, keys, and values in tiles. For each tile, it computes a block of scores, updates the softmax statistics, and accumulates the corresponding output using on-chip memory. It avoids storing the complete score matrix in global memory. PyTorch exposes optimized attention through scaled dot product attention : Call PyTorch scaled dot-product attention on query, key, and value tensors. import torch import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa kernel if not torch.cuda.is available : raise RuntimeError "This example requires a CUDA GPU" device = "cuda" dtype = torch.float16 q = torch.randn 2, 16, 512, 64, device=device, dtype=dtype k = torch.randn 2, 16, 512, 64, device=device, dtype=dtype v = torch.randn 2, 16, 512, 64, device=device, dtype=dtype with torch.no grad , sdpa kernel SDPBackend.FLASH ATTENTION, SDPBackend.CUDNN ATTENTION , set priority=True : output = F.scaled dot product attention q, k, v, is causal=True print output.shape 1234567891011121314151617181920 import torchimport torch.nn.functional as Ffrom torch.nn.attention import SDPBackend, sdpa kernel if not torch.cuda.is available : raise RuntimeError "This example requires a CUDA GPU" device = "cuda"dtype = torch.float16 q = torch.randn 2, 16, 512, 64, device=device, dtype=dtype k = torch.randn 2, 16, 512, 64, device=device, dtype=dtype v = torch.randn 2, 16, 512, 64, device=device, dtype=dtype with torch.no grad , sdpa kernel SDPBackend.FLASH ATTENTION, SDPBackend.CUDNN ATTENTION , set priority=True : output = F.scaled dot product attention q, k, v, is causal=True print output.shape Without the sdpa kernel context, PyTorch selects a backend based on the device, dtype, tensor layout, head dimension, and mask. The example explicitly prefers FlashAttention and then cuDNN attention. Hard-coding these backends can cause the code to fail when neither backend supports the inputs or hardware. FlashAttention is especially valuable for prefill because many query tokens are processed together. Decode has only one or a few new queries per request. It still benefits from specialized kernels, but reading the growing KV cache often becomes the larger concern. Paged Attention Kernels A request’s KV cache grows one token at a time and ends at an unpredictable length. Requiring one contiguous allocation per request leads to reserved but unused space and makes moving or resizing caches expensive. A paged KV cache divides storage into fixed-size blocks and separates the logical and physical representations of the cache. A logical sequence is represented by a block table: | Logical blocks: | 0 | 1 | 2 | 3 | | Physical blocks: | 17 | 4 | 29 | 8 | The attention kernel follows the table while reading keys and values. Logical token order remains contiguous even though physical memory is not. Blocks can be allocated on demand, returned to a pool when a request ends, and shared when sequences have a common prefix. Paged attention therefore combines a memory-management design with a specialized kernel. The kernel cannot assume that K t+1 is adjacent to K t ; it must translate logical positions through the block table. Paging does not reduce the number of logical KV elements in a sequence. It reduces fragmentation and reserved capacity. PyTorch can implement paged attention with FlexAttention and torch.compile . For an integrated serving example, vLLM allocates the KV cache in blocks and supplies a block table to its attention backend. The following example explicitly selects a cache block size of 16 tokens: Generate text with vLLM using a selected KV-cache block size. from vllm import LLM, SamplingParams attention backend = "TRITON ATTN" vLLM stores this decoder model's KV cache in 16-token blocks. llm = LLM model="Qwen/Qwen2.5-1.5B-Instruct", dtype="float16", block size=16, attention backend=attention backend, print "Attention backend:", attention backend Set up the sampling parameters. sampling params = SamplingParams temperature=0.7, top p=0.95, max tokens=64, Make a request. prompts = "Q: What is the capital of France?\nA:", "Q: Who wrote 'Pride and Prejudice'?\nA:", outputs = llm.generate prompts, sampling params for output in outputs: print f"Prompt: {output.prompt r}" print f"Output: {output.outputs 0 .text}\n" 12345678910111213141516171819202122232425262728293031 from vllm import LLM, SamplingParams attention backend = "TRITON ATTN" vLLM stores this decoder model's KV cache in 16-token blocks.llm = LLM model="Qwen/Qwen2.5-1.5B-Instruct", dtype="float16", block size=16, attention backend=attention backend, print "Attention backend:", attention backend Set up the sampling parameters.sampling params = SamplingParams temperature=0.7, top p=0.95, max tokens=64, Make a request.prompts = "Q: What is the capital of France?\nA:", "Q: Who wrote 'Pride and Prejudice'?\nA:", outputs = llm.generate prompts, sampling params for output in outputs: print f"Prompt: {output.prompt r}" print f"Output: {output.outputs 0 .text}\n" This example uses paged KV-cache management, but block size alone does not select a compute backend. The attention backend argument explicitly requests TRITON ATTN , whose forward path implements paged attention over the block-managed KV cache. During initialization, vLLM validates the backend against the model and hardware and raises an error if they are incompatible. If initialization succeeds, the printed value confirms that the paged-attention backend is in use. Fusion and CUDA Graphs with torch.compile Hand-written CUDA or Triton kernels provide control but require specialist work. PyTorch can instead capture a sequence of operations with TorchDynamo and compile the captured graph with TorchInductor, the default backend of torch.compile . TorchInductor performs optimizations such as graph rewriting, buffer reuse, and operator fusion. On GPUs, it can generate a single Triton kernel for compatible pointwise or reduction operations that eager PyTorch would launch separately. Therefore, torch.compile does perform fusion, but it does not fuse every operation. Fusion depends on the operations, tensor layouts, shapes, device, and compiler backend. Graph breaks also limit the regions that the compiler can optimize. Python-side data-dependent control flow, unsupported operations, and conversions from tensors to Python scalars can split one model into several compiled graphs. The following example compiles the GatedMLP used earlier. It includes the class definition again so the script runs on its own. The inductor backend enables compiler optimizations, including fusion. The reduce-overhead mode asks TorchInductor to use CUDA Graphs when the captured graph is eligible: python Compare compiled and eager output from the gated MLP. import torch import torch.nn as nn import torch.nn.functional as F class GatedMLP nn.Module : def init self, hidden size, intermediate size : super . init self.gate proj = nn.Linear hidden size, intermediate size, bias=False self.up proj = nn.Linear hidden size, intermediate size, bias=False self.down proj = nn.Linear intermediate size, hidden size, bias=False def forward self, x : gate = F.silu self.gate proj x up = self.up proj x return self.down proj gate up if not torch.cuda.is available : raise RuntimeError "This example requires a CUDA GPU" model = GatedMLP hidden size=1024, intermediate size=4096, .eval .to device="cuda", dtype=torch.float16 x = torch.randn 1, 128, 1024, device="cuda", dtype=torch.float16 compiled model = torch.compile model, backend="inductor", mode="reduce-overhead", with torch.inference mode : expected = model x The first calls may compile, warm up, and capture eligible CUDA work. for in range 3 : actual = compiled model x torch.cuda.synchronize torch.testing.assert close actual, expected, rtol=1e-2, atol=1e-2 print "compiled output matches eager output:", tuple actual.shape 1234567891011121314151617181920212223242526272829303132333435363738394041424344 import torchimport torch.nn as nnimport torch.nn.functional as F class GatedMLP nn.Module : def init self, hidden size, intermediate size : super . init self.gate proj = nn.Linear hidden size, intermediate size, bias=False self.up proj = nn.Linear hidden size, intermediate size, bias=False self.down proj = nn.Linear intermediate size, hidden size, bias=False def forward self, x : gate = F.silu self.gate proj x up = self.up proj x return self.down proj gate up if not torch.cuda.is available : raise RuntimeError "This example requires a CUDA GPU" model = GatedMLP hidden size=1024, intermediate size=4096, .eval .to device="cuda", dtype=torch.float16 x = torch.randn 1, 128, 1024, device="cuda", dtype=torch.float16 compiled model = torch.compile model, backend="inductor", mode="reduce-overhead", with torch.inference mode : expected = model x The first calls may compile, warm up, and capture eligible CUDA work. for in range 3 : actual = compiled model x torch.cuda.synchronize torch.testing.assert close actual, expected, rtol=1e-2, atol=1e-2 print "compiled output matches eager output:", tuple actual.shape The code uses the torch.inference mode context manager to disable gradient computation as well as view tracking and version-counter updates. It is more restrictive than torch.no grad . Fusion and CUDA Graphs remove different overheads. Fusion combines compatible operations into fewer kernels, reducing intermediate memory traffic and the number of launches. A CUDA Graph records the remaining sequence of GPU operations and replays it with one CPU submission, reducing Python, C++, driver, and kernel-launch overhead. CUDA Graphs do not create additional operator fusion; they complement the fused kernels produced by TorchInductor. CUDA Graph replay requires stable memory addresses and compatible shapes. It is most useful for repeated, small-batch workloads such as decode. It is less suitable when requests continually change shape or use data-dependent control flow. The reduce-overhead mode is not a guarantee: PyTorch uses CUDA Graphs only for eligible CUDA graphs, and it may consume more memory by retaining workspace allocations for replay. Set TORCH LOGS=perf hints before running the compilation block in your Python environment to investigate performance. PyTorch will log hints about CUDA Graph eligibility and other compiler choices. Rerun the same block with TORCH LOGS=output code to inspect TorchInductor’s generated output. These environment variables must be set before starting Python; setting them after importing PyTorch may not enable the requested logs. The generated code may contain Triton kernel names that list several fused operations. You can also use torch.profiler to compare eager and compiled execution. Fewer kernels provide evidence of fusion, but the final test is end-to-end performance. Eager and compiled execution may also produce slightly different results because compilation can change the order of floating-point operations. You should verify results with tolerances appropriate to the dtype. Dynamic tensor shapes may trigger recompilation and unexpectedly reduce fusion. Therefore, check the profiler trace, numerical output, and application-level performance rather than merely confirming that compilation completed. Further Reading Below are some resources you may find useful: - Tri Dao et al. 2022 . FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness https://arxiv.org/abs/2205.14135 . This paper introduces an IO-aware tiled algorithm for exact attention. - Tri Dao 2023 . FlashAttention-2 https://arxiv.org/abs/2307.08691 . This paper improves work partitioning and parallelism for attention kernels. - Woosuk Kwon et al. 2023 . Efficient Memory Management for Large Language Model Serving with PagedAttention https://arxiv.org/abs/2309.06180 . This paper describes paged KV cache management and the vLLM serving system. - PyTorch blog. FlexAttention for Inference https://pytorch.org/blog/flexattention-for-inference/ . This article demonstrates a PyTorch-native paged-attention implementation built with FlexAttention and torch.compile. - vLLM documentation. Attention Backend Feature Support https://docs.vllm.ai/en/latest/design/attention backends/ . This page explains automatic backend selection and the backends available on each hardware platform. - PyTorch documentation. Scaled dot product attention https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled dot product attention.html . This page documents the unified PyTorch attention API and its backend behavior. - PyTorch documentation. torch.compile https://docs.pytorch.org/docs/stable/generated/torch.compile.html . This page documents compilation modes, including reduce-overhead and its use of CUDA Graphs. - PyTorch blog. Why Is PyTorch Compile So Fast? Kernel Fusion https://pytorch.org/blog/why-is-pytorch-compile-so-fast-kernel-fusion/ . This article shows how TorchInductor combines compatible operations into generated Triton kernels. - PyTorch documentation. CUDA semantics https://docs.pytorch.org/docs/stable/notes/cuda.html cuda-graphs . This section explains CUDA Graph capture, replay, memory requirements, and constraints. Summary In this post, you learned that fast kernels improve inference by reducing memory traffic, intermediate tensors, and launch overhead. FlashAttention tiles exact attention so the complete score matrix is not stored in high-bandwidth memory. Paged attention follows block tables so KV cache memory can be allocated flexibly. The function torch.compile uses TorchInductor to optimize captured graphs and fuse compatible operations, while CUDA Graphs complement fusion by reducing the launch overhead of replaying the remaining kernels. These techniques depend on shapes, layouts, dtypes, and hardware support.