{"slug": "fast-attention-kernels-and-operator-fusion", "title": "Fast Attention Kernels and Operator Fusion", "summary": "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.", "body_md": "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.\n\nFast 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.\n\nIn this post, you will learn about:\n\n- How operator fusion reduces launches and intermediate memory traffic\n- How fast and paged attention kernels change execution and cache access\n- How CUDA graphs and compilation affect repeated inference steps\n\nLet’s get started.\n\n## Overview\n\nThis post is divided into four parts; they are:\n\n- Operator Fusion\n- FlashAttention-Style Kernels\n- Paged Attention Kernels\n- Fusion and CUDA Graphs with `torch.compile()`\n\n**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.\n\n## Operator Fusion\n\nThe SwiGLU block used by many decoder-only models provides an example:\n\n$$\n\n\\mathbf{y} = \\mathbf{W}_{\\text{down}}\\big(\\operatorname{SiLU}(\\mathbf{W}_{\\text{gate}}\\mathbf{x}) \\odot \\mathbf{W}_{\\text{up}}\\mathbf{x}\\big)\n\n$$\n\nA 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.\n\n``` python\nRun a small gated MLP as an example of operations that can be fused.\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass GatedMLP(nn.Module):\n    def __init__(self, hidden_size, intermediate_size):\n        super().__init__()\n        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)\n\n    def forward(self, x):\n        gate = F.silu(self.gate_proj(x))\n        up = self.up_proj(x)\n        return self.down_proj(gate * up)\n\nmodel = GatedMLP(hidden_size=16, intermediate_size=32).eval()\nprint(model(torch.randn(2, 4, 16)).shape)\n\n1234567891011121314151617181920\n\nimport 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)\n```\n\nThis 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.\n\nFusion 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.\n\nOther common transformer fusion opportunities include:\n\n- Adding a residual connection and applying RMSNorm or LayerNorm\n- Applying bias, activation, and element-wise multiplication\n- Applying rotary position embeddings to queries and keys\n- Updating the KV cache while transforming key and value tensors\n- Combining dequantization with matrix multiplication\n\nFor example, a fused residual and RMS normalization can conceptually perform:\n\n``` python\nCompute a residual update with RMS normalization.\n\nimport torch\n\ndef residual_rms_norm(x, residual, weight, eps=1e-6):\n    x = x + residual\n    variance = x.float().pow(2).mean(dim=-1, keepdim=True)\n    normalized = x * torch.rsqrt(variance + eps).to(x.dtype)\n    return normalized * weight, x\n\nnormalized, updated_residual = residual_rms_norm(\n    torch.randn(2, 4, 16), torch.randn(2, 4, 16), torch.ones(16)\n)\nprint(normalized.shape, updated_residual.shape)\n\n123456789101112\n\nimport 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)\n```\n\nThe 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.\n\nFusion 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.\n\n## FlashAttention-Style Kernels\n\nOrdinary attention with an additive attention mask can be written as:\n\n$$\n\n\\operatorname{Attention}(\\mathbf{Q},\\mathbf{K},\\mathbf{V}) = \\operatorname{softmax}\\big(\\frac{\\mathbf{Q}\\mathbf{K}^\\top}{\\sqrt{d}} + \\mathbf{M}\\big)\\mathbf{V}\n\n$$\n\nA 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.\n\nFlashAttention 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()`:\n\n```\nCall PyTorch scaled dot-product attention on query, key, and value tensors.\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn.attention import SDPBackend, sdpa_kernel\n\nif not torch.cuda.is_available():\n    raise RuntimeError(\"This example requires a CUDA GPU\")\n\ndevice = \"cuda\"\ndtype = torch.float16\n\nq = torch.randn(2, 16, 512, 64, device=device, dtype=dtype)\nk = torch.randn(2, 16, 512, 64, device=device, dtype=dtype)\nv = torch.randn(2, 16, 512, 64, device=device, dtype=dtype)\n\nwith torch.no_grad(), sdpa_kernel(\n    [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION], set_priority=True\n):\n    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)\n\nprint(output.shape)\n\n1234567891011121314151617181920\n\nimport 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)\n```\n\nWithout 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.\n\nFlashAttention 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.\n\n## Paged Attention Kernels\n\nA 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.\n\nA 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:\n\n| Logical blocks: | 0 | 1 | 2 | 3 | \n| Physical blocks: | 17 | 4 | 29 | 8 | \n\nThe 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.\n\nPaged 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.\n\nPyTorch 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:\n\n```\nGenerate text with vLLM using a selected KV-cache block size.\n\nfrom vllm import LLM, SamplingParams\n\nattention_backend = \"TRITON_ATTN\"\n\n# vLLM stores this decoder model's KV cache in 16-token blocks.\nllm = LLM(\n    model=\"Qwen/Qwen2.5-1.5B-Instruct\",\n    dtype=\"float16\",\n    block_size=16,\n    attention_backend=attention_backend,\n)\nprint(\"Attention backend:\", attention_backend)\n\n# Set up the sampling parameters.\nsampling_params = SamplingParams(\n    temperature=0.7,\n    top_p=0.95,\n    max_tokens=64,\n)\n\n# Make a request.\nprompts = [\n    \"Q: What is the capital of France?\\nA:\",\n    \"Q: Who wrote 'Pride and Prejudice'?\\nA:\",\n]\n\noutputs = llm.generate(prompts, sampling_params)\n\nfor output in outputs:\n    print(f\"Prompt: {output.prompt!r}\")\n    print(f\"Output: {output.outputs[0].text}\\n\")\n\n12345678910111213141516171819202122232425262728293031\n\nfrom 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\")\n```\n\nThis 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.\n\n## Fusion and CUDA Graphs with `torch.compile`\n\nHand-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.\n\nTherefore, `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.\n\nThe 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:\n\n``` python\nCompare compiled and eager output from the gated MLP.\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass GatedMLP(nn.Module):\n    def __init__(self, hidden_size, intermediate_size):\n        super().__init__()\n        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)\n\n    def forward(self, x):\n        gate = F.silu(self.gate_proj(x))\n        up = self.up_proj(x)\n        return self.down_proj(gate * up)\n\nif not torch.cuda.is_available():\n    raise RuntimeError(\"This example requires a CUDA GPU\")\n\nmodel = GatedMLP(\n    hidden_size=1024,\n    intermediate_size=4096,\n).eval().to(device=\"cuda\", dtype=torch.float16)\n\nx = torch.randn(1, 128, 1024, device=\"cuda\", dtype=torch.float16)\n\ncompiled_model = torch.compile(\n    model,\n    backend=\"inductor\",\n    mode=\"reduce-overhead\",\n)\n\nwith torch.inference_mode():\n    expected = model(x)\n\n    # The first calls may compile, warm up, and capture eligible CUDA work.\n    for _ in range(3):\n        actual = compiled_model(x)\n\ntorch.cuda.synchronize()\ntorch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)\nprint(\"compiled output matches eager output:\", tuple(actual.shape))\n\n1234567891011121314151617181920212223242526272829303132333435363738394041424344\n\nimport 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))\n```\n\nThe 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()`.\n\nFusion 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.\n\nCUDA 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.\n\nSet `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.\n\nRerun 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.\n\nThe 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.\n\nEager 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.\n\n## Further Reading\n\nBelow are some resources you may find useful:\n\n- 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.\n- Tri Dao (2023). [FlashAttention-2](https://arxiv.org/abs/2307.08691) . This paper improves work partitioning and parallelism for attention kernels.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.\n\n## Summary\n\nIn 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.\n\nThese techniques depend on shapes, layouts, dtypes, and hardware support.", "url": "https://wpnews.pro/news/fast-attention-kernels-and-operator-fusion", "canonical_source": "https://machinelearningmastery.com/fast-attention-kernels-and-operator-fusion/", "published_at": "2026-11-12 17:01:10+00:00", "updated_at": "2026-11-12 20:22:42.225295+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "ai-infrastructure", "ai-research", "developer-tools"], "entities": ["vLLM", "PyTorch", "torch.compile()", "FlashAttention", "SwiGLU", "CUDA"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/fast-attention-kernels-and-operator-fusion", "markdown": "https://wpnews.pro/news/fast-attention-kernels-and-operator-fusion.md", "text": "https://wpnews.pro/news/fast-attention-kernels-and-operator-fusion.txt", "jsonld": "https://wpnews.pro/news/fast-attention-kernels-and-operator-fusion.jsonld"}}