推論基礎設施
vLLM Supports 16K Embeddings on Ironwood TPUs, StepPool Preserves Chunked Prefill State
Google and vLLM have added tensor alignment, compilation warmup, and pooling state that accumulates across chunks for Qwen3 embedding models. Official results report roughly 84,000 tokens/s on four Ironwood TPUs, but no cost-equivalent GPU comparison or real-world retrieval quality testing has been provided.
Google Cloud has published vLLM-TPU support for long-context embeddings along with a reproducible deployment recipe, targeting the text-only Qwen3-Embedding-8B and multimodal Qwen3-VL-Embedding-8B models. Rather than introducing new model weights, the work enables a serving engine originally centered on autoregressive generation to reliably handle pooling, 16K-scale inputs, and multimodal prefill on Ironwood TPUs.
The first change addresses the tensor divisibility constraints of TPU MXUs. When the model’s vocabulary matrix is sharded with tensor parallelism, vLLM adds hardware-safe padding to prevent All-Gather failures caused by misaligned shapes. The second change improves lazy loading and warms up a shard-aware JAX/XLA compilation cache before the service accepts requests, reducing the risk that compilation will be triggered only by the first batch of requests in a multi-process deployment.
More importantly, the update introduces a hybrid StepPool. To avoid exhausting HBM during prefill, long sequences must be divided into multiple chunks. Embeddings, however, require pooling across the entire sequence. If intermediate aggregation state exists only for a single step, it may be lost when a request is preempted or resumed. The new version moves the relevant metadata into `CachedRequestState`, allowing aggregation state to accumulate across chunks and persist after preemption. The multimodal implementation currently chunks only the text portion, leaving the image path as a capacity boundary that still requires monitoring.
The official recipe uses a single 2×2×1 Ironwood node with four TPUs, BF16, TP=4, and random 16,384-token inputs. Google’s blog reports 83,996 tokens/s and 5.13 requests/s, while the GitHub recipe table lists 84,038.27 tokens/s and the same 5.13 requests/s. The slight discrepancy suggests that the figures may come from different benchmark runs. Accuracy validation compares vectors produced by the TPU and a reference backend using cosine similarity, with thresholds of 0.999 for text and 0.995 for multimodal inputs. This verifies numerical consistency between backends, but it does not demonstrate superior retrieval quality, tail latency, or cost compared with GPUs. Engineering teams should rerun the benchmarks using their own sequence-length distributions, text-to-image ratios, and GKE scaling behavior.