Investigación de rendimiento de Laya MLX
Fecha de la investigación: 2026-09-19. Objetivo: Apple M3 Max, 40 núcleos de GPU, 128 GiB de memoria unificada, MLX/MLX Metal 0.32.2. Esta es una revisión estática del runtime nativo, la implementación de MLX instalada, la documentación oficial y el JSON de benchmarks existente. Para esta investigación no se ejecutó ningún benchmark de GPU ni inferencia de modelo. Ninguna de las optimizaciones propuestas más abajo tiene una mejora medida en este informe.
Los primeros experimentos deberían ser la compilación de todo el modelo y una planificación representativa de lotes, seguidos de la multiplicación de matrices cuantizada selectiva. Estos abordan el trabajo repetido dominante. Un kernel especializado de atención local es un proyecto creíble a más largo plazo para entradas largas. La poda exacta de la última capa de cabezas de decisión es factible, pero su ahorro aritmético sobre todo el modelo es solo de unos pocos por ciento. Las mejoras grandes sin cambiar el checkpoint requerirán mejorar la columna densa, eliminar solicitudes genuinamente redundantes o encontrar un cuello de botella de implementación medido; sustituir una activación o habilitar otra bandera de atención probablemente no bastará.
Qué establecen las mediciones existentes
Las siguientes son latencias medianas de extremo a extremo existentes, que incluyen la preparación del prompt y el formateo de resultados, con finalización sincronizada de la GPU, cinco warmups y 50 iteraciones medidas. La carga y las descargas del modelo quedan excluidas. El benchmark permite un lote de 64 preguntas; el valor por defecto del runtime público es 16, así que el resultado de 50 preguntas no es la configuración por defecto de la API.
| Checkpoint / precisión | Breve 1 pregunta | Breve 10 preguntas | Breve 50 preguntas | Larga 1 pregunta | Larga 10 preguntas |
|---|---|---|---|---|---|
| Laya MLX FP16 | 13.421 ms | 71.068 ms | 336.030 ms | 44.927 ms | 420.987 ms |
| Laya MLX FP32 | 15.954 ms | 98.820 ms | 450.712 ms | 61.331 ms | 534.242 ms |
| Laya stock Torch MPS FP32 | 24.918 ms | 95.265 ms | 497.856 ms | 65.581 ms | 586.594 ms |
| Multilingual MLX FP16 | 7.390 ms | 27.386 ms | 127.565 ms | 37.635 ms | 389.487 ms |
| Multilingual MLX FP32 | 7.988 ms | 32.337 ms | 151.387 ms | 47.208 ms | 451.331 ms |
| Multilingual stock Torch MPS FP32 | 19.349 ms | 43.158 ms | 194.171 ms | 52.939 ms | 534.492 ms |
Fuentes: Laya FP16, Laya FP32, Laya MPS, multilingüe FP16, multilingüe FP32 y multilingüe MPS. Las entradas largas contienen 512 tokens para Laya y 1024 para el multilingüe; comparar sus latencias de entrada larga no es, por tanto, una comparación a igual longitud de secuencia. Las longitudes cortas con padding son 93 y 91, respectivamente. Las comparaciones con Torch deben conservar la etiqueta FP32: combinan un cambio de backend con un cambio de precisión al compararse con MLX FP16.
Existe una variabilidad de ejecución significativa. Por ejemplo, la ejecución multilingüe FP16 de 10 preguntas largas tiene p50 389.487 ms, p95 462.319 ms y máximo 619.663 ms. Su mediana de forward de una sola pregunta breve es 8.023 ms, mientras que la mediana de extremo a extremo medida de forma independiente es 7.390 ms. Restar esas medianas produciría un tiempo de preprocesamiento negativo sin sentido. Los archivos actuales no aíslan el tokenizador, el despacho de Python, los kernels individuales de GPU ni los costes de sincronización. Establecen líneas base útiles, no un diagnóstico de cuello de botella a nivel de kernel.
Los informes de validación registran 63/63 de coincidencia de argmax para cada uno de los tres checkpoints en FP32 y FP16, y 100 llamadas repetidas, finitas y deterministas por variante: 378/378 coincidencias de respuesta y 600 llamadas repetidas en total. Son comprobaciones de regresión sobre un corpus pequeño de fixtures, incluidas preguntas repetidas. No son evidencia de que futuros cambios de cuantización o de arquitectura preserven la exactitud general en las tareas.
Por qué la multiplicación de matrices densa merece prioridad
La implementación del modelo aplica proyecciones QKV y de salida, una MLP de codificador con gating y dos capas convencionales de Transformer de cabezas de decisión a cada token con padding. Sea D el tamaño oculto, I el tamaño intermedio del codificador, N el número de capas del codificador y H el número de capas de cabezas de decisión. El número de pesos de matriz usados por token en estos bloques es:
A = N * (4 * D^2 + 3 * D * I) + H * 12 * D^2
dense FLOPs per batch ~= 2 * B * L * A
dense attention FLOPs ~= 4 * B * (N + H) * L^2 * D
Estas estimaciones cuentan la multiplicación y la suma por separado y excluyen la normalización, las activaciones, los embeddings, la puntuación, el enmascaramiento, el movimiento de memoria y la sobrecarga de kernels. Son un modelo aritmético, no un perfil de runtime.
| Familia de checkpoint | D / I / capas del codificador | Pesos del embedding de tokens | Pesos de matriz principales por token A | Capas globales / locales del codificador |
|---|---|---|---|---|
| Laya / decisiones tipadas | 1024 / 2624 / 28 | 51,576,832 | 368,312,320 | 10 / 18 |
| Multilingüe | 768 / 1152 / 22 | 196,608,000 | 124,452,864 | 8 / 14 |
Los aproximadamente 322 millones de parámetros totales del checkpoint multilingüe incluyen 196.6 millones de parámetros de embedding. Para la inferencia solo se recogen filas seleccionadas del embedding; no es una proyección completa del vocabulario. Su carga de trabajo principal de matrices por token es aproximadamente un tercio de la del modelo en inglés, pese a que sus recuentos totales de parámetros parecen mucho más cercanos. Esto es coherente con la brecha de latencia medida en lotes cortos, aunque no demuestra un cuello de botella de hardware concreto. La cuantización solo del embedding reduciría predominantemente el tamaño residente de los pesos, sobre todo en el multilingüe; no tiene por qué mejorar la latencia de inferencia.
Con la ruta de atención densa actual, los productos de atención representan alrededor del 1.5% de los FLOPs modelados con L=93 en Laya, el 7.9% con L=512 en Laya y el 23.3% con L=1024 en el multilingüe. Su fracción del tiempo de reloj de pared puede diferir sustancialmente. Un perfilador debería distinguir los kernels de matrices, los kernels de atención, los kernels elementales, la construcción del grafo en CPU y los huecos de inactividad antes de comprometerse con trabajo Metal personalizado.
Experimentos priorizados
| Prioridad | Experimento | Mejor objetivo | Trade-off principal / condición de aceptación |
|---|---|---|---|
| P0 | Perfilar una forma breve y una larga; compilar el modelo o los bloques del codificador | Latencia de una sola solicitud y despacho de Python | Mantener salidas idénticas dentro de la tolerancia de precisión existente; medir la compilación de primer uso por separado |
| P0 | Ajustar el batching por presupuesto real de tokens y distribución de longitudes | Muchas preguntas distintas y tráfico de longitud mixta | Optimizar el rendimiento sujeto a un presupuesto de latencia p95 y de memoria; tener en cuenta la cola |
| P1 | Cuantizar capas lineales seleccionadas de la columna, empezando por 8 bits y luego 4 bits | Tráfico de pesos y, potencialmente, inferencia densa | Puertas de calidad y calibración; la velocidad real en M3 Max puede empeorar |
| P1 | Cachear la tokenización del estado compartido y plantillas de preguntas estables | Muchas preguntas que comparten un estado o rúbricas repetidas | Identidad exacta de tokens; caché acotada; separar los resultados cacheados de los no cacheados |
| P1 | Calcular solo las salidas requeridas de la capa de cabeza final | Toda carga de trabajo, sobre todo secuencias más largas | Poda exacta de dependencias; ahorro aritmético modesto sobre todo el modelo |
| P2 | Atención de ventana local genuina con límites de tiles | Cargas de trabajo de 512/1024 tokens | Nueva complejidad de kernel; preservar la ventana bidireccional y la semántica del padding |
| P2 | Fusionar residual/norm o GELU/gate solo donde el perfil lo justifique | Sobrecarga de kernels pequeños o tráfico de activaciones | Los kernels rápidos existentes ya cubren buena parte de esto; mantener la semántica exacta de GELU |
| Característica de producto aparte | Deduplicar entradas idénticas de forward | Cargas de trabajo que realmente repiten preguntas | Informar del recuento de inferencias únicas y de los aciertos de caché; no presentarlo como una mejora general de kernel |
Compilar la ruta de inferencia evaluada completa
Agent.forward() actualmente construye arrays, llama a self.model y evalúa el resultado. No hay ningún mx.compile envolvente a nivel de modelo ni de bloque. La compilación con forma fija puede reducir la construcción del grafo en Python y fusionar operaciones compatibles. MLX documenta la especialización de formas y la captura explícita de estado; también advierte de que la compilación sin forma (shapeless) no puede preservar de forma segura operaciones arbitrarias de Python dependientes de la forma. Consulta la guía oficial de compilación.
Empieza con un callable compilado creado después de cargar, castear y evaluar los pesos, usando una especialización de formas normal. Conserva el callable sin compilar existente para la comparación de paridad y la compatibilidad con CPU. Para un modelo de inferencia congelado, los pesos pueden permanecer capturados para esa instancia del modelo; si cambian los pesos o la estructura de módulos, reconstruye el callable o captura explícitamente el estado relevante. No reutilices un closure compilado entre reemplazos de checkpoint.
Prueba un envoltorio de todo el modelo y, si las limitaciones de trazado o el coste de compilación lo hacen poco atractivo, compila los bloques del codificador y la cabeza por separado. El código actual lee x.shape en enteros de Python, hace reshape con valores explícitos de batch/longitud, crea máscaras arange(length) e indexa marcadores usando un rango de filas derivado de la forma. Aplicar shapeless=True a este grafo completo sin rediseñarlo es inseguro. Mover la construcción dinámica de máscaras fuera de un bloque compilado y usar operaciones de flatten/unflatten independientes de la forma puede habilitar una variante shapeless posterior; verifícala con B, L y recuentos de marcadores cambiados.
Los buckets de forma pueden limitar el retrazado, pero el padding tiene un coste de cómputo. Rellenar L=93 a 96 añade aproximadamente 3.2% de trabajo por token; rellenarlo a 128 añade aproximadamente 37.6%. Compara la compilación con forma exacta con múltiplos de longitud pequeños y un conjunto acotado de buckets informados por la carga de trabajo. Incluye (batch size, padded length, marker slots, dtype, device/model instance) en las decisiones de política de caché, y mide la latencia de compilación en frío y la memoria retenida bajo cambios de forma.
El benchmark independiente de forward en worker.py llama directamente a agent.model. Si la compilación solo se añade en Agent.forward, el benchmark de forward actual la omitiría mientras que el benchmark de extremo a extremo la usaría. Ambas rutas deben seleccionar explícitamente la misma implementación candidata para una comparación significativa. Mantén mx.eval y la sincronización de GPU en el procedimiento de medición: medir solo la construcción del grafo no mediría la inferencia.
Agrupar por tokens útiles y luego inspeccionar la planificación de matrices
collate_items rellena a la derecha cada fragmento hasta su secuencia más larga; el runtime agrupa las preguntas en orden de inserción. Para tráfico heterogéneo, ordena o agrupa por longitud preparada, usa un presupuesto de tokens además de un techo de recuento de preguntas, y restaura los IDs de pregunta originales y el orden de salida. Compara lotes de 1, 2, 4, 8, 16, 32 y 64 solo donde sean representativos del servicio. Para solicitudes en línea, incluye el tiempo de espera de un lote; las preguntas/segundo en offline por sí solas pueden ocultar una latencia inaceptable.
La carga de trabajo breve actual de 50 preguntas desperdicia aproximadamente el 8.9% de los tokens de padding en Laya y el 5.8% en el multilingüe. Las filas largas del benchmark no tienen desperdicio de padding. En consecuencia, quitar el padding o solo ordenar tiene un beneficio aritmético limitado en estos fixtures. Una distribución de longitud mixta, incluida una pregunta larga entre muchas breves, es necesaria para revelar el beneficio en producción. Eliminar el padding debe preservar las posiciones RoPE por ejemplo, las posiciones de los marcadores y los límites de atención; concatenar ejemplos en una sola secuencia sin una máscara de aislamiento cambia el modelo.
Para los kernels densos, inspecciona las formas y los strides reales. QKV ya es una única proyección, y las dos ramas de entrada de la MLP del codificador ya comparten una proyección. Dividirlas indiscriminadamente añadiría lanzamientos. Compara el aplanado explícito de la entrada contigua [B,L,D] a [B*L,D] solo si el perfilador o la traza de despacho de MLX muestran GEMMs por lotes indeseados; el framework puede que ya aplana de forma eficiente. La inspección del código fuente por sí sola no justifica afirmar que se ha pasado por alto una optimización de GEMM.
No insertes una sincronización después de cada capa en la implementación de producción. El runtime actual evalúa una vez por fragmento. Esperas adicionales podrían eliminar el solapamiento CPU/GPU y oscurecer una mejora de planificación; el perfilado a nivel de capa debería ser una ejecución de diagnóstico aparte.
Cuantización: apuntar a la columna e implementar su contrato de almacenamiento
La implementación de capas cuantizadas de MLX 0.32.2 instalada proporciona nn.quantize(..., class_predicate=...) y QuantizedLinear, con multiplicación de matrices solo de pesos mediante mx.quantized_matmul. La cuantización afín agrupada admite experimentos de 8 y 4 bits. Empieza con las capas lineales del codificador con tamaño de grupo 64, manteniendo activaciones, normas, embeddings de tipo, scorer y cabeza de acción en FP16. Luego añade por separado las capas lineales de las cabezas de decisión y, opcionalmente, la cuantización del embedding. Mide cada variante; los kernels solo de pesos pueden perder frente a GEMMs en FP16 con recuentos de tokens altos.
Hay peligros de integración concretos en el cargador y el modelo actuales:
Agent.__init__castea todos los pesos almacenados a un dtype de coma flotante e instancia solo módulos densos antes de la carga estricta. Un checkpoint cuantizado requiere metadatos que describan los módulos seleccionados, el tamaño de grupo, el ancho de bits y el modo; instancia los módulos cuantizados correspondientes antes de cargar y preserva los pesos enteros empaquetados. Un casteo a coma flotante de pesos empaquetados no es una carga válida.- La primera capa lineal de la cabeza de acción tiene un ancho de entrada
D+4, es decir 1028 o 772, que no es divisible por un tamaño de grupo afín de 32, 64 o 128. Por tanto, una llamada de cuantización general no es adecuada. Comprueba el ancho de entrada de cada capa seleccionada antes de la conversión. DecisionModel.__call__elige el dtype de entrada de la cabeza de acción a partir deself.act_head.layers[0].weight.dtype. En una capa cuantizada, ese peso sería almacenamiento entero empaquetado, no el dtype de activación deseado. Excluir la cabeza de acción evita esta ruta de momento; soportarla más adelante requiere un contrato explícito de dtype de activación.- La cuantización cambia los logits y las probabilidades calibradas. La coincidencia existente en el pequeño fixture FP16 es evidencia insuficiente para la calidad de 4 bits. Usa tareas etiquetadas y reservadas de choice, score y noul, entradas multilingües, decisiones ajustadas, distintos recuentos de opciones y ejemplos de escalado. Sigue la coincidencia de argmax, la exactitud en la tarea, el error de puntuación, la deriva de probabilidad, la calibración y las probabilidades de acción. Unas salidas de acción saturadas pueden ocultar grandes cambios en los logits de acción.
Para escalas y offsets afines en FP16 con tamaño de grupo 64, el almacenamiento aproximado de la matriz es bits/8 + 4/64 bytes por parámetro: 1.0625 bytes a 8 bits y 0.5625 bytes a 4 bits, frente a 2 bytes en FP16. Son estimaciones de almacenamiento para matrices cuantizadas, que excluyen otros tensores y la sobrecarga de empaquetado; no son estimaciones de mejora. La API oficial de cuantización describe la divisibilidad de grupos y los formatos.
Los formatos más nuevos de pocos bits deberían evaluarse contra el backend real de M3 Max, y no darse por supuesto que usan hardware de chips Apple posteriores. La comprobación de disponibilidad de NAX de MLX requiere una generación de arquitectura más nueva que el dispositivo applegpu_g15s registrado. Consulta la comprobación de dispositivo de MLX 0.32.2.
Reutilizar la preparación en CPU donde las entradas son realmente idénticas
build_sequence serializa, sanea y tokeniza el mismo estado por separado para cada pregunta. También tokeniza por separado cada instrucción y cada opción. El tokenizador Rust ya se usa directamente; sustituir la tokenización de Transformers no es una optimización pendiente.
Serializa y sanea el estado una vez por llamada a prepare, codifícalo una vez y corta sus IDs de token al espacio disponible de cada pregunta. Cachea prefijos inmutables de preguntas preparadas cuando se use la misma rúbrica entre estados, con claves que incluyan la identidad/revisión del tokenizador, el tipo de pregunta, los criterios ordenados, la serialización de instrucciones, el saneado de tokens especiales y los presupuestos de tokens. Las cachés acotadas no deben reutilizar resultados tras cambios de tokenizador o de configuración. La codificación por lotes del tokenizador es otro experimento, siempre que su salida coincida exactamente con la secuencia actual de codificaciones independientes.
No tokenices un prompt recién concatenado como sustituto de concatenar piezas codificadas de forma independiente: los límites de subpalabras pueden cambiar. Verifica byte a byte los IDs de entrada, las máscaras de atención, las posiciones de los marcadores, los qtypes, el comportamiento de truncamiento, los criterios estructurados, los literales de máscara, las entradas vacías y el mapeo de salida.
En el benchmark largo, el estado repite una frase 200 veces antes del truncamiento. Evitar N codificaciones repetidas del estado podría ayudar a la preparación en CPU, pero la diferencia de tiempos actual entre extremo a extremo y forward no mide ese ahorro. Mide por separado prepare, la construcción de collate/arrays, el forward y el postprocesamiento, y luego confirma el resultado de extremo a extremo con un corpus reservado no repetitivo.
Una caché KV de decodificador no se aplica a este codificador. Su primera capa es atención bidireccional global; las representaciones de los tokens de estado dependen de la pregunta, de las opciones y de sus posiciones. Reutilizar estados ocultos del estado o K/V entre preguntas distintas cambia los resultados. La tokenización y los resultados idénticos de una entrada completa sí pueden cachearse; el estado contextual arbitrario del codificador no.
Podar exactamente las salidas finales de las cabezas de decisión
Después del último HeadLayer, solo se consumen el token [CLS] y los tokens de los marcadores de opción. Las capas de cabeza anteriores siguen debiendo producir todos los tokens, porque la capa final lee sus K/V. Solo en la capa final:
- Normaliza todos los tokens de entrada y calcula todos los K/V.
- Recoge Q en
[CLS]y en las posiciones de opción válidas, y ejecuta esas consultas contra la secuencia completa de K/V enmascarada. - Aplica la proyección de salida, el residual, la segunda norma y la red feed-forward solo a esas posiciones seleccionadas.
- Usa la salida seleccionada de
[CLS]para la cabeza de acción y las salidas seleccionadas de los marcadores para la puntuación; preserva el padding de los marcadores y el orden original.
Una primera implementación puede conservar la proyección QKV completa fusionada y recoger Q después. Una variante más agresiva divide sus pesos en una proyección KV de longitud completa y una proyección Q de tokens seleccionados. Eso ahorra más aritmética, pero puede hacer que la planificación de GEMM sea menos eficiente. Los índices de marcador con padding duplicados son inofensivos solo si los resultados enmascarados siguen siendo inobservables. Los casos de 1 opción, muchas opciones y marcadores variables necesitan comprobaciones de paridad explícitas. Como un menor número de consultas puede seleccionar un kernel SDPA distinto, la equivalencia matemática no implica resultados de coma flotante idénticos bit a bit.
Con R = 1 + number of option slots, conservar el QKV completo elimina aproximadamente 18 * B * (L-R) * D^2 FLOPs densos y 4 * B * L * (L-R) * D FLOPs de atención de la última capa de cabeza. Dividir Q/KV cambia el coeficiente denso de 18 a 20. En relación con la estimación aritmética de todo el modelo de arriba, usando R=5 se obtiene:
| Checkpoint / longitud | Conservar QKV completo fusionado | Calcular además solo la Q seleccionada |
|---|---|---|
| Laya, L=93 | 2.44% | 2.70% |
| Laya, L=512 | 2.60% | 2.86% |
| Decisiones tipadas, L=1024 | 2.66% | 2.90% |
| Multilingüe, L=93 | 4.03% | 4.47% |
| Multilingüe, L=1024 | 4.22% | 4.58% |
Estas son reducciones estáticas de FLOPs, no reducciones de latencia predichas. La técnica elimina la mayor parte del trabajo de una capa de cabeza, no la mayor parte del trabajo del modelo. Es útil porque preserva las dependencias y se puede implementar sin reentrenar, no porque prometa un múltiplo de la velocidad de todo el modelo.
Construir atención genuinamente local solo después de medir su contribución
Todas las llamadas de atención del codificador ya usan mx.fast.scaled_dot_product_attention; RoPE ya es mx.fast.rope; nn.LayerNorm llama a la primitiva rápida de normalización. La API de atención de MLX acepta máscaras booleanas y realiza el softmax en FP32. La dimensión actual de cabeza es 64. El despacho Metal de MLX 0.32.2 soporta esta forma con máscaras de array y no selecciona el fallback no fusionado para ella durante la inferencia. No hay evidencia de que la máscara booleana de Laya desactive la atención fusionada. force_fused=True, disponible en la versión instalada, es útil como aserción de diagnóstico, pero no debería anunciarse aquí como una nueva ruta rápida.
La limitación restante es la dispersión estructurada. La implementación construye una máscara booleana local densa de forma [B,1,L,L]. En el kernel de atención Metal convencional, el bucle no causal recorre todo el rango de tiles KV; la máscara de array se aplica a las puntuaciones después de la multiplicación QK. Preserva la semántica de atención local sin aprovechar un rango de tiles local.
Un kernel especializado exacto puede limitar cada tile de consultas a la ventana K/V solapada, mantener la acumulación de softmax en FP32 y evitar una máscara densa L por L. La ventana correcta es bidireccional e inclusiva: abs(query_position - key_position) <= 64. Las consultas interiores pueden ver 129 posiciones, pese al nombre de configuración local_attention=128. Las capas de atención completa y ambas capas de cabezas de decisión deben seguir siendo globales. Las claves con padding deben seguir excluidas, y las consultas con padding no usadas necesitan un comportamiento finito definido.
Un prototipo de menor esfuerzo puede agrupar bloques de consultas con porciones K/V solapadas y llamar al SDPA existente con una máscara exacta más pequeña. Usa Q/K con RoPE ya posicionado, o conserva explícitamente los offsets absolutos. Prefiere bloques por lotes antes que muchas llamadas en Python, y ten en cuenta la materialización duplicada de K/V. Este prototipo puede perder frente al kernel actual en longitudes cortas; es un experimento de corrección y de punto de equilibrio antes de mantener Metal personalizado.
La reducción máxima de FLOPs modelados de todo el modelo al eliminar todos los pares de atención local prohibidos es pequeña en entradas cortas y más prometedora en las largas:
| Checkpoint / longitud | Reducción total ideal de FLOPs por dispersión local exacta |
|---|---|
| Laya, L=93 | 0.086% |
| Laya, L=512 | 3.61% |
| Decisiones tipadas, L=1024 | 7.69% |
| Multilingüe, L=93 | 0.147% |
| Multilingüe, L=1024 | 11.92% |
Las estimaciones usan local_pairs = L*(2*r+1) - r*(r+1) para L>r, con r=64. Incluyen todas las proyecciones densas y ambas capas de cabezas de decisión de atención completa. Excluyen la generación de máscaras y el tráfico de memoria. El beneficio en runtime puede superar o quedar por debajo de la fracción de FLOPs, porque la atención y los GEMMs tienen eficiencias distintas; solo el perfilado puede establecerlo. Con 8192 tokens el trade-off sería distinto, pero los agentes distribuidos limitan las entradas a 512 o 1024, así que afirmar algo con 8192 tokens requeriría una carga de trabajo soportada por separado.
Fusión más allá de la compilación
El nn.gelu exacto instalado ya está decorado con compilación shapeless, y nn.Linear ya usa un addmm consciente del sesgo cuando procede. La compilación de todo el bloque aún puede fusionar el GELU con su multiplicación de gate, las sumas residuales, los casteos, las máscaras y las pequeñas operaciones de características de puntuación. Inspecciona el grafo de kernels compilado antes de implementar un kernel personalizado equivalente.
Si el tráfico de activaciones sigue siendo significativo, prototipa la fusión exacta GELU-y-gate o residual-y-LayerNorm. Conserva el GELU exacto basado en erf actual; una aproximación con tanh o sigmoide cambia el modelo y necesita mediciones de calidad aparte. Inspecciona los strides y las copias reales de Q/K/V antes de añadir conversiones de layout: la implementación de atención completa de MLX acepta una dimensión de cabeza contigua con otro striding y escribe un layout de salida conveniente para fusionar cabezas. Una copia contigua incondicional puede añadir trabajo.
El softmax de marcadores, la ordenación de los dos mayores, la entropía, la pequeña cabeza de acción y el formateo de resultados con NumPy son objetivos posteriores legítimos solo si se miden. Hay pocas posiciones de opción en comparación con cientos de operaciones de Transformer de ancho completo, así que optimizarlas primero probablemente no aborde la ruta dominante.
Proteger el significado del benchmark mientras se buscan ganancias agresivas
El generador de cargas de trabajo actual cicla tres definiciones de pregunta para construir 5, 10 o 50 preguntas. Hay como máximo tres entradas de modelo únicas en esos lotes. La deduplicación exacta por llamada puede evitar inferencia redundante en una aplicación real, pero mejoraría desproporcionadamente estos fixtures. Mantén esta característica separada de la optimización de kernels e informa de questions, unique_forward_inputs, los aciertos de caché y los tokens realmente evaluados. Reconstruye cada respuesta original usando sus propias etiquetas, criterios ordenados y metadatos de calibración. Conserva una suite con 50 preguntas genuinamente distintas y otra con entradas duplicadas deliberadamente.
No uses el cacheo de resultados entre llamadas para el benchmark principal de inferencia: llama repetidamente a exactamente la misma solicitud. Informa de cualquier experimento de caché como tal. Los benchmarks de caché del tokenizador deberían incluir tanto un escenario de rúbrica repetida como un escenario de entradas nuevas.
Para cada candidato, usa este diseño de experimento:
- Mantén el objetivo constante. Registra la revisión/hash del código fuente, la revisión del modelo, el dtype, las banderas de compilación, los módulos cuantizados seleccionados, los IDs de token o su hash, la forma, el recuento de marcadores, la política de batching, el warmup, la sincronización y el dispositivo. El
input_sha256existente hashea el estado/las preguntas en lugar de los tensores de token reales; no establece una identidad de tensores entre tokenizadores distintos. Vuelve a ejecutar la línea base desde la misma revisión final del código fuente, porque los hashes históricos del código fuente del JSON difieren. - Separa las clases de carga de trabajo. Usa formas fijas para comparaciones controladas de kernels; longitudes variables reales para planificación y compilación; preguntas distintas para el rendimiento; rúbricas repetidas para el cacheo legítimo en CPU; y preguntas duplicadas deliberadamente para la deduplicación. Incluye colas de longitud larga y distintos recuentos de opciones. Mantén el texto de origen y el truncamiento idénticos dentro de cada comparación de backends.
- Mide el comportamiento en frío y en caliente. Registra por separado la carga del modelo y la primera compilación. Mide la preparación, la construcción del grafo, el forward evaluado y sincronizado, la conversión de salida y la latencia de extremo a extremo sin tratar medianas independientes como aditivas. Perfila los kernels en una ejecución aparte, porque el trazado puede perturbar la latencia.
- Cambia una optimización a la vez. Criba con el recuento de iteraciones existente, luego repite los finalistas en bloques alternos de línea base/candidato y recoge suficientes muestras para un p95 creíble, por ejemplo al menos 200 solicitudes medidas por carga de trabajo. Usa un único benchmark de GPU activo, condiciones consistentes de energía/térmicas y conserva las muestras en bruto. Exige una mejora mayor que la variabilidad medida de las ejecuciones.
- Comprueba la corrección y la estabilidad. Compara con MLX al mismo dtype y con la referencia FP32 existente; impón la identidad de tokens para los cambios de planificación y de preparación en CPU. Ejercita longitudes alrededor de 64, 128 y los límites de buckets; lotes alrededor de los límites de fragmento; una y muchas opciones; entradas multilingües; padding; y cambios de forma tras la compilación. Repite solicitudes con forma cambiante, vigila tanto la memoria activa como la de caché después del warmup, y verifica resultados deterministas finitos dentro de cada configuración.
- Aplica puertas más estrictas a los cambios aproximados. La cuantización, la aproximación de activaciones, la poda de tokens, las salidas tempranas y la destilación requieren resultados de tarea y calibración reservados más allá del pequeño fixture de regresión. Conserva una identidad de modelo y una etiqueta de benchmark separadas al cambiar el comportamiento entrenado. Un checkpoint multilingüe o un modelo destilado más rápido es un modelo distinto, no una aceleración del mismo checkpoint Laya.
Para cualquier punto caliente medido que ocupe una fracción f del tiempo de extremo a extremo y se acelere por un factor s, usa la cota de Amdahl 1 / (1 - f + f/s) para evaluar el impacto total esperado. Usa una fracción de tiempo medida para f; las fracciones aritméticas de arriba no son sustitutas. La siguiente decisión de ingeniería concreta debería seguir a la ablación de compilación/batching y a un perfil de kernels de breve frente a largo, no a un multiplicador sin verificar.
Notas de documentación y reproducibilidad
La consulta de documentación usó el flujo de trabajo Context7 requerido: una resolución library MLX y luego consultas separadas de documentación oficial para compilación y atención rápida, usando /websites/ml-explore_github_io_mlx_build_html (tres comandos en total). Se inspeccionaron los archivos .pyi y las fuentes Python instalados para verificar el comportamiento de MLX 0.32.2, incluidos force_fused, las API de cuantización, el GELU compilado y el LayerNorm rápido. Se leyeron las fuentes C++/Metal upstream fijadas por versión para obtener detalles de despacho y del bucle de tiles. No se actualizó ninguna biblioteca para esta investigación.
Los dos archivos de línea base FP16 usados para las observaciones detalladas tenían estos digests SHA-256 en el momento de la inspección:
laya-mlx-float16.json
63146dd664d039dde1a728b17aad896e491bd01ea36ea0786953691180e55b09
laya-multilingual-mlx-float16.json
9af74bd5a11e4edc15e6a8c9dc929a7b9fd2d19cb06f348cd0e04a076f912473
El proceso principal de benchmark seguía produciendo artefactos adicionales durante esta investigación. La tabla cita deliberadamente archivos de línea base completos ya disponibles cuando se revisaron y no afirma nada sobre implementaciones candidatas no medidas.