Visko 如何在 Google Cloud TPU 上流式传输实时世界模型

Editor's note: Generative AI startups are moving rapidly beyond static clip generation to build interactive, real-time worlds. In this technical deep dive, we explore how AI startup Visko collaborated with Google Cloud TPU engineers to deploy Orbis — its continuously streaming Live World Model—on Cloud TPUs, going from initial kickoff to production handover in just 30 days.

---

What makes a live video world model so challenging to serve in production? The bottleneck isn't average throughput — it's beating a relentless real-time clock. Unlike offline video generators where a user submits a prompt and waits in a queue, a live world model streams continuously. Every new chunk of video and synchronized audio must finish rendering before the viewer finishes watching the previous one, or the world freezes.

Today, we're excited to share how Visko, an AI startup pioneering interactive world models, teamed up with Google Cloud to solve this exact challenge. Working side by side, Visko and Google Cloud engineers deployed Visko's end-to-end live inference pipeline — combining video generation, super-resolution upscaling, and synchronized audio—onto a single Cloud TPU slice. Together, the teams brought end-to-end generation down, rendering at 1.07x playback speed.

Here is an inside look at how we engineered the pipeline to hit that milestone in a single month, the architectural optimizations that moved the needle, and practical lessons for other startups building ultra-low-latency generative AI on Google Cloud.

Redefining interactive media with Live World Models

While most generative video platforms focus on rendering fixed, pre-baked clips, Visko is building what it calls Live Models: video world models that generate continuously and respond on the fly.

Visko's flagship model, Orbis, streams live video starting from a single text or image prompt. As the stream plays, viewers can enter new steering prompts in real time, dynamically evolving the scene without restarting or breaking continuity. For instance, in one test session, the livestream opens on a cat napping beside a crackling fireplace; when the viewer steers the scene to introduce a bowl of food, the bowl materializes seamlessly, and the cat stretches, stands up, and meows.

Since launching its public demo on September 1, 2026, Orbis has unlocked new possibilities across interactive entertainment, live commerce, AI companions, immersive education, and robotics simulation. To scale these interactive sessions reliably, Visko partnered with Google Cloud TPU engineers in August 2026 to port the Orbis live pipeline to Google's Tensor Processing Unit (TPU). Within 30 days, the full multimodal pipeline was running in real time on TPUs.

Visko’s ambition is to power a new generation of interactive worlds that people can step into and shape in real time. Streaming continuous video and audio under a strict clock is a fundamentally different system challenge than offline clip generation. Working side by side with Google Cloud's TPU engineers, we brought the full Orbis pipeline to TPUs in just 30 days — giving us the low-latency compute and memory bandwidth required to bring Live Models to creators and developers at scale.

"Visko is bringing an exciting new class of real-time, steerable workloads to the Cloud TPU ecosystem. By co-engineering their live multimodal pipeline on a singleTPU slice using JAX, torchax, and custom Pallas kernels, we proved how quickly a lean startup team can go from PyTorch prototype to render at 1.07x playback speed on Google Cloud's AI Hypercomputer."
– Dave Elliott, Head of Developer Advocacy for AI, Google Cloud 

Live video is a hard-deadline problem

When you serve an offline clip generator, latency is forgiving — the user waits once while the backend processes frames. A live world model has no waiting room.

Every single chunk must be denoised, decoded from latents, upscaled to 1664x936 (at 18 fps), and paired with generated audio faster than realtime playback speed, for as long as the session runs.

Two structural constraints make live world models far more demanding than standard batch inference:

• Batch size of one: Each live stream represents an individual viewer's interactive session. You cannot mask latency by grouping dozens of concurrent requests into a massive batch — the standard playbook for keeping AI accelerators saturated.

• Strict sequential dependency: Every chunk depends directly on the one before it. Because the video model maintains a rolling memory of recent video latents so that chunk N+1 smoothly continues the action from chunk N, you cannot pre-render chunks out of order or parallelize across future timestamps.

For startup infrastructure engineers coming from LLM serving, it helps to think of live world modeling not as a tokens-per-second throughput benchmark, but as an unforgiving per-step latency SLO.

Architecting the multimodal pipeline on a single TPU slice

To keep infrastructure lean and cost-effective as Visko scales, the team designed the live pipeline to chain three distinct generative models on a single Cloud TPU slice running in BF16 precision.

When a viewer submits an initial prompt or a live steering update, all three generative models — the core video diffusion model, the super-resolution upscaler, and the audio generator — share the same chips, taking turns on the hardware for each chunk.

To combine rapid startup iteration with bare-metal accelerator speed, the software stack is built on three open-source pillars:

• JAX: Google's high-performance Python array library for accelerators, compiled into optimized machine code by the XLA compiler.

• torchax: An interoperability layer that runs PyTorch models directly on Cloud TPUs via JAX, enabling startup teams to bring existing PyTorch components over without an upfront, ground-up rewrite.

• Pallas: The custom kernel programming language for TPUs in JAX — the TPU counterpart to writing custom Triton or CUDA kernels on GPUs.

Real-time performance breakdown: Hitting the clock 

By profiling and optimizing each stage on a single slice in BF16, the engineering team achieved a total per-chunk latency of under the target while — maintaining roughly 7% headroom to absorb variance during live streaming:

Because video denoising accounts for nearly two-thirds of the total compute budget, that stage became the primary focus of our kernel and parallelism optimizations.

Inside the stack: Three optimizations that unlocked real-time streaming

1. Ulysses context parallelism across all TPU chips

A video chunk expands into a massive sequence of spatiotemporal tokens inside the video transformer. To distribute that workload across the TPU chips at batch size 1, the team implemented Ulysses context parallelism.

Instead of splitting model weights, Ulysses splits the token sequence across chips. Immediately before the attention layer, a high-speed all-to-all interconnect exchange regroups the tensors so that each chip holds the entire video sequence, but only for one-quarter of the attention heads. After each chip computes attention over the full sequence for its assigned heads, a second all-to-all exchange restores the sequence-split layout for the feed-forward layers.

2. Custom Pallas splash attention from MaxDiffusion

To accelerate the attention computation itself, the team integrated a custom Pallas splash attention kernel from MaxDiffusion, Google Cloud's open-source reference library for diffusion models on Cloud TPU.

Splash attention brings FlashAttention-style tiling to TPU architecture: rather than materializing the full quadratic attention matrix in High Bandwidth Memory (HBM), it computes attention in compact tiles that remain inside ultra-fast on-chip VMEM. Startup teams can drop the exact kernel used in this port straight into their own pipelines via custom_splash_attention.py.

3. A 5.8x audio speedup by moving from torchax to pure JAX

We first brought all three Orbis models onto Cloud TPU with *torchax*, which runs PyTorch code on TPU through JAX and XLA without a rewrite. For the video and upscaling models, `torchax` paired with Pallas attention kernels met our latency target directly, and both stay in `torchax` in the final pipeline.

Only the audio model needed a rewrite. Moving the audio stage from `torchax` to *native JAX* made it *5.8x faster* and brought the full pipeline under the playback deadline.

Startup Builder Tip: Port strategically.
The takeaway isn't "skip torchax" — it's to use it as a fast bring-up bridge. torchax got our PyTorch models running on TPU in days rather than weeks, giving us the end-to-end telemetry needed to pinpoint true bottlenecks. Measure the full pipeline first, and invest engineering cycles rewriting only the components that dominate your latency budget.

Field notes for AI builders: Cloud TPU lessons

Takeaways from tuning live multimodal inference on Cloud TPU:

- *Ulysses context parallelism for single stream video:* At batch size 1, splitting the token sequence across chips with `all-to-all` interconnect exchanges was faster than both tensor parallelism and ring attention.

- *Native BF16 for continuous streams:* FP8 quantization cuts memory bandwidth pressure, but small quantization errors build up across chunks in a continuous stream. Native BF16 kept visual quality intact within the latency budget.

- *Selective JAX ports:* `torchax` kept the video and upscaling models in PyTorch, so the team only rewrote the audio bottleneck in native JAX.

- *XLA compilation caching:* The first boot spends a few minutes compiling XLA graphs. Caching compiled artifacts brings worker restarts under a minute.

Building real-time generative models on Cloud TPU? Start here

If your startup is pushing the boundaries of live video, world models, or ultra-low-latency multimodal inference, you can build with the same toolkit that powered Visko's port:

• Accelerate diffusion attention out of the box: Grab the open-source splash attention kernel used in this collaboration from MaxDiffusion, and explore the Pallas documentation to write custom TPU kernels.

• Migrate PyTorch fast, then optimize selectively: Bring your existing PyTorch models onto TPU quickly with torchax, pinpoint bottlenecks with the Cloud TPU Profiler, and port latency-critical stages to native JAX.

• Spin up TPUs on demand: Use TPU Flex-start VMs to provision TPU capacity for short development and benchmarking runs without a long-term reservation.

• Scale with Google for Startups: Early-stage AI startups looking for cloud credits, technical mentorship, and access to Google Cloud's AI Hypercomputer stack can apply to the Google for Startups Cloud Program.

---

^{Special thanks to the Visko.ai team for the close technical partnership, and to the Google Cloud TPU engineers who lead the project including: Yixin Shi, James Huang, and Elisa Tsai.}