Cloudflare Workers AI: Optimizaciones para Inferencia de Modelos LLM a Escala

La proliferación de agentes de IA y la creciente demanda de inferencia de modelos de lenguaje grandes (LLM) están generando una presión significativa sobre los costos y la eficiencia computacional. Los modelos propietarios, aunque potentes, se vuelven prohibitivos a escala de 'agentes personales' o para uso empresarial intensivo. Este artículo aborda el problema fundamental de cómo servir LLMs de frontera, que requieren grandes ventanas de contexto y alta capacidad de razonamiento, de manera costo-efectiva y escalable en una infraestructura distribuida.

La solución propuesta por Cloudflare Workers AI es integrar LLMs open-source de alto rendimiento, como Kimi K2.5, y optimizar su stack de inferencia para reducir el costo por token y mejorar la latencia. Esto se logra mediante técnicas avanzadas de optimización de GPU, como kernels personalizados y paralelización, junto con mejoras en la gestión de la carga de trabajo, como el caching de prefijos y APIs asíncronas. El objetivo es trasladar la complejidad de la ingeniería de ML a la plataforma, permitiendo a los desarrolladores enfocarse en la lógica del agente.

Arquitectura del Sistema

El stack de inferencia de Workers AI se basa en un motor propietario llamado Infire, diseñado para optimizar el rendimiento y la utilización de GPU. Para servir LLMs de gran escala como Kimi K2.5, se han implementado kernels personalizados que mejoran la eficiencia del modelo y la utilización del hardware. Esto va más allá de la ejecución 'out-of-the-box' de modelos open-source, aplicando técnicas de paralelización como data parallelism, tensor parallelism y expert parallelism. También se utiliza la estrategia de 'disaggregated prefill', donde las etapas de prefill (procesamiento de tokens de entrada) y generación (creación de tokens de salida) se separan en diferentes máquinas para optimizar el throughput y la utilización de GPU.

Para mejorar la eficiencia en cargas de trabajo conversacionales y agenticas, Workers AI implementa 'prefix caching'. Esta técnica almacena los tensores de entrada de solicitudes previas, permitiendo que en turnos subsiguientes de una conversación, solo los nuevos tokens de entrada requieran prefill, reduciendo significativamente el tiempo de procesamiento y mejorando el Time to First Token (TTFT) y Tokens Per Second (TPS). La afinidad de sesión se gestiona mediante un encabezado x-session-affinity en las solicitudes HTTP, asegurando que las solicitudes de una misma sesión se dirijan a la misma instancia de modelo para maximizar la tasa de aciertos del caché.

Además, se ha rediseñado la API asíncrona para manejar grandes volúmenes de inferencia que no son en tiempo real. Este sistema ahora es 'pull-based', donde las instancias de modelo extraen solicitudes de una cola a medida que tienen capacidad disponible, en contraste con un sistema 'push-based' anterior. Esto permite una ejecución más duradera y resistente a picos de carga, priorizando las solicitudes síncronas críticas mientras se procesan las asíncronas de manera eficiente, monitoreando la utilización de GPU en tiempo real.

curl -X POST \
"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/moonshotai/kimi-k2.5" \
-H "Authorization: Bearer {API_TOKEN}" \
-H "Content-Type: application/json" \
-H "x-session-affinity: ses_12345678" \
-d '{
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "What is prefix caching and why does it matter?"
}
],
"max_tokens": 2400,
"stream": true
}'

// (1.) Push a request in queue
// pass queueRequest: true
let res = await env.AI.run("@cf/moonshotai/kimi-k2.5", {
"requests": [{
"messages": [{
"role": "user",
"content": "Tell me a joke"
}]
}, {
"messages": [{
"role": "user",
"content": "Explain the Pythagoras theorem"
}]
}, ...{<add more requests in a batch>} ];
}, {
queueRequest: true,
});
// (2.) grab the request id
let request_id;
if(res && res.request_id){
request_id = res.request_id;
}
// (3.) poll the status
let res = await env.AI.run("@cf/moonshotai/kimi-k2.5", {
request_id: request_id
});
if(res && res.status === "queued" || res.status === "running") {
// retry by polling again
...
}
else
return Response.json(res); // This will contain the final completed response

Fundamentos Teóricos

La optimización de la inferencia de modelos de lenguaje grandes se conecta con principios fundamentales de la computación distribuida y la arquitectura de sistemas de alto rendimiento. Las técnicas de paralelización (data, tensor, expert parallelism) tienen sus raíces en la computación paralela y distribuida, buscando dividir una carga de trabajo computacionalmente intensiva entre múltiples unidades de procesamiento, un concepto explorado en trabajos como 'MapReduce: Simplified Data Processing on Large Clusters' (Dean & Ghemawat, 2004) para el procesamiento de datos a gran escala, y adaptado aquí para la inferencia de modelos.

El 'prefix caching' es una aplicación específica del principio general de caching, fundamental en la informática para reducir la latencia y el consumo de recursos al almacenar resultados de cómputos frecuentes. En el contexto de LLMs, esto se asemeja a la optimización de caché de instrucciones o datos en arquitecturas de CPU, donde la localidad temporal y espacial se explota para mejorar el rendimiento. La gestión de la afinidad de sesión para el caching es un patrón común en balanceadores de carga y proxies, asegurando que las solicitudes relacionadas se dirijan al mismo backend para mantener el estado o aprovechar cachés locales, un concepto bien establecido en la ingeniería de sistemas distribuidos.