KV cache: จดไว้ ไม่ต้องคิดใหม่
จำตึกในบทที่ 3 ได้มั้ย ตอนเขียนคำใหม่ มีแค่ token ใหม่ที่ขึ้นตึก แล้วของเก่าล่ะ? token ใหม่ต้องใช้ คิวรี-คีย์-แวลู (query, key, value)สามบทบาทใน attention คำถาม (query) ไปเทียบกับป้ายชื่อ (key) ของทุกคำ แล้วผสมข้อมูล (value) ตามน้ำหนักที่ได้ดูคำศัพท์ทั้งหมด ของทุก token ก่อนหน้า (บทที่ 2) ในทุกชั้น ถ้าไม่จด token ที่ 1,000 ต้องดัน 999 token เก่าผ่านทุกชั้นใหม่
นับดู คำตอบ 1,000 token (ไม่มี prompt):
- ไม่จดไว้: 1 + 2 + … + 1,000 = 500,500 ครั้งที่ดัน token ผ่านโมเดล เท่าช่องที่ไม่ถูกบังในตาราง การบังคำอนาคต (causal mask)กฎใน attention ที่ให้แต่ละ token ฟังได้แค่ตัวเองและคำก่อนหน้า โดยตั้งคะแนนของคำในอนาคตเป็น −∞ ก่อน softmaxดูคำศัพท์ทั้งหมด ของบทที่ 3 ()
- จดไว้: แค่ 1,000 ครั้ง
ทำไมจดได้? causal mask ทำให้ token เก่าไม่เคย "มอง" token ใหม่ K, V ของมันจึงไม่เปลี่ยน คำนวณครั้งเดียวแล้วเก็บใน แคชคีย์และแวลู (KV cache)ที่เก็บ key และ value ของ token ที่คำนวณแล้วในทุกชั้น จะได้ไม่ต้องคำนวณซ้ำทุกครั้งที่เขียน token ใหม่ ยิ่งข้อความยาวยิ่งใช้หน่วยความจำมากดูคำศัพท์ทั้งหมด
ที่ต้องจ่ายคือหน่วยความจำ ต่อ token = 2 (K และ V) × จำนวนชั้น × ชุด K/V × ตัวเลขต่อชุด × byte
โมเดลสมมติคือ 8B จากบทที่ 3 (32 ชั้น, 32 หัว, หัวละ 128 ตัวเลข) แต่หลายหัวใช้ K/V ชุดเดียวกัน (grouped-query attention, GQA) จึงเก็บแค่ 8 ชุดต่อชั้น
ลองคิดตาม (ตัวเลขละ 2 byte หรือ 16 บิต)
- ต่อ token: 2 × 32 × 8 × 128 × 2 = 131,072 byte
- 1,000 token: 0.13 GB
- 128,000 token: 16.8 GB มากกว่าน้ำหนักโมเดลที่ 16 บิต (16.0 GB หน้า 5)
KV cache เก็บ key กับ value ของทุก token ไว้ใช้ซ้ำ ไม่ต้องคำนวณใหม่ แต่กินหน่วยความจำเพิ่มเป็นเส้นตรงตามความยาวข้อความ
มังคุด
เก็บของเก่าไว้แล้ว ทำไมแชตยาวยังช้าลงอยู่ดีล่ะ?