Efficient Attention

Semua teknik sejauh ini mengatur ulang penyimpanan; teknik terakhir mengubah perhitungannya sendiri. Attention standar menghitung matriks skor berukuran n×n — n adalah panjang sequence — sehingga memori tumbuh kuadratik: teks 4x lebih panjang berarti skor 16x lebih besar. Pada konteks panjang (ribuan hingga ratusan ribu token), matriks skor inilah yang pertama kali membuat pelatihan kehabisan memori, jauh sebelum parameter menjadi masalah. Efficient attention memecahkannya lewat tiga rute: operator fused yang tidak pernah menuliskan matriks penuh (FlashAttention), pola sparse yang memangkas pasangan yang dihitung, dan varian linear yang menulis ulang rumus menjadi kompleksitas hampir linier. Chapter ini menghitung biaya kuadratiknya dengan angka nyata, membuktikan tiling baris menghasilkan keluaran identik dengan hitungan penuh, dan menunjukkan API scaled_dot_product_attention yang menjadi pintu masuk semuanya.

Mengapa Memori Attention Tumbuh Kuadratik #

Rumus attention yang sudah kita bedah di chapter Scaled Dot-Product Attention menghitung softmax(QK^T / √d) @ V. Langkah tengahnya — matriks skor QK^T — berukuran n×n untuk setiap pasangan batch dan head: batch 1, satu head, float32, berarti 4 × n² byte. Struktur memorinya tersusun sebagai berikut:

Matriks skor bersifat sementara — tidak dibawa ke backward sebagai parameter — tetapi puncak alokasinya tetap harus tersedia sekaligus. Berkas efficient_attention.py menghitung harga itu untuk tiga panjang sequence, membandingkan attention manual dengan API fused, lalu membuktikan tiling baris tidak mengubah hasil:

"""Efficient attention: biaya O(n^2) vs scaled_dot_product_attention dan tiling."""

import math

import torch
import torch.nn.functional as F

torch.manual_seed(9)

# ---------- 1. Attention manual vs SDPA fused ----------
B, H, n, d = 2, 3, 128, 32
q = torch.randn(B, H, n, d)
k = torch.randn(B, H, n, d)
v = torch.randn(B, H, n, d)

skor = q @ k.transpose(-2, -1) / math.sqrt(d)  # (B, H, n, n)
manual = skor.softmax(dim=-1) @ v
sdpa = F.scaled_dot_product_attention(q, k, v)
print("=== Attention manual vs F.scaled_dot_product_attention ===")
print(f"shape skor q@k^T     : {tuple(skor.shape)} "
      f"= {skor.numel() * skor.element_size():,} byte")
print(f"max |selisih hasil|  : {float((manual - sdpa).abs().max()):.2e}")
print(f"allclose (atol 1e-5) : {torch.allclose(manual, sdpa, atol=1e-5)}")

# ---------- 2. Memori matriks skor tumbuh kuadratik ----------
print()
print("=== Memori matriks skor (batch 1, 1 head, float32) ===")
print(f"{'panjang seq (n)':>15s} {'n^2 elemen':>12s} {'byte (4B)':>12s} {'setara':>9s}")
for n_seq in (1_000, 4_000, 16_384):
    byte = n_seq * n_seq * 4
    setara = f"{byte / 1024**2:.1f} MB" if byte < 1024**3 else f"{byte / 1024**3:.1f} GB"
    print(f"{n_seq:>15,} {n_seq * n_seq:>12,} {byte:>12,} {setara:>9s}")
print("4x panjang seq -> 16x memori skor (pertumbuhan kuadratik)")

# ---------- 3. Tiling baris: hasil identik, memori jauh lebih kecil ----------
print()
print("=== Tiling baris (analogi blok FlashAttention) ===")
n2, d2, blok = 4096, 64, 128
q2 = torch.randn(n2, d2)
k2 = torch.randn(n2, d2)
v2 = torch.randn(n2, d2)

penuh = ((q2 @ k2.T / math.sqrt(d2)).softmax(dim=-1) @ v2)
byte_penuh = n2 * n2 * 4

tiling = []
byte_puncak_blok = 0
for awal in range(0, n2, blok):
    q_b = q2[awal : awal + blok]  # hanya blok baris query
    skor_b = q_b @ k2.T / math.sqrt(d2)
    byte_puncak_blok = max(byte_puncak_blok, skor_b.numel() * 4)
    tiling.append(skor_b.softmax(dim=-1) @ v2)
gabung = torch.cat(tiling, dim=0)
print(f"skor penuh    : {n2}x{n2} = {byte_penuh:,} byte ({byte_penuh / 1024**2:.1f} MB)")
print(f"skor per blok : {blok}x{n2} = {byte_puncak_blok:,} byte "
      f"({byte_puncak_blok / 1024**2:.2f} MB)")
print(f"hemat puncak  : {byte_penuh / byte_puncak_blok:.0f}x lebih kecil")
print(f"hasil identik : {torch.equal(penuh, gabung)} (torch.equal)")

Output yang Dihasilkan:

=== Attention manual vs F.scaled_dot_product_attention ===
shape skor q@k^T     : (2, 3, 128, 128) = 393,216 byte
max |selisih hasil|  : 8.34e-07
allclose (atol 1e-5) : True

=== Memori matriks skor (batch 1, 1 head, float32) ===
panjang seq (n)   n^2 elemen    byte (4B)    setara
          1,000    1,000,000    4,000,000    3.8 MB
          4,000   16,000,000   64,000,000   61.0 MB
         16,384  268,435,456 1,073,741,824    1.0 GB
4x panjang seq -> 16x memori skor (pertumbuhan kuadratik)

=== Tiling baris (analogi blok FlashAttention) ===
skor penuh    : 4096x4096 = 67,108,864 byte (64.0 MB)
skor per blok : 128x4096 = 2,097,152 byte (2.00 MB)
hemat puncak  : 32x lebih kecil
hasil identik : True (torch.equal)

Mari kita telaah kode di atas baris per baris:

1. Baris 16-23 (manual vs fused): jalur manual menghitung skor (393.216 byte untuk B=2, H=3, n=128), softmax, lalu kalikan V; jalur F.scaled_dot_product_attention menghasilkan keluaran yang sama dengan selisih maksimum 8,34e-07 — setara pembulatan float32, bukan perbedaan algoritma.
2. Baris 29-33 (tabel kuadratik): n=1.000 memakan 3,8 MB; n=4.000 memakan 61,0 MB; n=16.384 (panjang konteks model bahasa modern) memakan 1,0 GB — per head, per sampel. Model dengan 32 head dan batch 8 pada konteks tersebut membutuhkan ratusan GB hanya untuk skor. Setiap 4x penambahan panjang berarti 16x memori.
3. Baris 43-53 (tiling baris): loop memproses 128 baris query sekali jalan — skor_b hanya berukuran 128×4096 (2,00 MB), softmax dihitung per blok, hasilnya dirangkai dengan torch.cat. Puncak memori skor turun dari 64,0 MB menjadi 2,00 MB (32x lebih hemat) dan torch.equal membuktikan hasil identik bit-per-bit.
4. Baris 54-58 (ringkasan): penghematan berskala n/blok — makin kecil bloknya, makin rendah puncak memori, dengan jumlah langkah yang bertambah proporsional. Tidak ada hasil yang berubah; yang berubah hanyalah urutan kerja.

FlashAttention, Sparse, dan Linear #

Tiling baris di atas adalah inti FlashAttention, dengan dua penyempurnaan penting pada implementasi aslinya. Pertama, blok diproses langsung di memori cepat GPU (SRAM) sambil melakukan softmax secara inkremental (online softmax) — statistik pembagi terkumpul bertahap tanpa pernah menyimpan satu baris skor penuh sekaligus. Kedua, desainnya sadar-IO: yang dioptimalkan adalah jumlah transfer antara memori HBM yang lambat dan SRAM yang cepat, bukan sekadar jumlah FLOP. Hasil matematisnya tetap attention penuh (bukan aproksimasi) — berbeda dengan dua rute lain:

F.scaled_dot_product_attention adalah pintu masuk tunggal PyTorch: operator ini otomatis memilih kernel tercepat yang tersedia (flash / mem-efficient / math) sesuai perangkat dan bentuk input.

Hingga titik ini kita telah mengumpulkan seluruh kotak perkakas: memahami anatomi memori, menumpuk gradien, menghitung ulang aktivasi, membagi data, model, pipa, dan parameter, hingga menulis ulang attention. Chapter penutup section ini merangkai semuanya menjadi daftar periksa praktis yang terurut — lengkap dengan alat pemantauan untuk mengukur sebelum dan sesudahnya.

Rujukan resmi API scaled_dot_product_attention tersedia di docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html ; makalah FlashAttention dapat dibaca di arxiv.org/abs/2205.14135 .