ZenteiQ scales physical AI inference 4x with GKE and llm-d

ZenteiQ, an AI startup spun out of the Indian Institute of Science and selected under the IndiaAI Mission, builds sovereign scientific foundation models from scratch. Unlike general-purpose language models trained mainly on web text and code, these models must interpret simulation outputs, evaluate equation formulations, and work through step-by-step mathematical derivations obeying physical laws across domains such as thermodynamics, computational fluid dynamics, and electromagnetics. Pre-training on large datasets of governing equations takes hundreds of billions of reasoning tokens, so ZenteiQ operates a batch synthetic data-generation pipeline that generates, evaluates, filters, and packages candidate derivations and solver validation problems before feeding them into MaxText pre-training jobs on Google Kubernetes Engine (GKE). Because the pipeline runs on a fixed compute allocation, every hour accelerators spend generating and verifying training data is an hour taken from pre-training; when data generation falls behind, multi-host accelerator clusters reserved for pre-training sit idle.

The pipeline separates prompt preparation on worker pods from model serving on vLLM across multi-node accelerator pools. Each batch moves through four stages: expanding scientific seed concepts into prompt batches, generating candidate reasoning traces with generator models, scoring responses using evaluator models, and running guardrail checks with a separate model family before packaging verified datasets for pre-training. Worker pods pull scientific seed documents and domain taxonomies from Cloud Storage and expand them into batches of multi-step prompts. Originally the model pools ran on Kubernetes with standard ClusterIP round-robin load balancing across inference endpoints. As the pipeline scaled to feed 9B- and 35B-parameter model schedules, two routing bottlenecks slowed dataset creation. Round-robin load balancing had no visibility into vLLM engine queues, so it kept sending requests to serving pods already busy processing long reasoning traces while neighboring nodes had spare capacity. In addition, prompts within a generation batch share roughly 70-80% of their instructional context and task specifications, but round-robin scattered requests with identical instructions across different serving instances, forcing each pod to recompute the same shared prompt prefix from scratch instead of reusing warm caches in accelerator memory. Together, uneven queues and redundant prefill work caused generation throughput to swing by as much as half within a single run, stretching dataset creation and delaying scheduled pre-training.

Following a recommendation from Google Cloud to address the routing bottlenecks without changing ZenteiQ‘s scientific methodology or pipeline design, the team deployed GKE Inference Gateway, powered by the open-source llm-d distributed inference framework, in front of its vLLM accelerator pools. A single gateway deployment manages traffic across all model families and serving nodes: rather than maintaining isolated routing for each model tier, generation requests, evaluation queries, and guardrail checks all route through the same gateway to their respective model pools. Routing decisions use live serving signals instead of static round-robin rotation. GKE Inference Gateway continuously monitors real-time queue depth and key-value (KV) cache utilization across vLLM endpoints. When a serving node is busy decoding a long equation derivation, the gateway routes incoming requests to nodes with open capacity. At the same time, requests with matching scientific instructions are routed to the serving node where that prefix’s KV cache is already warm in accelerator memory. Because 70-80% of each prompt is shared instructional context while the generated reasoning trajectories vary, combining prefix-cache locality with live queue balancing lets pods skip repeated prefill work without overloading any single replica.

In production clusters, the change delivered a 4.07x increase in aggregate inference throughput, from 16,964 to 69,114 tokens per second per 1,000 accelerators, growing daily output from 1.47 billion to 5.97 billion generated tokens. Individual model pools saw 3.5x to 5.5x gains, and the prefix-cache hit rate improved by 14 percentage points. For ZenteiQ‘s research and infrastructure teams, the most important outcome was schedule predictability: quadrupling daily output meant verified datasets were ready on time, so multi-host pre-training runs started on schedule instead of waiting on data generation. ZenteiQ plans to use the same GKE Inference Gateway and llm-d routing setup for post-training reinforcement learning and automated numerical solver evaluation loops, where multi-turn rollouts build on previous context and depend even more heavily on live queue balancing and warm prefix caches.

ZenteiQ scales physical AI inference 4x with GKE and llm-d

View Original