Embeddings en Cloud TPU con vLLM cómo servir Qwen3-Embedding-8B a escala enterprise

¡Hola Devs! Hablemos de embeddings en producción

Todos hablan de LLMs, pero los que realmente sostienen la operación son los modelos de embedding. Traducen texto, imágenes y audio a vectores densos que capturan significado — por eso "gato" queda matemáticamente cerca de "felino" y lejos de "sombrero".

Hacer un prototipo con un modelo chico es pan comido. Pero escalar eso a millones de queries al día con picos de tráfico dinámico ya es otra cosa. Los dos cuellos de botella clásicos:

• Capacidad elástica de aceleradores — no puedes levantar GPUs tan rápido como crece el tráfico.
• Eficiencia costo/performance — contextos largos e inputs multimodales se comen la HBM en segundos.

La respuesta de Google Cloud: meter soporte nativo de TPU directo en vLLM, el motor de serving open-source que la industria ya usa como estándar. ¿Resultado? Una sola stack para GPU y TPU, con elasticidad real vía Custom Compute Classes en GKE.

💡 Si te interesa cómo la IA está redefiniendo el flujo de trabajo de los devs, checa nuestro análisis sobre cómo Spotify escaló la experiencia de desarrolladores para equipos y agentes.

La ingeniería: correr embeddings en TPU no es trivial

Servir modelos de embedding de última generación implica manejar contextos de 4K+ tokens (texto) y 15K+ tokens (multimodal) — manteniendo paridad matemática estricta contra hardware de referencia. Mira lo que hizo el equipo en vLLM-TPU:

1. Padding de vocabulario hardware-safe

Las Matrix Execution Units (MXUs) del TPU imponen restricciones estrictas de divisibilidad al fragmentar matrices de vocabulario vía Tensor Parallelism. La solución: una estrategia unificada de padding que garantiza alineación exacta de tensores en el All-Gather.

2. Compatibilidad con el lazy-loader

vLLM hace lazy-load de pesos en TPU para reducir cold-start. Pero las transformaciones lazy rompían la inicialización. Fix: attribute promotion dentro del pipeline de unquantization, dejando la carga 100% compatible con el lazy-loader.

3. Pre-warming consciente de sharding

La compilación JIT en runtime en deploys multi-proceso causaba compilation traps. Fix: precalentar los caches JAX/XLA antes de la inferencia, estabilizando los rollouts.

4. StepPool para Chunked Prefill

Contextos ultra-largos exigen chunked prefill para no reventar la HBM — pero eso arriesga perder estado de pooling entre steps. Fix: un StepPool híbrido con metadata migrada a CachedRequestState, asegurando acumulación correcta y supervivencia a preempciones.

Ejemplo mínimo que funciona

from vllm import LLM

# Inicializa Qwen3-Embedding-8B en Cloud TPU usando el pooling runner nativo de vLLM
llm = LLM(
    model="Qwen/Qwen3-Embedding-8B",
    runner="pooling",              # Habilita salida de pooling denso
    tensor_parallel_size=2,        # Fragmentado sobre el mesh de topología del TPU
    max_model_len=16384,
    max_num_batched_tokens=512,
    dtype="bfloat16",
    trust_remote_code=True
)

# Extrae embeddings vectoriales densos de los inputs
prompts = ["Recuperación semántica enterprise en TPUs con vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding

Ya está: una sola llamada LLM() con runner="pooling" y tienes embeddings densos en TPU. Los scripts completos están en la recipe oficial Qwen3-Embedding-8B de AI-Hypercomputer en GitHub.

Paridad numérica: el requisito no negociable

Los equipos enterprise no quieren embeddings "más o menos". Quieren reproducibilidad entre hardwares. Si entrenaste en XPU y sirves en TPU, los vectores tienen que coincidir.

Cómo se mide la paridad

Dado un embedding en TPU (v_Tpu) y uno de referencia en XPU (v_Ref), se calcula la similitud coseno:

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

Umbrales de aprobación

Un score cercano a 1.0 significa paridad numérica casi perfecta — la stack optimizada para TPU genera los mismos vectores que la referencia golden, sin sacrificar precisión.

TPU vs. XPU: trade-offs

⚠️ Limitaciones a vigilar

• Cold-start de compilación todavía existe; el pre-warming lo mitiga, pero no lo elimina.
• Restricciones de topología (divisibilidad de MXU) obligan a padding en algunos modelos — overhead extra de memoria.
• Paridad multimodal (0.995) es un umbral más flojo que texto (0.999) — los embeddings de imagen siguen siendo más difíciles de alinear.
• Las recipes de la comunidad aún son jóvenes; espera asperezas en familias de modelos fuera de Qwen.

Si construyes sistemas agénticos que consumen estos embeddings, checa nuestro deep dive sobre el Universal Commerce Protocol (UCP) para comercio agéntico.

Qué significa esto para tu equipo

Si corres inferencia de embeddings a escala, el juego cambió:

1. Deja de tratar TPU como nicho. El pooling runner nativo de vLLM significa que el mismo código que sirve LLMs sirve embeddings — sin stack separada.
2. Exige métricas de paridad. Pide a tu vendor (o a tu equipo de infra) benchmarks de similitud coseno contra una referencia golden. Cualquier cosa debajo de 0.999 para texto es bandera roja.
3. Planea chunked prefill. Embeddings de contexto largo sin manejo de estado estilo StepPool van a perder precisión silenciosamente en las fronteras de step.

Próximos pasos para aprender

• Empieza aquí: Recipe Qwen3-Embedding-8B TPU — scripts de setup listos para copiar.
• Pásate a multimodal: Recipe Qwen3-VL-Embedding-8B para embeddings de texto + imagen.
• Entiende el motor: vllm-project/tpu-inference — código del backend TPU de vLLM.
• Provisiona hardware: Google Cloud TPU Portal.

El cuello de botella se movió. Ya no es si puedes servir embeddings — es si puedes hacerlo de forma elástica, barata y con precisión matemática entre hardwares heterogéneos. Ese es el nuevo estándar.

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