Fast Attention Kernels and Operator Fusion

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. Fast Attention Kernels and Operator Fusion.Photo by Chris Boland. Some rights reserved. 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. 1234567891011121314151617181920 import torchimport torch.nn as nnimport torch.nn.functional as Fclass 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: import torchdef 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, xnormalized, 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(): 1234567891011121314151617181920 import torchimport torch.nn.functional as Ffrom torch.nn.attention import SDPBackend, sdpa_kernelif not torch.cuda.is_available(): raise RuntimeError("This example requires a CUDA GPU")device = "cuda"dtype = torch.float16q = 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: 12345678910111213141516171819202122232425262728293031 from vllm import LLM, SamplingParamsattention_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: 1234567891011121314151617181920212223242526272829303132333435363738394041424344 import torchimport torch.nn as nnimport torch.nn.functional as Fclass 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. This paper introduces an IO-aware tiled algorithm for exact attention. Tri Dao (2023). FlashAttention-2. 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. This paper describes paged KV cache management and the vLLM serving system. PyTorch 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. This page explains automatic backend selection and the backends available on each hardware platform. PyTorch documentation. Scaled dot product attention. This page documents the unified PyTorch attention API and its backend behavior. PyTorch documentation. torch.compile. This page documents compilation modes, including reduce-overhead and its use of CUDA Graphs. PyTorch 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. 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. No comments yet.