TPU Inference v0.31.0, released by the vLLM project on October 8, adds hierarchical dispatch to reduce MoE communication and compressed-tensors W4A16 linear operations on the JAX path.
The former reduces data movement for expert models spread across chips; the latter fits models with low-bit weights into memory. Hierarchical dispatch can now also work with FP8 all-gather.
The two features target different bottlenecks. In MoE models, which activate only some experts for each input, time goes not only to computation but also to inter-chip communication that sends data to experts.
For Gemma 4 31B, the starting question was whether the model would fit in memory. The W4A16 development PR covers QAT models, whose training accounts for quantization error.
According to that PR, the bf16 model, which uses a 16-bit floating-point format, did not fit on one v6e chip.
Two cores sharing inter-chip transfers
With hierarchical dispatch, each of a chip's two cores gathers half the input features from other chips, then exchanges the rest within the chip. The PR author says this halves inter-chip transfer volume.
Support covers the grouped matrix multiplication (GMM) path for expert parallelism (EP), which distributes experts across devices. It requires paired cores from the same chip along the model axis and a hidden width divisible by 256.
The setting is MOE_HIERARCHICAL_DISPATCH, and it is disabled by default. Unsupported configurations fall back to the existing path.
The PR author measured Qwen3.8-2.4T-A95B-FP8 on v7x-32, with tp=4 to split tensor computation and attn_dp=8 to distribute attention requests. During sequential token decoding, dispatch time fell from 11.7ms to 7.2ms per step.
Total decoding step time, however, fell from 89.5ms to 85.5ms, a 4.5% reduction. These are the PR author's profiling results, not an independent verification of the internal traces linked from the public description.
The FP8 all-gather compatibility PR provides a separate comparison. With FP8 used to gather data across chips, decoding time on v7x-32 at batch size 384 fell from 77.38ms to 74.10ms, a 4.2% reduction.
That comparison takes its baseline from PR #3665. It differs from the baseline for the earlier 4.5% result, so the two improvements cannot be added together.
The same PR reports a slight drop in the GSM8K math evaluation score, from 0.978 to 0.9773. A claim of bitwise-identical outputs for the data movement path does not establish a quality guarantee when combined with FP8.
The scope of 4-bit weights that fit 31B
W4A16 stores weights at 4 bits and handles activations passed between layers at 16 bits. This JAX implementation supports symmetric int4 weights in the compressed-tensors pack-quantized format, without separately quantizing activations.
It supports group-wise and channel-wise quantization. Weights remain int4 in high-bandwidth memory (HBM), and the gmm_v2 kernel reconstructs only the tiles needed for computation in on-chip working memory.
Asymmetric formats, 8-bit formats and actorder=group are not supported.
Linear operation support is separate from int4 MoE expert support. In this PR, compressed-tensors MoE quantization formats other than FP8 produce an error when loading the model.
The PR author reported that Gemma 4 31B QAT W4A16 used 21.67 GiB of HBM on one v6e chip and delivered output throughput of 500 tokens/second.
The test used 16 concurrent requests, a fixed output of 256 tokens per request and ignore_eos to ignore end-of-sequence tokens. Maximum model length was 2048 and maximum concurrent sequences was 16; results are medians of three runs.
However, the measurement image includes vLLM 0.29.1rc1.dev468. These are development PR results supporting the feature, not fresh measurements of the final v0.31.0 release.
Fitting in memory does not mean preserving quality. In the PR author's 3,880-case public label-reading evaluation, bf16 → W4A16 scores were E2B 68.5% → 65.7%, E4B 73.1% → 71.4% and 12B 76.0% → 75.0%.
The author reported differences of −2.8, −1.6 and −1.0 percentage points, respectively. The E4B difference is the original reported value, not a recalculation from the rounded scores.
The 31B W4A16 score was 77.6%, but there is no bf16 score for comparison. The evaluation's composition and reproducibility have not been independently audited, so these figures cannot establish whether 31B preserves quality.
Reported by the PR author. Compared record by record on the 3,880-record public suite against the bf16 checkpoint served on the same vLLM image. Parentheses show the 95% bootstrap range.
| Model | bf16 | W4A16 | Difference (points) |
|---|---|---|---|
| E2B | 68.5% | 65.7% | −2.8 (−4.0 to −1.6) |
| E4B | 73.1% | 71.4% | −1.6 (−2.5 to −0.9) |
| 12B | 76.0% | 75.0% | −1.0 (−1.7 to −0.3) |
| 31B | — | 77.6% | — |
Startup time is another cost. The PR reports compilation taking about 2.7 times as long as the XLA path. For 31B, initial compilation took about 27 minutes, or about 9 minutes with the JAX compilation cache retained.
The author also observed that bf16 was faster for small projection operations. All W4A16 serving measurements used one v6e chip; multi-chip tensor parallelism with TP>1 and v5e were not tested.
For MoE, the important figure is the 4.5% reduction in total step time, rather than the halving of communication volume. Other computation remains even after data movement falls, so there is no basis for expecting an equal throughput gain.
The 31B W4A16 result is more about feasibility than a speed contest. Since bf16 did not fit on the same chip, 500 tokens/second shows achievable throughput under the stated concurrency conditions, not a speedup over bf16.
MoE operators therefore have separate questions from operators constrained by memory. The former can assess communication bottlenecks and parallel configurations; the latter, checkpoint compatibility, quality and startup costs.
Dispatch requirements and checkpoint formats
Assessing hierarchical dispatch involves more than checking whether MOE_HIERARCHICAL_DISPATCH is enabled. The EP GMM path, paired-core placement on the model axis and hidden width divisible by 256 also matter.
The documented fallback to the existing path outside those conditions provides the basis for this check.
For W4A16, the first checks are symmetric int4 weights, the pack-quantized checkpoint format and unsupported options. Group widths below the chip's matrix multiplication unit (MXU) column width use gmm_v2.
Larger or equal group widths, and channel-wise quantization, use the XLA quantized matrix multiplication path. That path errors on layers whose input dimension is split across devices because it cannot partition the quantization scales.
However, the public evidence does not support providing commands or log strings to determine whether the features are active. These checks are limited to comparing documented support conditions with a specific configuration.
The next questions are hierarchical dispatch's effect under real serving loads and W4A16's multi-chip scalability. Current public evidence supports communication improvements in specific configurations and model fit on a single chip.
The value of v0.31.0 is not in making every model faster, but in expanding the options for reducing MoE communication bottlenecks and model memory bottlenecks within supported configurations.