Google rende open source TPU Sync per il trasferimento di pesi e cache KV
15 novembre 2026 · 2 min di lettura
Google Cloud ha rilasciato in modalità open source TPU Sync, una libreria nativa in C++ progettata per spostare i blocchi della cache KV e i pesi dei modelli tra la memoria HBM delle TPU, la memoria DRAM dell'host, l'archiviazione condivisa e la rete del data center.
Il progetto si rivolge a due carichi di lavoro specifici in cui gli acceleratori attendono spesso i dati: l'inferenza di modelli linguistici di grandi dimensioni disaggregata e a contesto lungo, e i cicli di apprendimento per rinforzo che inviano i pesi aggiornati dai trainer ai campionatori di rollout.
Per quanto riguarda la disaggregazione tra prefill e decode, Google riferisce che la latenza mediana del trasferimento KV per il modello Qwen3 235B è scesa da diversi secondi a circa 230 millisecondi. Per tale modello, le riduzioni medie del tempo per il primo token sono state del 60%, mentre per un modello da 480B sono state del 90% con bassa concorrenza. Per la cache dei prefissi, l'azienda riporta una riduzione media del tempo per il primo token fino al 90% su carichi di lavoro con contesto di 64K.
Gli sviluppatori hanno a disposizione pacchetti Python per JAX e PyTorch, che includono interfacce di programmazione per il gestore della cache KV, l'archivio KV e il sincronizzatore dei pesi. TPU Sync funge da connettore KV per vLLM-TorchTPU tramite l'interfaccia KVConnector di vLLM, e come backend per la cache gerarchica e la disaggregazione per SGLang-TorchTPU. Inoltre, si associa a llm-d 1.0 su GKE per l'instradamento consapevole della cache dei prefissi. Il modulo di sincronizzazione dei pesi per l'apprendimento per rinforzo per Pathways e JAX multi controller è invece sperimentale.
La libreria è validata su TPU v7x e v6e. Il repository GitHub, che ha assorbito il precedente codice di TPU Raiden, avverte che le interfacce di programmazione cambiano rapidamente e che su PyPI sono disponibili solo le versioni di sviluppo, consigliando di fissare l'installazione all'ultimo commit funzionante noto. Il pacchetto PyTorch richiede una build di torch_tpu corrispondente che non è presente in un indice pubblico.
TPU Sync è disponibile sotto licenza Apache 2.0. Il repository include documentazione sull'architettura, interfacce in C++ e Python, ed esempi per l'offloading della cache KV dell'host e per l'inferenza disaggregata a host singolo e multiplo. Fino al rilascio di una versione stabile, i pacchetti JAX e PyTorch sono pubblicati solo come build di sviluppo, mentre il livello di archiviazione e la sincronizzazione dei pesi per l'apprendimento per rinforzo rimangono sperimentali.