# Cara kerja penyangga

Bayangkan Anda sedang berdiskusi dengan seseorang, dan alih-alih mengingat apa yang telah mereka katakan sebelumnya, mereka harus memulai dari awal setiap kali Anda merespons. Hal ini pasti terasa lambat dan tidak efisien, bukan?

Anda dapat memperluas analogi ini ke model transformer. Generasi model autoregresif dapat lambat karena membuat prediksi satu token pada satu waktu. Setiap prediksi baru bergantung pada seluruh konteks sebelumnya.

Untuk memprediksi token ke-1000, model memerlukan informasi dari 999 token sebelumnya. Informasi ini direpresentasikan sebagai perkalian matriks di antara representasi token.

Untuk memprediksi token ke-1001, Anda membutuhkan informasi yang sama dari 999 token sebelumnya ditambah informasi apa pun dari token ke-1000. Ini adalah banyak perkalian matriks yang harus dihitung ulang dan ulang oleh model untuk setiap token!

Cache key-value (KV) menghilangkan inefisiensi ini dengan menyimpan pasangan kv yang diturunkan dari lapisan perhatian token yang telah diproses sebelumnya. Pasangan kv yang disimpan diambil dari cache dan digunakan kembali untuk token berikutnya, sehingga menghindari kebutuhan untuk menghitung ulang.

> [!WARNING]
> Caching hanya harus digunakan untuk **inference**. Pengaktifannya selama pelatihan dapat menyebabkan error yang tidak terduga.

Untuk lebih memahami cara dan alasan kerja caching, mari kita lihat lebih dekat struktur matriks attention.

## Matriks perhatian

**scaled dot-product attention** dihitung seperti yang ditunjukkan di bawah ini untuk batch berukuran `b`, jumlah attention head `h`, panjang urutan sejauh ini `T`, dan dimensi per attention head `d_head`.

$$ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V $$

Matriks query (`Q`), key (`K`), dan value (`V`) adalah proyeksi dari input embeddings dengan bentuk `(b, h, T, d_head)`.

Untuk atensi kausal, mask mencegah model memperhatikan token di masa depan. Setelah sebuah token diproses, representasinya tidak pernah berubah terhadap token di masa depan, yang berarti $ K_{\text{past}} $ dan $ V_{\text{past}} $ dapat di-cache dan digunakan kembali untuk menghitung representasi token terakhir.

$$ \text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{ter-cache}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{ter-cache}}, v_{t}]) $$

Pada saat inferensi, Anda hanya perlu query dari token terakhir untuk menghitung representasi $ x_t $ yang memprediksi token berikutnya $ t+1 $. Pada setiap langkah, vektor key dan value baru akan **disimpan** di cache dan **ditempelkan** ke key dan value sebelumnya.

$$ K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t) $$

Attention dihitung secara independen di setiap lapisan model, dan caching dilakukan berdasarkan lapisan.

Lihat tabel di bawah ini untuk membandingkan bagaimana caching meningkatkan efisiensi.

| tanpa penyangga | dengan penyangga |
|---|---|
| untuk setiap langkah, hitung ulang semua `K` dan `V` sebelumnya  | untuk setiap langkah, hanya hitung `K` dan `V` saat ini |
| biaya perhatian per langkah bersifat **kuadratik** terhadap panjang sekuens | biaya perhatian per langkah bersifat **linear** terhadap panjang sekuens (memori tumbuh secara linear, tetapi komputasi per token tetap rendah) |

## Kelas Cache

Antarmuka KV cache yang dasar menerima tensor key dan value untuk token saat ini dan mengembalikan tensor `K` dan `V` yang telah diperbarui. Hal ini dikelola secara internal oleh metode `forward` dari sebuah model.

```py
new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)
```

Saat Anda menggunakan kelas [Cache](/docs/transformers/main/id/internal/generation_utils#transformers.Cache) dari Transformers, modul self-attention melakukan beberapa langkah penting untuk mengintegrasikan informasi masa lalu dan saat ini.

1. Modul perhatian menggabungkan pasangan kv saat ini dengan pasangan kv masa lalu yang disimpan dalam cache. Hal ini menghasilkan bobot perhatian dengan bentuk `(new_tokens_length, past_kv_length + new_tokens_length)`. Pasangan kv saat ini dan masa lalu pada dasarnya digabungkan untuk menghitung skor perhatian, sehingga memastikan model menyadari konteks sebelumnya dan input saat ini.

2. Ketika metode `forward` dipanggil secara iteratif, sangat penting agar bentuk attention mask sesuai dengan panjang gabungan dari pasangan kv masa lalu dan saat ini. Attention mask harus memiliki bentuk `(batch_size, past_kv_length + new_tokens_length)`. Hal ini biasanya ditangani secara internal di [generate()](/docs/transformers/main/id/main_classes/text_generation#transformers.GenerationMixin.generate), tetapi jika Anda ingin mengimplementasikan loop generasi Anda sendiri dengan [Cache](/docs/transformers/main/id/internal/generation_utils#transformers.Cache), perhatikan hal ini! Attention mask harus menyimpan nilai token masa lalu dan saat ini.

## Implementasi penyimpanan cache

Cache disusun sebagai daftar lapisan, di mana setiap lapisan berisi cache kunci dan nilai. Cache kunci dan nilai adalah tensor dengan bentuk `[batch_size, num_heads, seq_len, head_dim]`.

Layers dapat memiliki tipe yang berbeda (misalnya `DynamicLayer`, `StaticLayer`, `StaticSlidingWindowLayer`), yang sebagian besar mengubah cara panjang sekuens ditangani dan cara cache diperbarui.

Yang paling sederhana adalah `DynamicLayer` yang tumbuh seiring dengan diprosesnya lebih banyak token. Dimensi panjang sekuens (`seq_len`) bertambah dengan setiap token baru:

```py
cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)
```

Jenis layer lain seperti `StaticLayer` dan `StaticSlidingWindowLayer` memiliki panjang sekuens tetap yang diatur saat cache dibuat. Hal ini membuatnya kompatibel dengan `torch.compile`. Dalam kasus `StaticSlidingWindowLayer`, token yang sudah ada akan digeser keluar dari cache ketika token baru ditambahkan.

Contoh di bawah ini menunjukkan cara membuat loop generasi dengan [DynamicCache](/docs/transformers/main/id/internal/generation_utils#transformers.DynamicCache). Seperti yang dibahas, attention mask adalah konkatenasi dari nilai token masa lalu dan saat ini.

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"
```