
ZenteiQ借助GKE和llm-d将物理AI推理提升4倍

ZenteiQ是一家从印度科学理工学院(IISc)剥离的AI初创公司,入选IndiaAI Mission,目标是构建自主的科学基础模型。这类模型不同于用网页文本和代码训练的通用模型,它要理解热力学、计算流体力学、电磁学等领域里的物理规律、控制方程和数值求解器,还要能做符合物理定律的逐步推导。预训练这样的模型需要数千亿推理token,所以ZenteiQ除了使用精选文献和仿真数据集,还自己设计了一条大规模批量合成数据管道,生成、评估、过滤候选推导和求解器验证问题,再把验证过的数据喂给MaxText预训练任务。
ZenteiQ的算力配额固定,加速器花在生成数据上的时间就是预训练损失的时间;数据生成一慢,多主机加速器集群只能空转等待。管道把提示准备和服务拆开:工作pod负责扩展提示,vLLM在多节点加速器池上跑推理。每个批次经过四个阶段——把种子概念扩成提示批次、用生成器模型生成推理轨迹、用评估器模型评分、再用独立模型家族做guardrail检查,最后打包。最初的路由是Kubernetes标准的ClusterIP轮询,规模扩展到9B和35B参数模型后,两个瓶颈开始拖慢数据生成。
第一个瓶颈是队列不均:提示复杂度和响应长度差异很大,轮询看不到vLLM引擎队列深度,新请求继续发给正忙的pod,旁边节点却闲着。第二个瓶颈是冗余前缀计算:同一批提示约70%到80%是共享指令,轮询把相同指令的请求打散到不同实例,每个pod都重算同样的前缀,不复用加速器里的缓存。结果一次运行里吞吐波动最高达到一半,预训练开始时间被不断推迟。
解决方案是在vLLM加速器池前部署GKE Inference Gateway,它由开源分布式推理框架llm-d驱动,不改变底层科学方法论和管道设计。所有模型家族的生成请求、评估查询、guardrail检查都通过同一个网关路由到对应模型池。路由决策不再基于静态轮询,而是持续监控每个vLLM端点的实时队列深度和KV缓存利用率:节点在解码长推导时,新请求被转给有空闲容量的节点;匹配科学指令前缀的请求,则路由到该前缀缓存已经热在加速器内存里的节点。这样一来,pod既跳过重复prefill,又不会压垮单个副本。
结果是生产集群的聚合吞吐从每1000个加速器每秒16,964 tokens提高到69,114 tokens,约4.07倍;每日生成token从14.7亿增加到59.7亿。单个模型池的收益在3.5倍到5.5倍之间,前缀缓存命中率提高14个百分点。对团队来说,最重要的不是数字本身,而是预训练计划变得可预测:数据集按时准备好,多主机预训练任务按计划启动。
下一步,ZenteiQ要把模型从9B扩展到35B参数以上,计划沿用同一套GKE Inference Gateway和llm-d配置,用于后训练强化学习(RL)和自动化数值求解器评估循环。这类多轮rollout每一轮都依赖上一轮的上下文,对实时队列平衡和热前缀缓存的需求更高。


