软件栈
函数化(functionalization)
也称为 functionalization、函数化、显式状态传递、explicit state
先用大白话
函数化把原地修改对象的代码改写成以状态为输入、返回新状态的代码,这样编译器才能追踪它。
技术定义
函数化是一种程序变换,把对 Python 对象的原地修改(例如写入 KV cache)转换为纯函数的显式输入和输出,使 jax.jit 这类追踪式编译器能把程序捕获为计算图。
工程细节
vLLM 管理着每一步都在变化的状态,最重要的是 KV cache。在 TorchAX 路径中,模型包装器把权重和 KV cache 作为显式状态交给 JAX:每个推理步骤接收旧 cache,返回更新后的 cache。jax.jit 因此可以只捕获一次,然后反复执行编译好的图。PyTorch 在 AOTAutograd 内部也有自己的函数化 pass,为 torch.compile 提供同样的作用,这是 TorchTPU 路径能复用编译管线的原因之一。
为什么重要
这一变换在正常工作时不可见,出问题时却很棘手。服务引擎在编写时假定 Python 对象可变,它们与函数式编译器之间的不匹配是翻译层问题之一,正是这类问题推动 Google 从 TorchAX 转向原生 TorchTPU 后端。
如何在 InferenceX 中解读
TPU InferenceX 预览把函数化描述为此前 tpu-inference 后端把 vLLM 搬到 JAX 上的手段之一。TorchTPU 路径用 PyTorch dispatch 取代了这种手写的状态传递,因此 Qwen3.5 397B 的结果依赖编译器在内部处理它。