Optimización de la Inferencia de LLMs para Recomendaciones Generativas en Spotify con TensorRT-LLM

La personalización a escala en sistemas de recomendación ha evolucionado desde enfoques basados en recuperación y ranking a modelos generativos basados en LLMs. Este cambio fundamental introduce nuevos desafíos en la inferencia, particularmente en la eficiencia y la latencia para aplicaciones en tiempo real. La tesis central es que, si bien los LLMs ofrecen una capacidad de personalización sin precedentes, su despliegue en producción requiere una optimización profunda de la pila de inferencia para superar las limitaciones de rendimiento de las implementaciones iniciales y lograr una operación rentable a escala.

Tradicionalmente, los sistemas de recomendación se basan en pipelines multifase de recuperación de candidatos y ranking. Los LLMs generativos, como los explorados por Spotify, prometen un enfoque unificado y más contextualizado. Sin embargo, la complejidad computacional de la inferencia de LLMs, especialmente con técnicas como beam search y la inyección de personalización, puede llevar a cuellos de botella de CPU, subutilización de GPU y latencias inaceptables. La necesidad de integrar estas capacidades de manera eficiente en un único motor de inferencia es crítica para la viabilidad de los sistemas de recomendación generativa en entornos de producción de alto volumen.

Arquitectura del Sistema

La arquitectura de inferencia de Spotify para recomendaciones generativas se basa en LLMs especializados, inspirados en el paper PLUM. Inicialmente, se utilizó vLLM como motor de inferencia, destacando por su alto throughput, pero con limitaciones en la implementación de beam search, que operaba fuera del motor principal de GPU, generando cuellos de botella en la CPU y recomputaciones de bloques KV. Para el scoring de candidatos, se exploró SGLang, que ofrece una técnica eficiente de 'candidate masking' para evaluar múltiples candidatos en una sola solicitud, evitando el enfoque ingenuo de 'fan-out'.

La personalización se logra mediante 'Soft Tokens'. Esta técnica, detallada en el white paper 'Embedding to Prefix' de Spotify Research, proyecta embeddings de usuario existentes en el espacio de embeddings del LLM como un único 'soft token'. El LLM es pre-entrenado con un subconjunto de pesos para interpretar este token. En tiempo de inferencia, el vector de usuario se proyecta e inserta en el prompt antes del forward pass inicial. Otras aproximaciones como 'Plain User Text' (historial de usuario como texto) y 'User Semantic IDs' (extensión del vocabulario del modelo con tokens de ID de usuario) fueron exploradas, pero descartadas o consideradas para investigación futura debido a la latencia o la complejidad.

La colaboración con NVIDIA se centró en unificar estas funcionalidades en TensorRT-LLM (TRT-LLM). TRT-LLM ofrecía una superior utilización de hardware para beam search, pero carecía de la API de scoring de SGLang y la capacidad de inyectar soft tokens y devolver múltiples secuencias en beam search de vLLM. Se realizaron contribuciones a TRT-LLM, incluyendo soporte para múltiples secuencias en beam search, modelos 'out-of-tree' en trtllm-serve, paso de embeddings multimodales vía 'image_embeds' y scoring multi-ítem vía la API de Python. La implementación de beam search en TRT-LLM se integró en el sampler, reduciendo overheads del host y mejorando la utilización de GPU mediante 'overlap scheduling'. Para los soft tokens, se implementó un plugin de extensión Qwen3 personalizado utilizando el soporte de modelos 'out-of-tree' de TRT-LLM, que realiza la proyección del embedding de usuario y el reemplazo del soft token.

Fundamentos Teóricos

El problema de la personalización en sistemas de recomendación ha sido un pilar de la investigación en recuperación de información y aprendizaje automático durante décadas. La transición a modelos generativos se alinea con avances recientes en Large Language Models (LLMs) y su aplicación a tareas más allá del procesamiento de lenguaje natural puro. El artículo menciona el paper PLUM (Personalized Language Understanding Model) como una base para las arquitecturas de recomendación generativa, lo que sugiere una conexión con la investigación en modelos secuenciales y auto-atención para la generación de ítems relevantes.

La técnica de 'beam search' tiene raíces profundas en la inteligencia artificial, utilizada en algoritmos de búsqueda heurística desde la década de 1970 (Lowerre, 1976 para reconocimiento de voz). Su aplicación en LLMs para la generación de secuencias de tokens es una extensión directa de su uso en tareas como la traducción automática y la síntesis de voz, donde se busca explorar un espacio de búsqueda combinatorio de manera eficiente. La optimización de su implementación en hardware especializado, como las GPUs, es un problema de ingeniería de sistemas que busca cerrar la brecha entre la teoría algorítmica y la eficiencia computacional en la práctica. La idea de 'soft tokens' o 'prefix tuning' para la personalización se conecta con la investigación en adaptación de modelos de lenguaje, donde se busca inyectar información contextual sin reentrenar completamente el modelo, una técnica que ha ganado tracción en la comunidad de NLP (Li & Liang, 2021).