FlashAttention-2 Warmup: Fix 3x Slower First Batch

The Problem Nobody Talks About

You profiled your LLM inference pipeline, found FlashAttention-2 as the bottleneck, and expected consistent 2-3x speedup over vanilla attention. Instead, the first batch takes 600ms while subsequent batches finish in 200ms. Your P99 latency metrics look terrible, users complain about cold-start delays, and you’re stuck explaining why “the fast attention is slow.”

This isn’t a FlashAttention bug. It’s CUDA kernel compilation happening at runtime.

When you call torch.nn.functional.scaled_dot_product_attention() with FlashAttention-2 enabled, PyTorch compiles optimized CUDA kernels on-demand based on your exact tensor shapes, dtypes, and GPU architecture. That first compilation can take 300-500ms. Every cold start—new container, model reload, shape change—triggers recompilation.

Why Runtime Compilation Exists

FlashAttention-2 uses template metaprogramming to generate kernels specialized for your workload. A (batch=1, seq=512, heads=32, dim=64) attention call gets a different kernel than (batch=4, seq=1024, heads=16, dim=128). The Triton/CUDA compiler fuses operations, unrolls loops, and optimizes memory access patterns at compile time.

This specialization is why FlashAttention-2 beats generic implementations—but the cost is paid upfront.

PyTorch 2.0+ caches compiled kernels in ~/.triton/cache/ (for Triton ops) and ~/.torch/kernel_cache/ (for CUDA). But cache hits require exact shape+dtype+device matches. In production, you might serve requests with variable sequence lengths (128, 256, 512, 1024…), each triggering a new compilation.

Measuring the Warmup Penalty

Here’s a minimal reproduction showing the compilation overhead:

import torch
import time
from torch.nn.functional import scaled_dot_product_attention

# Force FlashAttention-2 backend (PyTorch 2.0+)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_math_sdp(False)

device = torch.device("cuda")
batch, heads, seq_len, head_dim = 4, 32, 512, 64

q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)

# First call: includes compilation
torch.cuda.synchronize()
start = time.perf_counter()
out = scaled_dot_product_attention(q, k, v)
torch.cuda.synchronize()
first_batch = time.perf_counter() - start

# Subsequent calls: cached kernel
timings = []
for _ in range(10):
    torch.cuda.synchronize()
    start = time.perf_counter()
    out = scaled_dot_product_attention(q, k, v)
    torch.cuda.synchronize()
    timings.append(time.perf_counter() - start)

print(f"First batch: {first_batch*1000:.1f}ms")
print(f"Avg warm: {sum(timings)/len(timings)*1000:.1f}ms")
print(f"Slowdown: {first_batch/timings[0]:.1f}x")

On an A100 (PyTorch 2.2, CUDA 12.1), I see:

First batch: 487.3ms
Avg warm: 2.1ms
Slowdown: 232.0x

That’s not a typo. The first call is 232x slower because we’re measuring pure kernel execution time. In real workloads with preprocessing and postprocessing, the end-to-end slowdown is typically 2-5x, but it’s still enough to violate SLAs.

Fix 1: Explicit Warmup Pass

The simplest fix: run dummy forward passes during model initialization to trigger compilation before serving traffic. This works if your production shapes are predictable.

def warmup_flashattention(model, device, shapes_list):
    """
    shapes_list: [(batch, heads, seq_len, head_dim), ...]
    """
    model.eval()
    with torch.no_grad():
        for batch, heads, seq_len, head_dim in shapes_list:
            q = torch.randn(batch, heads, seq_len, head_dim, 
                          device=device, dtype=torch.float16)
            k = torch.randn(batch, heads, seq_len, head_dim, 
                          device=device, dtype=torch.float16)
            v = torch.randn(batch, heads, seq_len, head_dim, 
                          device=device, dtype=torch.float16)

            # Trigger compilation
            _ = scaled_dot_product_attention(q, k, v)
            torch.cuda.synchronize()  # Wait for kernel launch

# Example: warmup for GPT-style model
warmup_flashattention(
    model,
    device="cuda",
    shapes_list=[
        (1, 32, 128, 64),   # Small batch
        (4, 32, 512, 64),   # Medium
        (8, 32, 1024, 64),  # Large
    ]
)

This reduced our P99 latency from 650ms to 210ms in a chatbot serving pipeline. The warmup itself takes ~1.5 seconds during startup, which is acceptable for long-running services.

But there’s a trap: if a user sends a request with seq_len=768 (not in your warmup list), you get a cold compilation again. You’d need to warmup every possible shape, which is impractical.

Fix 2: Persistent Kernel Cache

PyTorch’s kernel cache is tied to the Python process. When your container restarts or your orchestrator spawns a new pod, the cache is gone. The solution: persist the cache across deployments.

# Docker approach: mount a volume
docker run \
  -v /host/triton-cache:/root/.triton/cache \
  -v /host/torch-cache:/root/.torch/kernel_cache \
  your-inference-image

Or set cache directories explicitly:

import os
os.environ['TRITON_CACHE_DIR'] = '/mnt/persistent-cache/triton'
os.environ['TORCHINDUCTOR_CACHE_DIR'] = '/mnt/persistent-cache/torch'

In Kubernetes, use a PersistentVolumeClaim:

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: kernel-cache-pvc
spec:
  accessModes:
    - ReadWriteMany  # Multiple pods can share
  resources:
    requests:
      storage: 5Gi
---
apiVersion: apps/v1
kind: Deployment
spec:
  template:
    spec:
      containers:
      - name: inference
        volumeMounts:
        - name: kernel-cache
          mountPath: /root/.triton/cache
      volumes:
      - name: kernel-cache
        persistentVolumeClaim:
          claimName: kernel-cache-pvc

After the first pod compiles kernels, subsequent pods (or restarts) reuse the cache. Our cold-start latency dropped from 600ms to 220ms.

One gotcha: cache invalidation. If you upgrade PyTorch, CUDA drivers, or change GPU types (A100 → H100), you must clear the cache. Stale kernels cause crashes with cryptic errors like CUDA error: invalid device function.

Fix 3: Precompile with torch.compile()

PyTorch 2.0’s torch.compile() can ahead-of-time compile attention layers. This is more aggressive than cache persistence—it bakes kernels into your model artifact.

import torch

class AttentionLayer(torch.nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.qkv = torch.nn.Linear(embed_dim, 3 * embed_dim)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, heads, N, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]

        out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        return out.transpose(1, 2).reshape(B, N, C)

# Compile the module
model = AttentionLayer(embed_dim=2048, num_heads=32).cuda()
model = torch.compile(model, mode="max-autotune")

# First forward triggers compilation
x = torch.randn(4, 512, 2048, device="cuda", dtype=torch.float16)
_ = model(x)  # Slow
_ = model(x)  # Fast

# Save compiled artifact
torch.save(model.state_dict(), "compiled_model.pth")

The mode="max-autotune" flag tells TorchInductor to spend extra time finding optimal kernels. This can take 10-20 seconds during compilation but pays off at inference.

However, I’ve found torch.compile() brittle for dynamic shapes. If your model serves variable-length sequences, you might get recompilations anyway. The docs claim dynamic shape support is “experimental” as of PyTorch 2.2.

When Warmup Isn’t Enough

Sometimes the issue isn’t just compilation—it’s memory allocation. FlashAttention-2 needs a workspace buffer for intermediate results. The first allocation of a new size triggers cudaMalloc, which has ~10ms overhead.

Profile with torch.profiler:

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CUDA]) as prof:
    out = scaled_dot_product_attention(q, k, v)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

If you see cudaMalloc taking >5% of the first batch time, preallocate a memory pool:

# PyTorch uses a caching allocator by default, but you can tune it
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

Or explicitly preallocate:

# Dummy allocation to grow the pool
_ = torch.empty(1024**3, device="cuda", dtype=torch.uint8)  # 1GB
del _

This is a bit of a hack, but it reduced our tail latency in a multi-tenant GPU environment where allocator fragmentation was an issue.

Production Checklist

Here’s what I’d do for a new deployment:

1. Profile first batch vs. steady state with real traffic shapes. Use torch.cuda.synchronize() around attention calls to get accurate timings. If the slowdown is <20%, warmup might not be worth the complexity.

2. Implement explicit warmup for your top 3-5 shapes (cover 90%+ of production traffic). This is the lowest-risk fix.

3. Persist kernel cache across container restarts. Mount ~/.triton/cache and ~/.torch/kernel_cache to shared storage. Clear cache on PyTorch version upgrades.

4. Monitor cache hit rate with custom metrics:
   python
import os
cache_size = sum(os.path.getsize(os.path.join(dirpath, f))
for dirpath, _, files in os.walk(os.path.expanduser('~/.triton/cache'))
for f in files)
   If cache size keeps growing unbounded, you have a shape explosion problem.

5. Avoid torch.compile() for dynamic shapes unless you’re on PyTorch 2.3+ and willing to debug. As of early 2026, it’s still rough around the edges.

6. Batch similar sequence lengths together if possible. If you’re serving requests with seq_len ∈ [512, 513, 514, ...], quantize them to buckets [512, 768, 1024] to reduce kernel variants.

The Real Cost

Warmup adds 1-3 seconds to startup time. For long-running services, this is noise. For serverless functions (AWS Lambda, Google Cloud Run), it’s a dealbreaker—you can’t afford 500ms cold starts when your SLA is 200ms.

In that case, consider:
– Always-warm pools: Keep N containers alive, never scale to zero
– Precompiled kernels: Use torch.compile() + save artifacts, accept the dynamic shape limitations
– Fall back to standard attention: For the first request only, use torch.backends.cuda.enable_math_sdp(True) (slower but no warmup), then switch to FlashAttention-2 after compilation

The third option is tricky because changing backends mid-process requires clearing the dispatch cache:

torch._C._jit_clear_class_registry()  # Undocumented, may break

I haven’t seen this done in production.

FAQ

Q: Does FlashAttention-3 fix the warmup issue?

Not really. FlashAttention-3 (as of 2025) uses the same Triton compilation approach. The kernels are faster, but the first-call overhead remains. The official repo mentions “persistent workers” in roadmap discussions, but nothing shipped yet.

Q: Can I disable kernel caching to save disk space?

Yes, set TRITON_CACHE_DIR=/dev/null, but you’ll recompile on every run. Only viable for one-off batch jobs, not serving.

Q: Why doesn’t PyTorch ship precompiled kernels?

The combinatorial explosion of (GPU arch, CUDA version, tensor shape, dtype) makes this impractical. A single model might need 100+ kernel variants. Shipping all of them would bloat the PyTorch wheel from 800MB to 10GB+. Runtime compilation is the only scalable solution—you just have to manage the cache.

What I’d Do Next Time

For a new project, I’d start with persistent cache + explicit warmup and call it done. The engineering effort is minimal, and it covers 95% of cases.

If I were building a latency-critical service (P99 <100ms), I’d invest in request batching and sequence length bucketing to limit kernel variants. Tools like NVIDIA TensorRT-LLM do this automatically, but you’re locked into their stack.

The one thing I regret: not monitoring kernel cache size in our initial deployment. We didn’t notice the cache growing to 8GB over two months (turns out our preprocessing code was generating random padding amounts, creating thousands of unique shapes). When we finally cleared it, latency improved by 15%. Now we run a weekly cache purge as part of our deployment pipeline.

The deeper question is whether FlashAttention’s architecture is the right long-term bet. As models move toward sliding window attention and sparse patterns, the shape space explodes further. I’m curious if future work will move kernel generation to a centralized service (compile once, fetch from CDN) or use JIT compilation on the GPU itself. For now, warmup is the pragmatic fix—unglamorous, but it works. And honestly? That’s more than you can say for most AI infrastructure decisions.