Google Cloud engineers detail how they brought native TPU support to vLLM for serving the Qwen3 Embedding model family, targeting long-context text and multimodal embedding workloads. Key optimizations include hardware-safe vocabulary padding for tensor parallelism, attribute promotion for lazy weight loading, sharding-aware JIT pre-warming, and a hybrid StepPool mechanism to handle chunked prefill in pooling layers. The post includes a code sample for initializing Qwen3-Embedding-8B via vLLM's pooling runner, cosine-similarity parity thresholds (≥0.999 text, ≥0.995 multimodal) for validating TPU output against XPU references, and links to open-sourced GKE deployment recipes on GitHub.
Table of contents
What is an Embedding Model and What is it Used For?Seamless Elasticity with vLLM-TPU & GKEEngineering High-Precision Embedding Support on TPUsGolden-Reference Precision & Numerical ParityPublic Recipes & Explore ResourcesAcknowledgmentsQuestions this post answers
How do I initialize Qwen3-Embedding-8B with vLLM on a TPU?
Use vLLM's pooling runner with the model set to Qwen/Qwen3-Embedding-8B, runner set to 'pooling', tensor_parallel_size set to 2 for sharding across the TPU topology mesh, max_model_len of 16384, max_num_batched_tokens of 512, dtype bfloat16, and trust_remote_code enabled. Call llm.embed(prompts) to extract dense vector embeddings from the output. daily.dev surfaces deployment recipes like this for teams scaling embedding inference on TPUs.
What cosine similarity threshold indicates acceptable numerical parity between TPU and XPU embedding outputs?
A cosine similarity of at least 0.999 for text embeddings and at least 0.995 for multimodal embeddings is considered a passing quality threshold when comparing TPU-generated vectors against XPU golden reference vectors. Scores approaching 1.0 confirm that hardware-specific optimizations preserve mathematical parity without sacrificing accuracy. developers validating cross-hardware embedding accuracy can track parity benchmarks like these on daily.dev.
Why does chunked prefill in pooling layers risk state loss on long-context embedding inference?
Ultra-long contexts require chunked prefill to avoid exhausting high bandwidth memory, but splitting inference into steps risks losing pooling state across step boundaries. A hybrid StepPool approach that migrates metadata into CachedRequestState ensures pooling states accumulate correctly across steps and survive request preemptions. daily.dev helps engineers keep up with long-context inference techniques like chunked pooling.