KAIST elimina la calibración del decodificado MoE y acelera los bloques hasta 2,69× frente a vLLM

El decodificado de Mixture-of-Experts con batches pequeños no está limitado por el cómputo, sino por la memoria: para cada token hay que cargar los pesos de los expertos enrutados. Los investigadores de KAIST AI Geunhyeon Kim y Sung Ju Hwang han demostrado que una técnica de aceleración muy extendida, la selección de canales importantes dentro de un experto, puede funcionar sin datos de calibración. Su lector, construido únicamente a partir de los pesos del modelo, conserva el 98,1% de la precisión del modelo completo al leer el 17,8% de los pesos de los expertos. Los bloques MoE con kernels de Triton funcionan entre 1,60 y 1,90 veces más rápido que vLLM optimizado con batch one.

El umbral promete un presupuesto, pero no lo cumple

La idea de la selección de canales es sencilla. En cada experto, la mayoría de los canales intermedios apenas influye en un token concreto. Un selector predice qué canales son importantes y lee de la memoria solo sus pesos. Lo habitual es quedarse con los canales cuya puntuación supera un umbral, que se ajusta con una muestra offline.

Los autores señalan dos debilidades de este esquema. Primero, el umbral no fija el volumen de lectura: el número de canales que superan el corte varía de una entrada a otra. Prox, con umbrales calibrados para un presupuesto del 18%, leyó en realidad un 12,9% de media en sus benchmarks. El ingeniero que planifica latencia y rendimiento obtiene una cifra distinta de la que esperaba.

Segundo, tampoco está garantizada la ganancia de precisión para un presupuesto dado.

Top-k en lugar de umbral

La alternativa no necesita datos en absoluto: para cada token se conservan los k canales con mayor puntuación. El presupuesto de lectura se conoce de antemano y coincide con el declarado. Al comparar sobre las mismas puntuaciones, la versión top-k del selector igualó o superó a la de umbral en 40 de 43 casos con presupuestos del 13% y el 18%. En tareas de generación ganó en las 23 comparaciones.

Hay también un efecto cualitativo. Los umbrales calibrados llevaron a Gemma 4 a colapsar en bucles de repetición, mientras que las versiones top-k lo evitan. La conclusión de los autores: la calibración del umbral no aporta ninguna de las dos cosas por las que se hace, ni un presupuesto de lectura conocido ni una mejora estable de precisión con ese presupuesto.

Un lector basado solo en pesos

La segunda parte del trabajo es un lector propio que no necesita ni datos de calibración ni ajuste fino. Consta de dos pasos:

• los canales se descartan de forma aproximada mediante un boceto ternario (ternary sketch) de los pesos;
• los canales omitidos se compensan con una proyección de bajo rango, que también se construye a partir de los pesos del modelo.

Cifras clave en tres modelos MoE posentrenados:

• precisión: 98,1% respecto al modelo completo;
• volumen de lectura de pesos de los expertos enrutados: 17,8%;
• velocidad de los bloques MoE con kernels de Triton y batch one: entre 1,60 y 1,90 veces superior a vLLM optimizado;
• con batches grandes, la aceleración llega a 2,69 veces en los mismos tres modelos;
• Prox, calibrado para el 18%, leyó en la práctica un 12,9%;
• top-k frente a umbral: 40 de 43 comparaciones con presupuestos del 13% y el 18%, 23 de 23 en generación.

A quién le sirve

Los principales beneficiados son quienes ejecutan grandes modelos MoE en hardware con poca capacidad de memoria. En una conversación en Hugging Face, uno de los autores calificó expresamente de acuciante el problema de la poda de activaciones en GPU con memoria lenta, como DGX Spark. En ese hardware, el ancho de banda de memoria, y no los FLOPS, determina la velocidad de generación de tokens. Según sus palabras, el siguiente régimen de dispersión (sparsity) vendrá precisamente de la dispersión de activaciones en el despliegue sobre dispositivos con un límite estricto de ancho de banda.

La ventaja práctica de prescindir de la calibración es que el método se puede aplicar a cualquier modelo open-weights reciente nada más publicarse. No hay que buscar una muestra representativa ni se corre el riesgo de que un umbral ajustado en un dominio rompa la generación en otro, como ocurrió con Gemma 4. Un presupuesto de lectura fijo simplifica los cálculos: el 17,8% planificado se traduce en un ahorro predecible de tráfico de memoria y de latencia.

El código está publicado en el repositorio calibration-free-moe de GitHub. Para los equipos que escriben sus propios kernels en Triton u optimizan la inferencia local, es un punto de partida listo para usar. Conviene probarlo con los modelos propios, sobre todo en escenarios interactivos con batch one, donde la memoria es el mayor cuello de botella.