One of the interesting talks for me at Ray Summit was vLLM on TPU with TorchTPU and Ray. Layering should be obvious
Ray → vLLM → PyTorch → TorchTPU → TPU
Ray handles placement and distributed orchestration. vLLM handles scheduling, continuous batching and KV-cache management. TorchTPU tries to make TPU look like another native PyTorch accelerator underneath that stack.
A few things stood out.
1. Fused Eager
PyTorch developers expect relatively dynamic, eager execution. XLA gets much better performance when it sees larger chunks of computation and optimize them together.
TorchTPU’s Fused Eager approach tries to bridge the two. Instead of sending every PyTorch op independently to the TPU, it delays execution long enough to identify useful groups of operations, fuse them and then give XLA a larger region to optimize.
So something like:
matmul → bias → activation → multiply
can become a much larger fused computation on the TPU, while the developer still gets a PyTorch-like programming model.
That feels especially important for inference frameworks like vLLM, where forcing everything into a static execution model would be difficult.
2. Dynamic shapes
While serving batch sizes keep changing, sequence lengths keep changing, requests enter and leave continuously. You cannot afford to recompile every time the shape changes.
TorchTPU is actively building out bounded dynamism and dynamic-shape support so vLLM-style serving can stay dynamic without paying a constant recompilation penalty.
This sounds like a compiler detail, but it is really a serving-system requirement.
3. You still need highly optimized kernels for attention, MoE, quantization and other hot paths.
The emerging TPU stack is interesting:
PyTorch → Helion (abstraction) → Pallas (low level kernel programming) → TPU
alongside:
PyTorch → TorchTPU → XLA → TPU
This starts to give the PyTorch TPU ecosystem something closer to what Triton has done for GPUs: keep the programming model relatively high level, while still having a path to accelerator-specific kernels where performance matters.
4. Ray is solving a different layer of the problem
TPUs placement and topology matter because devices are organized into slices with high-speed ICI connectivity.
TorchTPU is also being built out for the techniques that matter for large models: tensor, pipeline, data and expert parallelism, along with FP8 and lower-precision execution.
To watch in the roadmap ahead:
Fused Eager and bounded dynamism need to work well under highly dynamic vLLM workloads. Kernel coverage has to catch up across attention, MoE, quantization and speculative decoding. And operationally, moving from GPU to TPU or vice versa should not require rebuilding observability, profiling, autoscaling and debugging practices from scratch.