TPU Sync của Google giảm độ trễ suy luận bằng cách di chuyển dữ liệu bộ nhớ đệm KV

Google vừa giới thiệu TPU Sync, một framework phần mềm mã nguồn mở mới được thiết kế để tăng tốc suy luận AI trên các Tensor Processing Units (TPU) thế hệ thứ tám. Được công bố vào ngày 9 tháng 11 năm 2026, công cụ này giải quyết các nút thắt cổ chai về bộ nhớ trong việc phục vụ các mô hình ngôn ngữ lớn bằng cách tối ưu hóa cách dữ liệu bộ nhớ đệm key-value (KV) di chuyển giữa các tầng bộ nhớ khác nhau. Công ty tuyên bố rằng phương pháp tiếp cận này có thể giảm thời gian phản hồi token đầu tiên hơn 90% trong các kịch bản tham số cao cụ thể.

Chuyện gì đã xảy ra

Bản phát hành này đi kèm với các cập nhật phần cứng rộng lớn hơn của Google, nhắm mục tiêu đặc biệt vào các nhà phát triển chạy các mô hình lớn với cửa sổ ngữ cảnh hàng triệu token hoặc khối lượng công việc học tăng cường (reinforcement-learning). Khi độ dài ngữ cảnh tăng lên, chi phí và độ trễ trong việc quản lý bộ nhớ đệm key-value (KV) trở thành những nút thắt cổ chai đáng kể. TPU Sync nhằm giải quyết vấn đề này bằng cách di chuyển dữ liệu hiệu quả giữa bộ nhớ máy gia tốc, DRAM host, lưu trữ và mạng mà không cần sao chép thừa.

Trong các bài kiểm tra nội bộ, Google đã chứng minh những cải thiện hiệu suất đáng kể. Đối với mô hình Qwen3 có 235 tỷ tham số, việc di chuyển bộ nhớ đệm KV giảm từ vài giây xuống còn khoảng 230 mili giây. Sự giảm thiểu này hạ thấp thời gian trung bình phản hồi token đầu tiên đi 60% dưới tải nhẹ. Khi thử nghiệm một mô hình lớn hơn với 480 tỷ tham số, mức giảm độ trễ đạt tới 90%. Những cải thiện tương tự cũng được quan sát thấy trong các bài kiểm tra bộ nhớ đệm prefix với ngữ cảnh 64.000 token, nơi thông lượng tăng hơn 50% và độ trễ P99 giữa các token giảm 30%.

Tuy nhiên, những kết quả này đi kèm với các điều kiện hạn chế. Hiệu suất thay đổi tùy thuộc vào khối lượng công việc cụ thể, mức độ tái sử dụng dữ liệu được cache và băng thông mạng khả dụng. Phần mềm chưa được coi là sẵn sàng cho môi trường sản xuất (production-ready) đối với mọi trường hợp sử dụng. Google cảnh báo rằng các API vẫn đang phát triển và khuyên các nhà phát triển nên ghim (pin) các bản triển khai của họ vào các commit xác minh cụ thể hoặc các phiên bản release có tag, thay vì theo dõi nhánh chính (main branch).

Cách thức hoạt động

TPU Sync được viết bằng C++ và hoạt động bằng cách quản lý các chuyển giao dữ liệu qua phân cấp bộ nhớ của các host TPU. Nó sử dụng các chuyển giao zero-copy và các descriptor phần cứng PjRtBuffer gốc để loại bỏ sự trùng lặp dữ liệu thừa. Phương pháp này cho phép hệ thống tiếp cận các giới hạn lý thuyết của giao diện phần cứng. Trong các bài kiểm tra trên host TPU Ironwood, thông lượng kết hợp từ host-to-device và device-to-host đạt 223 GB/s, gần với giới hạn 256 GB/s của PCIe Gen5.

Framework này tích hợp với các hệ sinh thái suy luận hiện có, bao gồm TorchTPU, vLLM và SGLang. Tính tương thích này có nghĩa là các nhà phát triển có thể chạy các khối lượng công việc tách rời (disaggregated workloads) trên TPUs mà không cần xây dựng lại toàn bộ stack phục vụ của họ. Việc tách rời (Disaggregation) phân tách các giai đoạn prefill và decode của quá trình suy luận, cho phép mỗi giai đoạn được tối ưu hóa độc lập. Trước đây, việc phục vụ TPU bên ngoài thiếu cơ sở hạ tầng để xử lý hiệu quả sự phân tách này, nhưng TPU Sync cung cấp các cơ chế cần thiết để quản lý trạng thái và di chuyển dữ liệu giữa các giai đoạn này.

Chi tiết chính

• TPU Sync là một framework mã nguồn mở viết bằng C++ để quản lý việc di chuyển dữ liệu bộ nhớ đệm KV.
• Các bài kiểm tra benchmark cho thấy mức giảm lên đến 90% thời gian phản hồi token đầu tiên đối với các mô hình có 480 tỷ tham số.
• Thông lượng trên các host TPU Ironwood đạt 223 GB/s, gần với giới hạn lý thuyết của PCIe Gen5.
• Phần mềm hỗ trợ Cloud TPU v7x và v6e với các yêu cầu phiên bản JAX và PyTorch cụ thể.
• Có sẵn tích hợp cho các engine suy luận phổ biến như vLLM và SGLang thông qua TorchTPU.
• Google khuyên dùng các commit xác minh hoặc các phiên bản release có tag vì các API vẫn đang thay đổi.

Tại sao điều này quan trọng

Đối với các nhóm kỹ thuật xây dựng sản phẩm với các mô hình ngôn ngữ lớn, độ trễ suy luận ảnh hưởng trực tiếp đến trải nghiệm người dùng và chi phí vận hành. Thời gian phản hồi token đầu tiên là một chỉ số quan trọng đối với các ứng dụng tương tác, và việc giảm nó với biên độ như vậy có thể làm cho các cửa sổ ngữ cảnh trước đây không thực tế trở nên khả thi. Bằng cách giải quyết lớp phần mềm quản lý việc di chuyển bộ nhớ, Google đang cố gắng thu hẹp khoảng cách với hệ sinh thái CUDA của Nvidia, vốn đã thống trị suy luận AI từ lâu nhờ các công cụ trưởng thành và tối ưu hóa hiệu suất.

Sự phát triển này cũng nhấn mạnh tầm quan trọng ngày càng tăng của các kiến trúc phục vụ tách rời (disaggregated serving architectures). Khi các mô hình mở rộng quy mô, việc xem xét prefill và decode như các vấn đề riêng biệt cho phép tận dụng tài nguyên tốt hơn. TPU Sync cung cấp cơ sở hạ tầng để làm điều này khả thi trên phần cứng của Google, tiềm năng mang lại một lựa chọn thay thế tiết kiệm chi phí so với các hệ thống GB200 và GB300 của Nvidia cho một số khối lượng công việc nhất định. Tuy nhiên, các so sánh trực tiếp với các thiết lập Nvidia được tinh chỉnh đầy đủ vẫn còn hạn chế, và các nhà phát triển phải cân nhắc lợi ích so với sự chưa hoàn thiện hiện tại của phần mềm.

Bạn có thể làm gì

• Đánh giá xem các khối lượng công việc hiện tại của bạn có liên quan đến các cửa sổ ngữ cảnh lớn hoặc các mô hình tham số cao sẽ hưởng lợi từ việc tối ưu hóa bộ nhớ đệm KV hay không.
• Thử nghiệm TPU Sync trong môi trường không phải sản xuất (non-production) bằng cách sử dụng các instance Cloud TPU v7x hoặc v6e để đo lường mức cải thiện độ trễ thực tế.
• Ghim các phụ thuộc của bạn vào các phiên bản release có tag hoặc các commit xác minh cụ thể để tránh các thay đổi gây hỏng hóc từ các API đang phát triển.
• Xem xét lại stack suy luận hiện có của bạn để xem liệu nó có thể tích hợp với TorchTPU, vLLM hoặc SGLang để dễ dàng áp dụng hay không.
• Giám sát băng thông mạng và tỷ lệ tái sử dụng dữ liệu được cache, vì các yếu tố này ảnh hưởng đáng kể đến mức cải thiện hiệu suất mà bạn sẽ thấy.
• Theo dõi các bản cập nhật trong tương lai từ Google về tính ổn định của API và hỗ trợ phần cứng rộng rãi hơn trước khi lên kế hoạch di chuyển sang môi trường sản xuất.