Google เปิดซอร์ส TPU Raiden: ไลบรารีโอเพนซอร์สจัดการ KV-Cache ที่ท้าชน NIXL ของ NVIDIA โดยตรง

TPU Raiden จัดการ KV-Cache Transfer อย่างไร

ฟังก์ชันหลักของ Raiden แบ่งออกเป็น 3 ความสามารถสำคัญ:

การเคลื่อนย้าย KV-cache ระหว่างอินสแตนซ์ Raiden ถ่ายโอนเทนเซอร์ key-value จากอินสแตนซ์ TPU ที่ทำ prefill ไปยังอินสแตนซ์ TPU ที่ทำ decode ในสถาปัตยกรรมการให้บริการแบบแยกส่วน ทำให้สามารถแยกเฟสเหล่านี้ไปทำงานบนฮาร์ดแวร์เฉพาะได้ ซึ่งเป็นงานเดียวกับที่ NIXL ของ NVIDIA ทำในการปรับใช้บน GPU

การออฟโหลด (Offloading) Raiden มีกลไกพื้นฐานสำหรับย้ายข้อมูล KV-cache ออกจากหน่วยความจำ TPU ไปยังสตอเรจภายนอก รองรับการจัดการแคชแบบหลายชั้น (hierarchical cache management) คล้ายกับที่ NIXL ทำกับแบ็คเอนด์ NVMe และ RDMA ซึ่งจำเป็นต่อการขยายความจุแคชให้เกินกว่าหน่วยความจำบนชิปที่มีราคาแพง

การขนส่งแบบ TPU-native แตกต่างจากสแต็ก GPUDirect RDMA ของ NVIDIA ตรงที่ Raiden ทำงานบนโครงข่ายเชื่อมต่อ TPU (ICI) ของ Google และถูกปรับให้เหมาะกับฮาร์ดแวร์ TPU v5e ขึ้นไป ซึ่งหมายความว่ามันถูกออกแบบมาสำหรับสถาปัตยกรรมตัวเร่งความเร็วของ Google โดยเฉพาะ ไม่ใช่การดัดแปลงจากแนวทางที่เน้น GPU

TPU Raiden เทียบกับ NVIDIA NIXL: การเปรียบเทียบฟีเจอร์

ตารางต่อไปนี้สรุปการเปรียบเทียบของทั้งสองไลบรารีในมิติสำคัญ:

เทรนด์อุตสาหกรรม: การเปิดซอร์สโครงสร้างพื้นฐาน Inference

การปล่อย TPU Raiden เป็นส่วนหนึ่งของเทรนด์ที่ชัดเจนและเร่งตัวขึ้นในการเปิดซอร์สสแต็กซอฟต์แวร์สำหรับการอนุมาน

ผู้ให้บริการ Hyperscaler ทั้งสองรายแข่งกันเปิดซอร์สสแต็กของตน NVIDIA เปิดซอร์ส NIXL พร้อมกับเฟรมเวิร์ก Dynamo ที่งาน GTC 2025 Google ก็ทยอยเปิดซอร์สซอฟต์แวร์ TPU อย่างต่อเนื่อง เริ่มจากการผสาน vLLM TPU, ตามด้วยเฟรมเวิร์ก llm-d ที่ทำงานบน Kubernetes สำหรับการอนุมานแบบกระจาย และล่าสุดกับ Raiden

การอนุมานแบบแยกส่วนกลายเป็นสถาปัตยกรรมการให้บริการมาตรฐาน การแยกเฟส prefill และ decode ช่วยปรับปรุงเวลา Time-to-First-Token (TTFT) และการใช้ทรัพยากร แต่ก็ทำให้การถ่ายโอน KV-cache ที่รวดเร็วกลายเป็นคอขวดด้านประสิทธิภาพ ทั้ง Raiden และ NIXL มีไว้เพื่อแก้ปัญหานี้โดยเฉพาะ

การล็อกผู้ใช้กับผู้ขายรายใดรายหนึ่งกำลังถูกต่อสู้ที่ชั้นการเคลื่อนย้ายข้อมูล NIXL ถูกอธิบายว่า "ไม่ผูกติดกับผู้ขาย" (vendor-agnostic) และรองรับ AWS EFA ไม่ใช่แค่โครงข่ายเชื่อมต่อของ NVIDIA เอง Raiden ก็เสียบเข้ากับเฟรมเวิร์กแบบเปิด (vLLM, SGLang, llm-d) เช่นกัน ไลบรารีทั้งสองถูกออกแบบมาเพื่อทำให้ระบบนิเวศฮาร์ดแวร์ของตนน่าสนใจยิ่งขึ้นสำหรับนักพัฒนา แทนที่จะบังคับใช้ API ที่เป็นกรรมสิทธิ์

ระบบนิเวศกำลังมาบรรจบกันที่อินเทอร์เฟซ KV-Connector แบบเสียบปลั๊กได้ API KVConnector ของ vLLM รองรับคอนเนกเตอร์ NIXL และ Mooncake แล้ว และสถาปัตยกรรมถูกออกแบบให้รองรับแบ็คเอนด์ใด ๆ รวมถึง Raiden ทำให้ผู้ปฏิบัติงานสามารถสลับชั้นการขนส่งได้โดยไม่ต้องเปลี่ยน inference engine ความสามารถในการเสียบปลั๊กนี้มีความสำคัญอย่างยิ่งในสภาพแวดล้อมแบบหลายตัวเร่งความเร็วที่ความยืดหยุ่นของฮาร์ดแวร์เป็นสิ่งสำคัญ

ทำไม TPU Raiden จึงสำคัญต่อโครงสร้างพื้นฐาน AI

TPU Raiden จัดการกับความท้าทายในการขยายขนาดขั้นพื้นฐาน เมื่อโมเดลภาษาขนาดใหญ่เติบโตขึ้น KV cache ที่สร้างขึ้นระหว่างการอนุมานก็มหาศาลตามไปด้วย ซึ่งมักจะเกินหน่วยความจำบนชิปของตัวเร่งความเร็วแต่ละตัว การย้ายข้อมูลนี้ระหว่างโหนด prefill และ decode อย่างมีประสิทธิภาพคือความแตกต่างระหว่างระบบอนุมานที่ตอบสนองได้ดีกับระบบที่ติดขัดเพราะคอขวดในการถ่ายโอนข้อมูล

การเปิดซอร์ส Raiden ของ Google ไม่ใช่แค่การเล่นตามเกมของ NIXL จาก NVIDIA เท่านั้น แต่ยังเป็นการส่งสัญญาณว่าการอนุมานบน TPU เป็นตัวเลือกที่จริงจังสำหรับงาน AI ในระบบผลิต ไลบรารีนี้มอบเครื่องมือระดับเดียวกันให้กับผู้ใช้ TPU อย่างที่ผู้ใช้ GPU มีตั้งแต่ NIXL เปิดตัว: การควบคุมอย่างละเอียดว่าข้อมูล KV-cache เคลื่อนที่ระหว่างตัวเร่งความเร็ว ลำดับชั้นหน่วยความจำ และสตอเรจภายนอกอย่างไร

สำหรับอุตสาหกรรม AI ในวงกว้าง การปล่อย Raiden หมายความว่าการอนุมานแบบแยกส่วนบน TPU ตอนนี้เป็นทางเลือกที่ใช้งานได้จริงพร้อมเครื่องมือที่เหมาะสม นักพัฒนาที่ใช้ฮาร์ดแวร์ TPU ของ Google สำหรับการอนุมานในระบบผลิตสามารถใช้ประโยชน์จากรูปแบบสถาปัตยกรรมแบบแยกส่วนเดียวกันที่กลายเป็นมาตรฐานบนคลัสเตอร์ GPU ของ NVIDIA โดยไม่ถูกล็อกกับ API การเคลื่อนย้ายข้อมูลที่เป็นกรรมสิทธิ์