Functionalization
Also known as functionalization, functional state threading, explicit state
In plain English
Functionalization rewrites code that mutates objects in place into code that takes state as input and returns new state, so a compiler can trace it.
Technical definition
Functionalization is the transformation that converts in-place mutations of Python objects, such as writes into a KV cache, into explicit inputs and outputs of a pure function so that a tracing compiler like jax.jit can capture the program as a graph.
Engineering details
vLLM manages state that changes every step, most importantly the KV cache. In the TorchAX path, the model wrapper presents weights and KV cache to JAX as explicit state; each inference step takes the old cache in and returns the updated cache. jax.jit can then capture the step once and replay the compiled graph. PyTorch has its own functionalization pass inside AOTAutograd that serves the same purpose for torch.compile, which is one reason the TorchTPU path can reuse the compile pipeline.
Why it matters
The transformation is invisible when it works and painful when it does not. Serving engines were written assuming mutable Python objects, and the mismatch with a functional compiler is one of the translation-layer problems that pushed Google from TorchAX toward the native TorchTPU backend.
How to read it in InferenceX
The TPU InferenceX preview describes functionalization as part of how the previous tpu-inference backend got vLLM onto JAX. The TorchTPU path replaces that hand-written state threading with PyTorch dispatch, so the Qwen3.5 397B results depend on the compiler handling it internally.