Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU

Summary

Embedding models convert text, images, and audio into dense vectors used for semantic search, recommendation, intent classification, content discovery, and clustering. The post describes production support for long-context Qwen3 Embedding models on Cloud TPUs through native vLLM TPU integration and GKE Custom Compute Classes, allowing prioritized autoscaling across TPU and GPU capacity pools. Its implementation pads vocabulary tensors for hardware-safe sharding, hardens lazy loading, pre-warms sharding-aware JAX/XLA compilation caches, and uses a hybrid StepPool with CachedRequestState to preserve pooling state during chunked prefill and preemption. Parity evaluations target cosine similarity of at least 0.999 for text and 0.995 for multimodal inputs; on TPU Ironwood, Qwen3-Embedding-8B reached 83,996 total tokens per second and 5.13 requests per second at 16K+ sequence length with bf16 and TP=4.

Takeaways

1. 1

   GKE Custom Compute Classes can apply priority rules to scale across TPU capacity and secondary GPU spot or on-demand pools when the primary TPU reservations are fully utilized.

2. 2

   The TPU implementation uses vocabulary padding for tensor alignment, promotes attributes during unquantization for lazy loading, and pre-warms sharding-aware JAX/XLA compilation caches.

3. 3

   The stated parity thresholds are at least 0.999 cosine similarity for text and 0.995 for multimodal inputs; Qwen3-Embedding-8B reached 83,996 total token/s and 5.13 req/s on TPU Ironwood with bf16, TP=4, and 16K+ sequence length.