Kuantisasi KV Cache dengan Error Budget yang Eksplisit

Inference transformer autoregresif menyimpan tensor key dan value dari token sebelumnya agar setiap token baru dapat melakukan attention ke konteks terdahulu tanpa menghitung ulang proyeksi tersebut. Ketika context length dan jumlah sequence concurrent meningkat, KV cache ini dapat menjadi bagian besar dari memori accelerator.

Kuantisasi KV cache menyimpan tensor tersebut pada presisi lebih rendah dan merekonstruksi aproksimasi ketika attention menggunakannya. Perhitungan memorinya menarik, tetapi error yang dihasilkan bukan perturbasi generik seperti pada model weight. Key yang terkuantisasi memengaruhi attention score sebelum softmax, sedangkan value terkuantisasi memengaruhi weighted sum setelah attention probability terbentuk.

Perbedaan ini memberi cara yang lebih berguna untuk menilai presisi cache daripada menganggap semua tensor cache memiliki sensitivitas yang sama.

Ukuran cache mengikuti jumlah elemen tersimpan dan bit width

Untuk decoder dengan L transformer block, anggap setiap block menyimpan H_kv key-value head, setiap head berdimensi D, dan sebuah sequence saat ini memiliki T token dalam cache. Dengan mengabaikan metadata dan allocator overhead, cache menyimpan jumlah scalar berikut:

2 * L * T * H_kv * D

Faktor dua mewakili key dan value. Mengalikannya dengan byte per elemen tersimpan memberikan payload size untuk satu sequence.

Mengurangi cache dari format 16-bit ke representasi 8-bit kira-kira membagi dua payload elemen tersebut. Beralih ke penyimpanan 4-bit kira-kira menjadikannya seperempat. Memori aktual yang dialokasikan dapat berbeda karena format low-bit mungkin memerlukan scale, zero point, padding, packing, alignment, atau buffer dequantization sementara.

Karena itu, perbandingan yang berguna bukan hanya nominal bit width, melainkan total byte yang dipertahankan per cached token pada representasi konkret yang digunakan inference engine.

Error pada key mengubah input softmax

Untuk satu attention head, query vector q memberi skor pada cached key k_i melalui scaled dot product:

s_i = (q dot k_i) / sqrt(D)

Jika kuantisasi merekonstruksi key sebagai k_i + e_i, skornya menjadi:

s_i' = s_i + (q dot e_i) / sqrt(D)

Dampak error key bergantung pada proyeksinya terhadap query saat ini. Error rekonstruksi key dengan norm besar dapat memiliki dampak skor kecil untuk suatu query jika error hampir orthogonal terhadap query tersebut. Error yang lebih kecil tetapi searah dengan query dapat mengubah skor lebih banyak.

Skor yang terganggu kemudian melewati softmax. Softmax menghubungkan antarposisi: perubahan satu skor dapat mengubah normalized probability yang diberikan ke beberapa cached token. Karena itu, error budget praktis untuk key terkait dengan distorsi attention score, bukan hanya elementwise reconstruction error.

Ini juga berarti format cache yang dievaluasi hanya dengan tensor mean squared error dapat melewatkan perilaku penting bagi attention. Reconstruction metric tetap berguna sebagai diagnostik, tetapi tidak secara langsung menyatakan downstream score perturbation untuk query yang benar-benar dihasilkan model.

Error pada value masuk setelah attention probability

Value berada pada posisi berbeda dalam komputasi. Dengan attention probability p_i, output head adalah:

output = sum_i p_i * v_i

Jika cached value direkonstruksi sebagai v_i + r_i sementara probability tetap, kontribusi langsungnya terhadap output error adalah:

sum_i p_i * r_i

Value pada posisi dengan attention probability sangat kecil memberi kontribusi error langsung yang kecil untuk query tersebut. Error pada posisi yang diberi bobot besar memberi kontribusi lebih besar. Error dari beberapa posisi juga dapat saling memperkuat atau sebagian saling meniadakan dalam vector sum.

Pemisahan ini tidak berarti value selalu aman untuk dikuantisasi lebih agresif daripada key. Sensitivitas bergantung pada model, format cache, calibration data, konteks, dan operating point. Namun, hal ini menunjukkan bahwa presisi key dan value tidak harus diperlakukan sebagai satu parameter yang tidak dapat dipisahkan.

Granularitas kuantisasi mengendalikan rentang yang harus dicakup setiap scale

Kode integer low-bit hanya memiliki sejumlah level representasi. Memetakan tensor ke level tersebut memerlukan scale dan terkadang zero point. Kelompok elemen yang berbagi parameter itu menentukan rentang yang harus direpresentasikan satu quantizer.

Satu scale untuk tensor besar murah dari sisi metadata, tetapi satu magnitude ekstrem dapat memperlebar rentang dan menyisakan lebih sedikit level efektif untuk nilai yang lebih kecil. Group yang lebih halus memungkinkan scale menyesuaikan rentang lokal, dengan biaya metadata dan kerja kuantisasi tambahan.

Untuk KV cache, pilihan grouping dapat mengikuti token, head, channel, atau fixed-size element group, tergantung engine dan format. Pilihan ini tidak dapat dipertukarkan begitu saja. Per-token scale beradaptasi ketika setiap token baru datang, sedangkan scale yang dibagi untuk banyak token harus mengakomodasi nilai dari posisi berbeda.

Representasi juga memerlukan kebijakan untuk nilai ekstrem. Clipping mempersempit rentang dan meningkatkan resolusi di dalamnya, tetapi nilai di luar rentang akan saturate. Memperluas rentang menghindari clipping tersebut namun membuat jarak antarlevel kuantisasi lebih besar. Keseimbangan yang sesuai bergantung pada distribusi cache aktual dan error yang dapat ditoleransi model.

Pertumbuhan cache dinamis mengubah persoalan kalibrasi

Model weight tetap selama inference, sehingga weight quantization dapat menentukan parameter dari tensor yang sudah diketahui sebelum request tiba. Entri KV cache dihasilkan secara online dari hidden state yang bergantung pada request. Nilai cache masa depan belum tersedia ketika request dimulai.

Hal ini membuat static range menjadi asumsi tentang activation di masa depan. Jika runtime value melewati rentang tersebut, clipping atau saturation dapat meningkat. Dynamic scale beradaptasi terhadap cache entry yang diamati, tetapi menambah komputasi scale dan metadata; grouping menentukan seberapa sering pekerjaan itu dilakukan.

Long context menambah dimensi lain. Skema kuantisasi yang bekerja baik dekat awal sequence tetap perlu dievaluasi pada context length yang digunakan aplikasi. Cache entry bertahan, sehingga keputusan kuantisasi awal dapat tetap relevan setelah banyak iterasi decoding.

Biaya dequantization harus masuk ke model latency

Penyimpanan terkompresi tidak menjamin end-to-end latency lebih rendah. Attention kernel harus mengonsumsi representasi cache dengan suatu cara. Implementasi dapat melakukan dequantize ke tipe yang lebih lebar sebelum aritmetika, menggabungkan unpacking dan scaling ke attention kernel, atau menggunakan operasi hardware yang mendukung format tersimpan secara lebih langsung.

Jalur tersebut memiliki karakteristik memory traffic, arithmetic, temporary storage, dan kernel launch yang berbeda. Format yang menghemat kapasitas cache tetap dapat menambah conversion work hingga memperburuk latency pada device atau sequence shape tertentu.

Perbandingan juga berubah berdasarkan batch size dan context length. Pada konteks pendek, cache traffic mungkin hanya bagian kecil dari total decoding work. Pada konteks panjang, membaca cached key dan value menjadi operasi yang lebih besar. Klaim mengenai kecepatan karena itu memerlukan pengukuran pada sequence length, concurrency, kernel, dan hardware yang sesuai dengan target deployment.

Evaluasi harus memisahkan peningkatan kapasitas dari dampak model

Eksperimen cache lebih mudah ditafsirkan ketika melaporkan dua kelompok hasil secara terpisah. Pertama adalah perilaku sistem: byte yang dipertahankan per token, context atau concurrency maksimum yang memungkinkan, decode latency, dan temporary memory yang muncul akibat conversion. Kedua adalah perilaku model pada prompt dan decoding configuration yang sama.

Perbandingan juga perlu memakai cache policy yang sama selama satu run. Mencampur presisi antar-layer, mempertahankan recent window pada presisi lebih tinggi, atau mengecualikan tensor tertentu dapat menjadi desain yang valid, tetapi masing-masing mengubah jalur error dan persamaan memori. Kebijakan tersebut harus dianggap sebagai bagian dari representasi, bukan detail implementasi tersembunyi.

Kuantisasi KV cache paling berguna ketika presisinya dinyatakan sebagai resource budget dan error budget yang eksplisit. Stored bit width hanya menentukan bagian pertama. Grouping, range selection, distorsi key score, rekonstruksi value, conversion kernel, dan context length menentukan arti bit width tersebut bagi decoder yang sebenarnya.