Google targets AI’s next infrastructure bottlenecks with two new TPUs
Network World ·
AI models are becoming more performant, efficient, and sophisticated by the day; but this can be incredibly taxing for the infrastructure running them. Google Cloud is addressing these growing memory, security, and performance demands with two additions to its eighth-generation tensor processing unit (TPU) portfolio. TPU 8tx is optimized for secure and confidential inference, while TPU 8i Max is optimized for large context windows and high-throughput, low-latency inference. “As AI applications evolve and transform, the infrastructure that powers them needs to transform as well,” Mark Lohmeyer , VP and GM for AI and computing structure at Google Cloud, explained during a briefing. Incorporating security into the hardware layer Earlier this year, Google Cloud rolled out the eighth generation of its TPUs : TPU 8t optimized for large scale training, and TPU 8i optimized for low latency and cost-effective inference. The company said that TPU 8t and 8i offer up to 2.7 times better training performance and 80% better inference price performance, respectively, compared to previous generation TPUs. The TPU systems bring together the chips and the compute, networking, and storage surrounding them. “For the first time, we had bifurcated the portfolio, so instead of just a single system and a single chip, we actually had two chips and systems,” Lohmeyer said. TPU 8tx and TPU 8ix Max expand on the two products. TPU 8tx is Google Cloud’s first TPU to feature confidential computing capabilities, in partnership with AMD. Security is added directly into the hardware layer, keeping sensitive data and model weights “protected and isolated, even when they’re in active use,” he explained. Many enterprises want to use AI with sensitive information, whether that be proprietary or internal business data or information that has to be protected for regulatory or operational reasons. The idea with TPU 8tx is to secure the entire workflow, encrypt the flow of data from the host to the TPUs, and secure the TPUs themselves as they’re running active inference. “Then that entire system is attested so that we know it is what it says it is,” Lohmeyer said. At the same time, Google Cloud has brought the low latency networking architecture from TPU 8i to TPU 8tx, he added. This helps reduce communication latency between chips in the system and support more efficient inference. Lohmeyer said this results in 2x lower latency than TPU 8t, even though the system is using the same chip. “Because we’ve reduced the diameter of the network topology, the inference just gets much faster, and that’s together with the hardware level protection for confidential inference,” he said. “You shouldn’t have to trade off between security and performance.” Adding more memory lanes to the highway Increasingly, AI agents need to reason, retrieve information, and respond in parallel. But they also require large context windows to give them all the data they need to consider as they respond, and they have to have enough bandwidth to move that data quickly back and forth. “Those context windows are growing exponentially as agents work across documents and code bases and many other data and contexts,” Lohmeyer noted. “And those concurrent workloads put significant pressure on the memory capacity and bandwidth.” TPU 8ix Max addresses the “memory wall,” building on the same low latency architecture of TPU 8i, but essentially “doubles the silicon per package,” Lohmeyer said. The model has 384 gigabytes of high bandwidth memory (33% more than TPU 8i), giving it the ability to store larger context windows and key value (KV) caches. It has double the memory bandwidth (at 17.2 terabytes per second) and double the on-chip static random access memory (SRAM), at 768 megabytes. “It’s like adding more lanes to the highway,” he observed. “We’re not necessarily making each car travel faster, we’re creating more capacity so more traffic can move through the system at the same time.” Handling bursts, idling However, an agentic workload doesn’t behave exactly like a traditional inference workload. “It is bursty, long running, but often idle,” Lohmeyer noted. While it might generate outputs and call tools, it still has to wait for a response, maybe from a human, “literally for days.” During that time, it doesn’t need to do anything; it just needs to preserve context. But if resources are tied up for that entire multi-day period, they’re being wasted “90% plus of the time,” he said. Google Cloud’s open source TPU Sync is designed to address these issues by unifying memory, storage, and networking in software. The model-state layer for TPUs turns chip high bandwidth memory (HBM), host DRAM, shared storage, and the data center network into a “single, co-designed memory hierarchy,” the company said. It automatically tiers KV caches off the chip in milliseconds, streams context across nodes, and synchronizes RL weights to cut transfer latency and Time to First Token (TTFT) by more than 90%. Effectively, this moves the working memory off the chip when an agent pauses, then quickly restores it when the process resumes, Lohmeyer explained. This helps free up resources for other active workloads, drive down costs for token use and agent transactions, and improve performance at a “much higher effective usage of the underlying hardware.” The software is available for v7, v6, and v5 TPUs as well as for the current v8 TPUs. Addressing key shifts in AI architecture AI is moving beyond individual interactions and simple chat and response, Lohmeyer noted; agents can now take on more and more complex tasks to work toward a broader objective. In doing so, they break workflows into multiple steps: Calling models and tools, retrieving information, and writing, compiling, and executing code. “Dozens or even hundreds of operations are happening behind the scenes, and that has some pretty significant implications for the infrastructure that is powering all of this,” he said. He described three fundamental shifts occurring now in AI architecture. The first is the move from human scale to agentic scale: Agents are spinning off tens or hundreds of sub-agents that are executing in parallel, driving up compute demand by 10x, 100x, or more. The second is a shift from systems-of-record to systems-of-action, and the third is increasingly complex agentic workflows. “Putting all this together, you start to see a very different set of infrastructure challenges,” Lohmeyer said, noting that the conversation used to center around available compute, but now builders want to know how quickly data can move, how efficiently systems can use memory, how latency can be minimized, and how to manage the cost and security of running many workloads at scale. He pointed out that, while the silicon is still important and at the core, so is the surrounding memory, the networking, the storage, and the software that orchestrates all of it, as well as the data centers it sits within. “If any one of those components becomes a bottleneck, it can hinder the whole system performance,” Lohmeyer said.