Embeddings em Cloud TPU com vLLM como servir Qwen3-Embedding-8B em produção

E aí, dev! Bora falar de embeddings em produção?

Todo mundo ama falar de LLMs, mas quem realmente segura a barra nos bastidores são os modelos de embedding. Eles traduzem texto, imagem e áudio em vetores densos que capturam significado semântico — é por isso que "gato" fica matematicamente perto de "felino" e longe de "sapato".

Fazer um protótipo com um modelo pequeno? Moleza. Agora escalar isso pra milhões de queries por dia com picos de tráfego dinâmico? Aí a conversa muda. Os dois gargalos clássicos:

• Capacidade elástica de aceleradores — você não consegue subir GPUs na velocidade que o tráfego exige.
• Eficiência custo/performance — contextos longos e inputs multimodais torram HBM rapidinho.

A resposta do Google Cloud: colocar suporte nativo a TPU direto no vLLM, o motor de serving open-source que a galera já usa como padrão. Resultado? Uma stack única pra GPU e TPU, com elasticidade real via Custom Compute Classes no GKE.

💡 Se você curte entender como IA está mudando o próprio fluxo de trabalho dos devs, dá uma olhada no nosso texto sobre como o Spotify escalou a experiência de desenvolvedores para times e agentes.

A engenharia por trás: rodar embeddings em TPU não é trivial

Servir modelos de embedding de última geração significa lidar com contextos de 4K+ tokens (texto) e 15K+ tokens (multimodal) — mantendo paridade matemática estrita contra hardware de referência. Olha o que o time fez no vLLM-TPU:

1. Padding de vocabulário hardware-safe

As Matrix Execution Units (MXUs) do TPU impõem restrições rígidas de divisibilidade ao fatiar matrizes de vocabulário via Tensor Parallelism. A solução: estratégia unificada de padding que garante alinhamento exato de tensores no All-Gather.

2. Compatibilidade com o lazy-loader

O vLLM faz lazy-load de pesos no TPU pra reduzir cold-start. Só que transformações lazy quebravam a inicialização. Correção: attribute promotion dentro do pipeline de unquantization, deixando o carregamento 100% compatível com o lazy-loader.

3. Pre-warming ciente de sharding

Compilação JIT em runtime em deploys multi-processo causava compilation traps. Correção: pré-aquecer os caches JAX/XLA antes da inferência, estabilizando os rollouts.

4. StepPool pra Chunked Prefill

Contextos ultra-longos exigem chunked prefill pra não estourar HBM — mas isso arrisca perder estado de pooling entre steps. Correção: um StepPool híbrido com metadados migrados pro CachedRequestState, garantindo acúmulo correto e sobrevivência a preempções.

Exemplo mínimo que funciona

from vllm import LLM

# Inicializa o Qwen3-Embedding-8B em Cloud TPU usando o pooling runner nativo do vLLM
llm = LLM(
    model="Qwen/Qwen3-Embedding-8B",
    runner="pooling",              # Habilita saída de pooling denso
    tensor_parallel_size=2,        # Fatiado pelo mesh de topologia do TPU
    max_model_len=16384,
    max_num_batched_tokens=512,
    dtype="bfloat16",
    trust_remote_code=True
)

# Extrai embeddings vetoriais densos dos inputs
prompts = ["Recuperação semântica enterprise em TPUs com vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding

É isso: uma chamada LLM() com runner="pooling" e você tem embeddings densos em TPU. Os scripts completos estão na recipe oficial Qwen3-Embedding-8B do AI-Hypercomputer no GitHub.

Paridade numérica: o requisito inegociável

Time enterprise não quer embedding "mais ou menos". Quer reprodutibilidade entre hardwares. Se você treinou em XPU e serve em TPU, os vetores têm que bater.

Como se mede paridade

Dado um embedding no TPU (v_Tpu) e um de referência no XPU (v_Ref), calcula-se a similaridade cosseno:

cos_sim(v_Tpu, v_Ref) = (v_Tpu · v_Ref) / (||v_Tpu|| × ||v_Ref||)

Limiares de aprovação

Um score perto de 1.0 significa paridade numérica praticamente perfeita — a stack otimizada pra TPU gera os mesmos vetores que a referência golden, sem perder acurácia.

TPU vs. XPU: trade-offs

⚠️ Limitações pra ficar de olho

• Cold-start de compilação ainda existe; pre-warming mitiga, mas não elimina.
• Restrições de topologia (divisibilidade da MXU) exigem padding em alguns modelos — overhead de memória extra.
• Paridade multimodal (0.995) é um alvo mais frouxo que texto (0.999) — embeddings de imagem ainda são mais difíceis de alinhar.
• Recipes da comunidade ainda são novas; espere arestas em famílias de modelos fora do Qwen.

Se você constrói sistemas agênticos que consomem esses embeddings, vale conferir nosso deep dive sobre o Universal Commerce Protocol (UCP) para comércio agêntico.

O que isso muda pro seu time

Se você roda inferência de embeddings em escala, o jogo mudou:

1. Pare de tratar TPU como nicho. O pooling runner nativo do vLLM significa que o mesmo código que serve LLMs serve embeddings — sem stack separada.
2. Exija métricas de paridade. Peça ao seu vendor (ou ao time de infra) benchmarks de similaridade cosseno contra uma referência golden. Qualquer coisa abaixo de 0.999 pra texto é sinal vermelho.
3. Planeje chunked prefill. Embeddings de contexto longo sem tratamento de estado estilo StepPool vão perder acurácia silenciosamente nas fronteiras de step.

Próximos passos pra estudar

• Comece aqui: Recipe Qwen3-Embedding-8B TPU — scripts de setup prontos pra copiar.
• Vá pro multimodal: Recipe Qwen3-VL-Embedding-8B pra embeddings de texto + imagem.
• Entenda o motor: vllm-project/tpu-inference — código do backend TPU do vLLM.
• Provisione hardware: Google Cloud TPU Portal.

O gargalo mudou de lugar. Não é mais se você consegue servir embeddings — é se consegue fazer isso de forma elástica, barata e com precisão matemática entre hardwares heterogêneos. Esse é o novo padrão.

Fonte: Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU — Google Developers Blog.