Resumen Técnico: Cuando la Cuantización es Gratuita

Enunciado del Problema

La inferencia de transformadores de contexto largo está típicamente limitada por la huella de memoria y los requisitos de ancho de banda de la memoria caché de claves y valores (KV). A medida que aumentan las longitudes de contexto (por ejemplo, 128K tokens), los tensores KV acumulados en fp16 pueden exceder el tamaño de los propios pesos del modelo, creando un cuello de botella donde el ancho de banda de la memoria, y no la capacidad de cálculo, limita la velocidad de decodificación.

La práctica estándar de la industria aborda esto mediante estrategias de "rotar-luego-cuantizar" (por ejemplo, QuaRot, SpinQuant), que aplican rotaciones ortogonales aleatorias para gaussianizar las activaciones antes de cuantizarlas a pocos bits (int4/int8). Sin embargo, este enfoque se formula tradicionalmente como una compensación entre calidad y latencia: la sobrecarga computacional de la rotación y la desrotación, combinada con los errores de cuantización, típicamente resulta en un aumento del 5–20% en la latencia de decodificación, incluso si se ahorra memoria.

Los autores postulan que esta compensación se invierte en Silicio de Apple (arquitectura de memoria unificada). Hipotetizan que los ahorros de ancho de banda de una caché comprimida 3× pueden superar el costo computacional del kernel de cuantización, siempre que el kernel esté altamente optimizado y fusionado.

Metodología

1. Transformada de Fourier Aleatorizada por Signos (SRFT)

El artículo adopta SRFT (Transformada Rápida de Fourier Aleatorizada por Signos) como mecanismo de rotación, definida como $SRFT(x) = \text{pack}(F \cdot \text{diag}(s) \cdot x)$, donde $F$ es la DFT unitaria y $s$ es un vector de signos aleatorios fijos.

• Racional: Mientras que las transformadas de Hadamard (SRHT) son comunes, SRFT ofrece flexibilidad para dimensiones que no son potencias de dos y se alinea con operaciones de multiplicación de matrices de 16×16 (AMX, Tensor Cores).
• Hallazgo: Los autores demuestran que SRFT y SRHT son estadísticamente indistinguibles en términos de calidad de perplejidad (PPL) en cinco modelos (SmolLM2, Qwen2.5, Gemma-3) y diversos presupuestos de bits.

2. Kernel Metal Fusionado

Para eliminar la sobrecarga de despacho, los autores implementan un único kernel de cálculo Metal fusionado que realiza toda la tubería:

1. Carga y Multiplicación por Signos: Aplicar el diagonal de signos aleatorios en registro.
2. FFT: Ejecutar $\log_2(d/2)$ pases de mariposa radix-2 Stockham.
3. Empaquetado: Desempaquetar la simetría hermitiana a una representación real.
4. Escalar y Cuantizar: Realizar reescalado $\lambda$ por canal (aprendido o calibrado), reducción de máximo absoluto por grupo y redondeo a int4/int8 con signo.
5. Empaquetado de Nibbles: Empaquetar dos valores int4 en un solo byte usando mezclas SIMD.

Este kernel se expone mediante torch.mps.compile e integra como una subclase Cache de HuggingFace (SRFTInt4Cache).

3. Estrategias Avanzadas de Escalamiento

• Por Canal + Por Grupo: Para modelos con dimensiones de cabeza pequeñas (por ejemplo, Qwen2.5-1.5B, $d=128$), el escalamiento por token falla catastróficamente debido a coordenadas atípicas. Los autores combinan escalamiento por canal (una escala por coordenada) con escalamiento por grupo (grupos de 16) para recuperar la calidad.
• Rotaciones Aprendidas: El artículo investiga la calibración post-entrenamiento donde se añaden una rotación aprendible $R$ y una escala $\lambda$. Un hallazgo clave es que aprender $R+\lambda$ sin la base SRFT minimiza significativamente más el MSE de calibración (84.9% vs 50.3%) pero produce un peor PPL aguas abajo, lo que sugiere que la rotación aleatoria fija SRFT actúa como un regularizador.
• Rotaciones de Householder: En $d=256$, las rotaciones de Householder con $k=d/2$ reflectores logran una cuantización efectivamente sin pérdidas ($\Delta PPL \approx 0.000$) con la mitad de parámetros que las rotaciones de Cayley completas.

4. Ruta de CPU (AMX)

Una ruta de CPU complementaria utiliza el coprocesor de matrices AMX de Apple mediante cblas sgemm. Al agrupar DFTs de longitud 8 como multiplicaciones de matrices reales de 16×16, los autores alcanzan hasta el 91% del pico de AMX de un solo núcleo, ofreciendo una alternativa de baja latencia para regímenes de lotes pequeños donde la sobrecarga de despacho de GPU es dominante.

Resultados Clave

Rendimiento (Latencia)

En Apple M1 (16 GB de memoria unificada), la caché KV int4 fusionada supera a fp16 en la latencia de decodificación de extremo a extremo:

• Gemma-3 1B: Se ejecuta 3–8% más rápido (sobrecarga de latencia negativa) en prefijos de 256–4096 tokens.
• Qwen2.5-1.5B: Se ejecuta 0.7–2.6% más rápido en contextos cortos (≤1K tokens).
• Mecanismo: La sobrecarga del kernel (~25 ns/vec) es menor que los ahorros de ancho de banda de memoria logrados al transferir 3× menos datos por paso de decodificación.

Calidad (Perplejidad)

• SRFT vs. SRHT: No hay diferencia estadísticamente significativa en $\Delta PPL$ en todos los modelos probados.
• Qwen2.5-1.5B: El kernel fusionado con escalamiento por canal + por grupo reduce el $\Delta PPL$ de 4 bits de un +7975 catastrófico (línea base por token) a +638.6 (una reducción de 12.5×), igualando la implementación de referencia en Python.
• Gemma-3 1B: El $\Delta PPL$ de 4 bits mejoró de +18.7 (por token) a +3.6 con el kernel escalado fusionado.
• Régimen sin Pérdidas: En $d=256$, las rotaciones de Householder aprendidas logran $\Delta PPL = 0.000$.

Eficiencia de Memoria

• Compresión: Logra una compresión de memoria persistente de 3.0× a 3.56× en capas de atención completa.
• Impacto en el Mundo Real: Para un contexto de 128K en Llama-3-70B, la caché int4 reduce la memoria de ~41 GB (fp16) a ~10.9 GB.

Significado y Afirmaciones

El artículo afirma que la cuantización es "gratuita" (positiva en rendimiento) en hardware de memoria unificada, invirtiendo la suposición estándar de la industria de que la cuantización de pocos bits siempre incurre en una penalización de latencia.

1. Inversión de la Compensación: La contribución principal es demostrar que, en Silicio de Apple, los ahorros de ancho de banda de una caché comprimida 3× superan el costo computacional del kernel de cuantización. Esto hace que la caché KV int4 sea estrictamente superior a fp16 para inferencia de contexto largo sensible a la latencia en portátiles y dispositivos de borde.
2. Diseño Consciente del Hardware: El trabajo destaca que la elección de la rotación (FFT vs. Hadamard) es menos sobre superioridad matemática y más sobre alineación con el hardware (por ejemplo, la estructura radix-8 de FFT que coincide con los mosaicos 16×16 de AMX).
3. Separación entre Calibración y PPL: El artículo aísla un fenómeno novedoso donde minimizar el MSE de calibración no garantiza un mejor PPL aguas abajo, sugiriendo que la rotación aleatoria fija proporciona un prior de regularización necesario.
4. Viabilidad de Despliegue: Al fusionar toda la tubería en un solo kernel Metal e integrarlo con HuggingFace, los autores demuestran que la cuantización agresiva de 4 bits puede desplegarse sin sacrificar la calidad del modelo ni la velocidad, siempre que la arquitectura del hardware (memoria unificada) y la optimización del kernel (fusión) estén alineadas.

Los autores concluyen que, si bien la compensación estándar de cuantización se mantiene para GPUs con HBM separada y altas sobrecargas de despacho, se invierte para el régimen de memoria unificada donde opera la inferencia en portátiles y dispositivos de borde.