Ray's new TPU support is aimed at your GPU bill

Ray 2.55 adds official TPU support with atomic gang scheduling for TPU slices, aiming to cut GPU costs by enabling Ray-based workloads to run on Google's TPUs. Google's motive is to boost external TPU demand, following Anthropic's October 2025 deal for up to a million TPU chips reported at tens of billions of dollars.

Cloud & Infra https://sourcefeed.dev/c/cloud Article Ray's new TPU support is aimed at your GPU bill Ray 2.55 gives TPU slices atomic gang scheduling, and Google's motive goes well beyond developer convenience. Emeka Okafor https://sourcefeed.dev/u/emeka okafor The knock on TPUs was never the hardware. Google's chips have posted strong price-performance numbers for years, and interest ticks up every time an Nvidia invoice lands. The knock was everything around the hardware: you couldn't bring your stack. If your infrastructure sits on Ray https://www.ray.io/ , and in the GPU world a lot of training and serving infrastructure does, adopting TPUs https://cloud.google.com/tpu meant leaving that layer behind for xpk scripts, GKE JobSets, or Vertex AI's managed opinions. Ray 2.55 changes that. TPUs are now in Ray's official release pipeline, with prebuilt rayproject/ray: -tpu Docker images carrying the JAX and libtpu stack, and supported APIs across Core, Data, Train, and Serve. The plumbing isn't new KubeRay has had multi-host TPU scaffolding since v1.1.0, usable with Ray 2.32 and later , but until now it was scaffolding you assembled yourself, with a community webhook and containers you built and debugged alone. The difference between "technically possible" and "supported" is the difference between a spike and a platform decision. Why 16 chips isn't 16 chips The interesting engineering here is scheduling, not kernels. A v6e slice with a 4x4 topology is 16 chips spread across four host machines, wired together by the Inter-Chip Interconnect. That fabric only exists inside one physical slice. Ask a naive scheduler for 16 chips and it may hand you eight from one slice and eight from another; your tensor-parallel model then waits on a mesh that isn't there. The job doesn't fail loudly. It hangs, and the meter keeps running. Ray's answer has three layers. On GKE, the Ray operator add-on ships a TPU webhook that labels every worker pod with its physical placement: ray.io/tpu-slice-name , ray.io/tpu-topology , ray.io/tpu-worker-id . Above that, ray.util.tpu exposes slice placement group , which matches on those labels to reserve an entire slice atomically, all hosts or none. And above that, the AI libraries make the call for you, so most users never touch it: trainer = JaxTrainer train loop per worker=train func, scaling config=ScalingConfig num workers=4, use tpu=True, topology="4x4", accelerator type="TPU-V6E", , That topology="4x4" line replaces what used to be a block of hand-written host coordination. Ray Serve applies the same idea to inference: a topology field in the accelerator config gang-schedules multi-host vLLM deployments, and if you omit it, the deployment sits in DEPLOYING forever while TPU-hours burn. Ray Data adds iter jax batches , which feeds workers device-sharded JAX arrays directly instead of NumPy you convert on the hot path. I'd flag one thing for people who don't care about TPUs at all: this scheduling shape is coming for you anyway. An NVL72 rack is a 72-GPU NVLink domain with the same constraint, all workers inside one fabric or nothing. Ray's new label-selector machinery, which on GKE can express fallback chains try a v6e-32, fall back to v5e; try Spot, fall back to reserved capacity , is a general answer to topology-domain scheduling. Google built it because TPU slices forced the issue first. Google's motive is bigger than DX Google needs external TPU demand, badly. In October 2025 Anthropic signed for access to up to a million TPU chips in a deal reported at tens of billions of dollars. Anthropic can afford bespoke infrastructure. The next thousand customers can't, and they all run the same open-source stack: vLLM https://docs.vllm.ai/ for serving, Ray for orchestration, JAX https://docs.jax.dev/ or PyTorch for the model. Notice the direction of the port. Google isn't asking the world to adopt its internal tooling; it's spending engineering effort jointly with Anyscale, Ray's steward to make the GPU world's default stack run unmodified on its chips. vLLM's TPU backend came first, and Ray is the scheduling layer beneath it. Read "first-class accelerator" as a distribution play: TPU capacity becomes substitutable for GPU capacity without an infrastructure rewrite, which is exactly what a price-performance argument needs to convert. Before you port anything In practice this is a GKE story. The webhook and the operator add-on live there; you can run Ray on TPU VMs directly, but then you're back to assembling parts. If your platform team is on EKS, this release does nothing for you, and that asymmetry is the point of it. The API caveats are real. slice placement group ships at alpha stability, so expect surface churn between releases. The polished training path is JAX-shaped: JaxTrainer handles multi-host initialization for you, while PyTorch on TPU still routes through torch xla with rougher edges. Observability got a real upgrade TensorCore utilization, duty cycle, and HBM usage in the Ray dashboard, plus libtpu logs , but your team's CUDA debugging instincts won't transfer, and that cost is paid in incidents, not docs. So the adoption order writes itself. Inference first: if you're already on Ray Serve with vLLM and your serving bill is the pain, benchmark your actual model on a v6e slice, because that path is the most finished and the savings are measurable per token. Data preprocessing second, since iter jax batches mostly removes the format tax. Training last, and only if you're a JAX shop or willing to become one. My read: genuine shift, narrow beneficiary. Nothing here moves a CUDA-kernel codebase, and Nvidia loses no sleep this quarter. But for the specific reader who runs Ray, serves open-weight models, and winces at the invoice, the ergonomics excuse for ignoring TPUs was the top excuse, and it's mostly gone. When the scheduler stops caring which vendor's mesh it's placing onto, price-performance gets to speak for itself. That's new. Sources & further reading 1. Run Ray on TPU, Part 1: The foundations https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/ — developers.googleblog.com 2. Run Ray on TPU, Part 2: Ray AI libraries https://developers.googleblog.com/run-ray-on-tpu-part-2-ray-ai-libraries/ — developers.googleblog.com 3. Ray on TPUs with GKE: a more native experience https://cloud.google.com/blog/products/containers-kubernetes/ray-on-tpus-with-gke-a-more-native-experience — cloud.google.com 4. Use TPUs with KubeRay https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html — docs.ray.io 5. Anthropic to Expand Use of Google Cloud TPUs and Services https://www.googlecloudpresscorner.com/2025-10-23-Anthropic-to-Expand-Use-of-Google-Cloud-TPUs-and-Services — googlecloudpresscorner.com Emeka Okafor https://sourcefeed.dev/u/emeka okafor · Security Editor Emeka has spent over a decade tracking threat actors, vulnerability disclosures, and the evolving landscape of application security, bringing a sharp continent-spanning perspective to his reporting. He's known for translating dense CVE advisories into clear, actionable context that developers and security teams alike actually read. Discussion 0 No comments yet Be the first to weigh in.