Skip to content

DESIGN.md

Latest commit

 

History

History
100 lines (72 loc) · 3.37 KB

DESIGN.md

File metadata and controls

100 lines (72 loc) · 3.37 KB
Raw

Design Document -- attention-kernel-profiler

1. Motivation

Attention is the central computation of transformer-based LLMs. Its cost scales quadratically with sequence length in the naive implementation, but modern kernels (FlashAttention, mem_efficient_attention) reduce this to near-linear by avoiding materializing the full n x n attention matrix.

This project measures how much difference the kernel choice makes on real hardware.

2. Hardware Context

The RTX 2070 is a Turing architecture GPU (sm75). FlashAttention requires sm80+ (Ampere, e.g., A100, RTX 3090). PyTorch's SDPA automatically falls back to mem_efficient_attention on sm75. This project documents the performance of mem_efficient_attention as the practical FlashAttention substitute on consumer-grade hardware.

3. Attention Variants

naive_math

Manual implementation: scores = Q @ K^T / sqrt(d_head) scores = causal_mask(scores) attn = softmax(scores) out = attn @ V

Complexity: O(n^2) time, O(n^2) memory (materializes full attention matrix).

sdpa_default

torch.nn.functional.scaled_dot_product_attention with auto backend selection. On sm75: uses mem_efficient_attention. On sm80+: uses FlashAttention.

sdpa_flash

Forces FlashAttention backend. Falls back to sdpa_default on sm75.

sdpa_math

Forces the quadratic math backend (no tiled attention). Used as a baseline to measure the benefit of kernel optimization.

sdpa_memeff

Forces memory-efficient attention backend explicitly.

sliding_w128

Sliding window attention: tokens outside a window of 128 are masked. Current implementation allocates full n x n matrix and masks. A true sliding window kernel would compute only the window.

gqa_4kv

Grouped Query Attention: 12 query heads, 3 KV heads (group size = 4). Expands K/V via repeat_interleave before SDPA.

4. Key Results

Time (us) at batch_size=1

Variant       seq=256  seq=512  seq=1024  seq=2048
naive           288      549      1264      4947
sdpa_default     96      120       255       647
sdpa_math       489     1186      3766      8483
sdpa_memeff     115      257       228       707
gqa_4kv         147      175       332       945

Speedup vs naive at seq=2048

sdpa_default:  7.64x
sdpa_memeff:   7.00x
sdpa_flash:    7.02x
gqa_4kv:       5.24x
sdpa_math:     0.58x (worse than naive!)

Memory scaling (batch_size=1, seq 64->2048, 32x longer)

naive:         8.7 -> 217 MB  (25x growth, O(n^0.93))
sdpa_default:  8.5 ->  20 MB  (2.4x growth, O(n^0.25))
sdpa_math:     9.7 -> 489 MB  (50x growth, O(n^1.13))

The sublinear memory scaling of sdpa_default confirms that mem_efficient_attention achieves near O(n) memory on this hardware, closely matching FlashAttention's theoretical O(n) guarantee.

5. Connections to Prior Projects

Finding                          Connection
──────────────────────────────────────────────────────────────
O(n^2) memory = KV cache pressure  kv-cache-compaction-lab
Attention cost scales with seq     latency-breakdown-simulator
FlashAttention reduces memory      kv-cache-disaggregation-sim
GQA reduces KV footprint           real-model-profiler
sm75 vs sm80+ matters             real-model-profiler calibration