Investigación de ingeniería: ¿puede este port de MLX ser otro 10× más rápido?

Fecha: 2026-09-19. Máquina: Apple M3 Max, 40 núcleos de GPU, 128 GiB de memoria unificada, macOS 27.2, MLX / MLX Metal 0.32.2, inferencia FP16. Este informe contiene experimentos locales reales, incluido un kernel Metal escrito a mano. No cambia el runtime de producción ni publica pesos cuantizados.

Los cambios de ingeniería probados no ofrecen 10×. Las mediciones intercaladas respaldan mejoras modestas y dependientes de la forma a partir de la compilación y de podar salidas no usadas de la última capa de cabezas de decisión. Casos seleccionados mejoraron aproximadamente un 3–8% usando medianas emparejadas por ronda. Algunos intervalos de lotes más grandes no incluyen mejora. Un kernel GELU/gate de erf exacto personalizado tuvo éxito numérico pero no aportó un beneficio adicional consistente de extremo a extremo sobre la compilación de MLX. La cuantización ingenua de la columna a 8 y 4 bits redujo el almacenamiento, no logró acelerar las cargas de trabajo piloto más grandes y cambió las predicciones o las probabilidades calibradas.

Los límites matemáticos y los trade-offs de aproximación se examinan por separado en MATH_10X_RESEARCH.md. La revisión original de la implementación está en PERFORMANCE_RESEARCH.md; el benchmark del checkpoint publicado sigue siendo BENCHMARKS.md.

Controles y límites experimentales

Todo el trabajo de GPU de la investigación se ejecutó en serie. Otro trabajo del agente usó solo CPU/sistema de archivos/red. La máquina estaba conectada a la corriente, sin ninguna advertencia térmica/de rendimiento de pmset registrada y sin uso de swap reportado durante el experimento. La actividad normal del escritorio continuó. No es una cámara térmica controlada ni una máquina de benchmark dedicada por lo demás inactiva.

Las primeras ejecuciones de cribado ejecutaron cada candidato en un proceso nuevo, con 4–5 warmups y 12–16 muestras. Revelaron una deriva considerable entre ejecuciones. Por ejemplo, el piloto inglés de una sola pregunta sugería una mejora de compilación de 1.24×, mientras que el experimento intercalado posterior encontró solo alrededor de 1.03×. Las latencias del piloto secuencial son, por tanto, evidencia de cribado, no la afirmación causal principal de mejora.

El script de confirmación paired.py:

• Rota el orden de los candidatos dentro de cada ronda y usa las mismas entradas para cada candidato de esa ronda.
• Cambia el texto real del estado entre rondas. Genera hasta 16 variantes de estado y retiene las variantes con la misma forma de tensor; los casos breves multilingües tienen 10 de esas variantes, mientras que los demás casos reportados tienen 16.
• Usa preguntas distintas en lenguaje natural, incluidas 50 instrucciones diferentes para la carga de trabajo breve más grande. Comprueba los hashes de entrada y no cachea respuestas, no deduplica preguntas ni reutiliza estados contextuales del codificador.
• Evalúa los resultados y sincroniza la GPU antes de detener cada temporizador. Mide tanto las llamadas de forward preparadas como la ruta pública de predicción, incluida la tokenización y el formateo de salida. La carga del modelo queda excluida.
• Ejecuta 32 rondas medidas para el experimento inglés de cabeza/compilación y 16 para los experimentos multilingüe y de Metal personalizado, tras el warmup. Cada candidato ve el mismo número de rondas y la misma secuencia de entrada.

Estas entradas difieren de los fixtures de línea base publicados. Las comparaciones de abajo son dentro del experimento de investigación, no comparaciones antes/después obtenidas dividiendo tablas no relacionadas. El límite de lote de la investigación es 64, mientras que la API publicada usa 16 por defecto. Las variantes de estado repetidas son mediciones repetidas intencionales; no hay caché de resultados.

analyze.py calcula ratios por ronda eager_time / candidate_time e intervalos bootstrap percentiles exploratorios para su mediana, usando 2,000 remuestreos de índices de ronda. Esos intervalos no tienen en cuenta todas las fuentes de ruido del sistema operativo ni la correlación serial y no sustituyen a una replicación en varias sesiones. Un ratio de valores p50 calculados de forma independiente puede diferir de la mediana del ratio emparejado.

El JSON en bruto incluye todos los tiempos, hashes de entrada, metadatos del entorno, métricas de paridad y la huella del código fuente registrada en el momento de la medición. Los scripts del experimento se formatearon y ampliaron después con candidatos opcionales disjuntos; las huellas anteriores describen esas versiones anteriores de los scripts.

Compilación y poda exacta de la cabeza final

Se compararon cuatro rutas:

1. Eager: el DecisionModel FP16 publicado.
2. Compilado: mx.compile alrededor del modelo cargado, evaluado y congelado, usando especialización de formas normal.
3. Q seleccionada + compilado: en la última capa de cabeza, conserva la proyección QKV de longitud completa y los K/V, pero emite solo las consultas de atención de CLS/marcadores de opción. Ejecuta la proyección de salida y la FFN solo sobre estas salidas seleccionadas.
4. Atención completa + salidas seleccionadas + compilado: conserva la llamada original QKV y SDPA de longitud completa, y luego recoge las salidas de CLS/opción antes de la proyección de salida y la FFN. Esto retiene la forma original del kernel de atención a la vez que elimina la mayor parte del trabajo denso no usado de la cabeza final.

Ambos prototipos de poda preservan las dependencias matemáticas del modelo. Siguen calculando todas las proyecciones QKV; no realizan el ahorro adicional de la proyección solo de Q del cálculo matemático de cota superior. Cambiar las formas de GEMM y SDPA puede cambiar el redondeo en coma flotante. Ninguno de los prototipos es una caché de decodificador, una salida temprana ni una aproximación que descarte capas de Transformer anteriores.

Latencia p50 de extremo a extremo, en milisegundos:

Fuentes: datos emparejados en inglés y datos emparejados multilingües.

Para la ruta de atención completa/salidas seleccionadas, la mejora mediana emparejada y los intervalos exploratorios del 95% incluyen:

Los intervalos inglés de 50 preguntas y multilingüe de lote largo incluyen el 1. No establecen una mejora repetible. La ruta de Q seleccionada es algo mejor para los casos multilingües breve-16 y largo-1, pero ninguna ruta de poda domina todas las formas. Todos los intervalos de candidatos, las mediciones de forward y los ratios brutos por ronda están en paired_analysis.json.

La compilación coincidió exactamente con los logits eager, los logits de acción y las probabilidades calibradas en las 1,530 comparaciones de preguntas con entrada cambiada a lo largo de las dos familias de modelos en este experimento de cabeza/compilación. Ambas rutas de poda coincidieron en las 1,530 decisiones de argmax, con una diferencia máxima de probabilidad calibrada de 0.0001883. La ruta de poda de atención completa también superó la suite separada de 63 preguntas de cada modelo: 126/126 de coincidencia, con diferencias máximas de probabilidad de 4.31e-5 para inglés y 6.48e-6 para multilingüe. Son comprobaciones de regresión, no una afirmación de exactitud de tarea en 1,530 ejemplos etiquetados de forma independiente.

Se cribaron tanto la compilación de todo el modelo como la compilación por bloque. El experimento por bloque también preservó las 63 salidas del fixture inglés, pero no estableció una ventaja material sobre la compilación de todo el modelo. La especialización de formas debe estar acotada en un servicio. El modelo usa reshapes y máscaras dependientes de la forma en Python, así que aplicar shapeless=True indiscriminadamente es inseguro. La guía oficial de compilación documenta la especialización de formas y la captura de estado.

La primera llamada del candidato inglés de todo el modelo tardó 2,166.7 ms, seguida de unos 12.75 ms de p50 de forward en caliente en ese piloto; una primera llamada con una forma B16 nueva tardó 272.4 ms. El campo JSON se llama cold_forward, pero significa la primera llamada del candidato tras la inferencia de referencia eager, no una aplicación totalmente en frío ni un driver Metal recién inicializado. Los candidatos posteriores reutilizaron kernels Metal compilados previamente, así que sus tiempos de primera llamada no son una clasificación controlada del coste de arranque en frío. La memoria MLX activa/pico del short-1 inglés compilado era de aproximadamente 803.6/918.6 MiB en el piloto; el multilingüe era de aproximadamente 614.1/676.9 MiB. Estas mediciones del asignador no incluyen todas las asignaciones del compilador del lado del host y no establecen límites de memoria bajo cambios de forma ilimitados. Consulta piloto de compilación en inglés y piloto de compilación multilingüe.

Cuantización selectiva: ahorros de almacenamiento útiles, no adecuada como afirmación de velocidad

El prototipo llama a nn.quantize después de cargar el modelo denso en FP16. Selecciona solo módulos lineales encoder.layers.*, con tamaño de grupo afín 64, y luego compila el modelo resultante. Los embeddings, las normas, la cabeza de decisión, el scorer y la cabeza de acción permanecen en FP16. Esto evita castear pesos enteros empaquetados a través del cargador denso actual y evita el ancho de entrada no divisible 1028/772 de la cabeza de acción. No se publica ningún formato de checkpoint cuantizado ni contrato de carga. La implementación oficial de capas cuantizadas de MLX proporciona este mecanismo de selección.

El resultado multilingüe de 4 bits ilustra por qué la pequeña suite de fixtures por sí sola es insuficiente: sus 63 argmax del fixture siguieron igual, pero 7 de 26 decisiones de carga de trabajo distintas cambiaron. Son mediciones de coincidencia frente a FP16, no mediciones de exactitud de referencia. Un cambio absoluto de probabilidad de 0.8008 es 80.08 puntos porcentuales.

Con entradas piloto inglesas short-16, el p50 de extremo a extremo eager/compilado en FP16 fue 91.26/87.94 ms; el compilado a 8/4 bits fue 96.66/93.20 ms. La cuantización short-1 parecía algo más rápida en esa ejecución de cribado, mientras que las formas más grandes no. El cribado multilingüe de formas grandes tampoco mostró una ganancia de velocidad, pero sus ejecuciones secuenciales tuvieron una deriva sustancial. Estas observaciones justifican rechazar una afirmación de mejora o de publicación sin matices, no asignar factores de ralentización precisos sin replicación cuantizada intercalada. El trabajo de cuantización adicional necesita calibración o ajuste fino conscientes de las activaciones y una suite de calidad etiquetada representativa.

Fuentes en bruto: inglés 8-bit, inglés 4-bit, multilingüe 8-bit, multilingüe 4-bit.

Metal escrito a mano: se implementó y probó una fusión exacta GELU/gate

kernels.py implementa un kernel Metal personalizado real que lee las dos ramas de la MLP concatenadas, calcula el mismo GELU basado en erf, lo multiplica por el gate y escribe una única salida. No sustituye el tanh-GELU ni una aproximación con sigmoide. El kernel usa los propios helpers erf y expm1 de MLX v0.32.2, preservando sus licencias y avisos en vendor/README.md. Soporta explícitamente solo FP16 y usa el modo de matemáticas seguro de Metal. La guía oficial de kernels personalizados describe esta API y sus controles de modo de matemáticas.

En ocho formas de activación representativas, 27,958,016 elementos de salida FP16 generados aleatoriamente tuvieron valores exactamente iguales a la operación original. El microbenchmark compara igualdad numérica, no el bit de signo del cero. Las pruebas de modelo completo con entrada cambiada también coincidieron exactamente: 474/474 comparaciones de preguntas a lo largo de las dos familias de modelos, más ambas suites de fixtures de 63 preguntas, con cero diferencia en logits, logits de acción o probabilidad calibrada.

Este resultado de corrección no se tradujo en una ventaja de velocidad consistente sobre la expresión compilada fusionada de MLX. Por ejemplo, con 1,312 tokens y ancho intermedio 2,624, el tiempo de activación sincronizado por llamada fue 0.378 ms para GELU-y-gate eager, 0.268 ms para mx.compile y 0.280 ms para el kernel personalizado. Con 8,192 tokens y ancho 1,152, los valores correspondientes fueron 0.846/0.764/0.714 ms. Estos microbenchmarks incluyen sobrecarga de despacho y sincronización y son sondas de cribado; no son mediciones del tiempo de ejecución aislado del dispositivo. Las entradas completas, los tiempos en bruto y las comprobaciones de igualdad están en microbench.json.

El kernel personalizado se instaló después en cada MLP del codificador y se midió en el modelo completo con orden de candidatos rotatorio y entradas cambiantes:

Las ejecuciones emparejadas completas del kernel personalizado usan una segunda instancia del modelo con pesos idénticos, de modo que las implementaciones sin modificar y personalizada coexisten sin mutación ni capturas compiladas obsoletas. Sus tiempos absolutos no deben compararse con la ejecución anterior de poda de la cabeza. Los resultados mixtos y modestos no respaldan publicar el kernel personalizado como una mejora general de rendimiento. Fuentes: datos emparejados de Metal en inglés y datos emparejados de Metal multilingües.

Dónde merecería más investigación la ingeniería personalizada

El modelo ya llama a mx.fast.scaled_dot_product_attention, mx.fast.rope y a la normalización de capas optimizada. Su ruta SDPA con máscara booleana D64 está fusionada; no hay ningún interruptor de Flash Attention que falte que explique una brecha de 10×. La atención local sigue recorriendo tiles densos de clave/valor. Un kernel real de ventana bidireccional podría omitir esos tiles preservando la distancia inclusiva <=64 y la semántica del padding, pero su oportunidad aritmética sobre todo el modelo es pequeña en entradas cortas y está acotada en las formas largas publicadas. La revisión de código existente y el informe matemático cuantifican esta distinción.

Los siguientes proyectos útiles, con sus requisitos de evidencia, son:

• Atención de ventana para entradas largas: especializa los límites de tiles para D64, la ventana bidireccional real y los lotes con padding. Compara con la SDPA densa fusionada a 512/1024 tokens y luego en el modelo completo. Este kernel no se ha construido ni medido en este informe.
• Epílogos y planificación de kernels densos: investiga fusionar el epílogo de la MLP con gating en el GEMM o mejorar la planificación de matrices con M pequeño. MLX ya usa implementaciones Metal de GEMM especializadas, así que sustituirlas requiere un perfil real de despacho/kernel y ganancias medidas para las formas exactas de M/N/K. El resultado aislado de activación muestra por qué otro kernel elemental por sí solo es insuficiente.
• Batching consciente de la longitud y preparación compartida en CPU: preserva IDs de entrada exactos mientras tokenizas una vez el texto de estado compartido antes de construir cada secuencia de pregunta, y evita rellenar elementos pequeños con elementos largos no relacionados. El piloto multilingüe long-8 pasó unos 13.1 ms preparando entradas, frente a cientos de milisegundos de extremo a extremo. Incluso eliminar por completo esa preparación no produciría 10× en esta carga de trabajo. La latencia de cola y el recuento de inferencias únicas deben formar parte de cualquier afirmación sobre batching.
• Un estudiante más pequeño que responda conjuntamente: si 10× es un requisito de producto, destila o rediseña el modelo para eliminar la mayor parte del trabajo denso o responder muchas preguntas fijas con una única codificación contextual. Esto cambia el modelo aprendido y necesita entrenamiento/evaluación etiquetados y representativos; no es una optimización exacta del port. Reutilizar un estado/KV contextual arbitrario entre preguntas en el codificador bidireccional actual no es válido.

Ocho sondas GEMM independientes de proyección de entrada del codificador en FP16 lograron 0.66–11.55 TFLOP/s, incluida la sincronización por llamada. La sonda inglesa grande M=4096, N=5248, K=1024 logró 11.55 TFLOP/s; la sonda multilingüe M=8192, N=2304, K=768 logró 7.96 TFLOP/s. Estos son valores de rendimiento observados, no especificaciones de pico de hardware ni cotas superiores del rendimiento del grafo completo. Las mediciones con M pequeño están especialmente dominadas por los costes de envío y sincronización; un grafo en streaming los amortiza de forma distinta. Muestran qué formas merecen perfilado, no una prueba de que no pueda existir un kernel mejor. Los presupuestos de rendimiento de 10× con el mismo trabajo del informe matemático siguen siendo requisitos teóricos, no capacidades medidas del dispositivo.

Reproducción y decisión de publicación

Los scripts usan el .venv existente y los checkpoints locales fijados. Ejecuta los comandos de GPU de forma secuencial, nunca junto al benchmark formal:

# Screening: repeat for eager, compiled, blocks, q8, q4, selected-compiled.
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant compiled --iterations 12 --warmup 4 --quality \
  --output experiments/engineering/reproduced-compiled.json

# Primary confirmation, including 50 genuinely different questions.
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 32 \
  --output experiments/engineering/reproduced-laya-paired.json
.venv/bin/python -m experiments.engineering.paired \
  --model laya-multilingual --iterations 16 --cases short1,short16,long1,long8 \
  --output experiments/engineering/reproduced-multilingual-paired.json

# Hand-written kernel microbench and complete-model comparison.
.venv/bin/python -m experiments.engineering.microbench
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 16 --cases short1,short16,long1 --metal \
  --output experiments/engineering/reproduced-metal-paired.json
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant metal-compiled --iterations 5 --warmup 3 \
  --cases short1 --quality --output experiments/engineering/reproduced-metal-quality.json

# CPU-only paired analysis.
.venv/bin/python -m experiments.engineering.analyze

Todos los archivos Python experimentales pasan los controles de formato y lint de Ruff. El runtime estable, los resultados de benchmark originales y los checkpoints FP16 publicados siguen siendo los artefactos de publicación. La compilación y la poda exacta de la cabeza final son optimizaciones futuras opcionales creíbles tras la política de formas en frío/caché y una validación de calidad más amplia; las ganancias medidas no justifican añadir en silencio latencia de compilación o un kernel personalizado a la ruta por defecto. No se afirma ninguna mejora de 10×, ningún checkpoint cuantizado listo para producción ni ninguna ganancia medida del kernel de ventana local.