Amal-David/mlx-porting-skill

Portabilidad y optimización de modelos MLX

Misión

Producir una implementación MLX correcta, reproducible y consciente de la arquitectura. La corrección es lo primero. Cada afirmación de velocidad o memoria debe nombrar el hardware, las versiones de software, la carga de trabajo, la línea base y la puerta de calidad.

Reglas no negociables

1. No ejecutar código de modelo no confiable durante la ingesta. Inspeccionar estáticamente JSON, encabezados safetensors, archivos fuente y licencias. Tratar auto_map, módulos personalizados, hooks de instalación y trust_remote_code como puertas de revisión.
2. Fijar el origen. Registrar repositorio, revisión, archivos de modelo, revisión del tokenizador/procesador, licencia y suma de verificación o manifiesto de artefactos.
3. Construir un oráculo fuente antes de portar. Congelar fixtures deterministas y capturar tensores intermedios en límites significativos.
4. Portar primero la ruta eager más pequeña. Sin cuantización, compilación, kernels personalizados, lotes o decodificación especulativa hasta que se cumpla la paridad básica.
5. Cambiar solo una dimensión de optimización a la vez. Mantener un registro de medición y reversión.
6. Preferir operaciones nativas de MLX. Probar operaciones fusionadas integradas, cambios de diseño, diseño de caché y mx.compile antes de un kernel Metal personalizado.
7. No traducir folklore de CUDA mecánicamente. Una técnica de CUDA es solo un candidato de investigación hasta que se demuestre su cuello de botella e implementación en Metal/MLX.
8. Nunca ocultar regresiones de calidad detrás del rendimiento. Para modelos de audio, lenguaje, visión y generativos, usar controles de calidad específicos de la tarea además de tolerancias de tensores.
9. No publicar pesos convertidos sin verificación de licencia y procedencia. Conservar la tarjeta del modelo original, atribución, configuración de generación, archivos de tokenizador/procesador y receta de conversión.
10. La automatización diaria de investigación es solo para revisión. Puede recopilar y clasificar candidatos, pero no debe reescribir silenciosamente manuales ni fusionar recomendaciones.

Flujo de trabajo

1. Inspeccionar y clasificar

Ejecutar:

python3 scripts/inspect_model.py MODEL_OR_DIRECTORY --output inspection.json
python3 scripts/make_port_plan.py inspection.json --output PORT_PLAN.md
python3 scripts/recommend_optimizations.py inspection.json --markdown OPTIMIZATIONS.md

Leer ingesta y enrutamiento. Confirmar:

• tarea y dominio del modelo;
• familia de arquitectura y estado/caché recurrente;
• framework fuente y operaciones personalizadas;
• número de parámetros, dtypes, pesos compartidos/atados, fragmentos y adaptadores;
• preprocesamiento, tokenización, muestreo y posprocesamiento;
• licencia y riesgo de código remoto;
• Mac de destino, presupuesto de memoria, objetivo de latencia/rendimiento y objetivo de calidad.

No comenzar la implementación si la arquitectura, la revisión fuente o el objetivo de evaluación siguen siendo ambiguos. Registrar las incertidumbres en PORT_PLAN.md en lugar de adivinar.

2. Seleccionar la referencia MLX probada más cercana

Consultar mapa de soporte de modelos y assets/architectures.yaml.

Usar este orden:

1. implementación oficial de MLX / MLX-LM;
2. implementación activa de MLX-VLM o MLX-Audio con pruebas;
3. implementación MLX de terceros activa;
4. implementación fuente upstream más artículo de arquitectura;
5. prototipo de investigación que requiere una nueva implementación MLX.

Reutilizar patrones arquitectónicos, no suposiciones copiadas. Verificar la semántica de configuración y los diseños de tensores contra la fuente fijada.

3. Establecer el oráculo fuente

Seguir paridad y pruebas:

• establecer semillas deterministas y modo de inferencia;
• guardar entradas exactas después del preprocesamiento;
• capturar forma/dtype/estadísticas y tensores seleccionados después de embeddings/frontends, cada grupo de bloques, cuello de botella, actualización de caché, logits/latentes y salida del decodificador/vocoder;
• guardar salidas fuente en fixtures portátiles .npz, .json o .wav;
• registrar tolerancias por dtype y operación;
• incluir al menos un caso mínimo, uno ordinario y uno límite.

4. Implementar el gráfico MLX eager mínimo

Leer método de portabilidad central y el manual de arquitectura correspondiente:

• Transformer decodificador denso
• Transformer de mezcla de expertos
• Transformer codificador
• Transformer codificador-decodificador
• SSM, recurrente e híbrido
• Difusión y flujo
• Visión-lenguaje y omni
• Códec de audio neuronal
• LM de audio autorregresivo / TTS
• TTS de flujo o difusión
• Vocoder
• ASR
• Habla en streaming
• Separación y mejora

Restricciones iniciales de implementación:

• solo pesos de punto flotante;
• tamaño de lote uno a menos que el lote sea intrínseco;
• sin decorador de compilación;
• sin kernels personalizados;
• objetos de estado y caché explícitos;
• aserciones de forma alrededor de cada transformación no trivial;
• manifiesto de mapa de pesos reversible.

5. Convertir pesos de manera determinista

Crear un mapa de pesos con clave fuente, clave objetivo, forma fuente, forma objetivo, transformación, dtype y regla de enlace/compartición. Validarlo:

python3 scripts/validate_weight_map.py   --source source-manifest.json   --target target-manifest.json   --mapping WEIGHT_MAP.json

Nunca confiar en strict=False en tiempo de carga para ocultar pesos faltantes o adicionales. Categorizar cada excepción como ignorada intencionalmente, generada, compartida o no soportada.

6. Subir la escalera de paridad

Usar scripts/compare_tensors.py y pasar, en orden:

1. paridad de configuración y preprocesamiento;
2. cobertura de pesos y paridad de forma transformada;
3. paridad de primitivo/bloque único;
4. paridad intermedia escalonada;
5. paridad determinista de extremo a extremo;
6. paridad de calidad de tarea;
7. paridad de caché/estado y generación incremental;
8. paridad de límite, entrada larga, streaming y lote.

Cuando falle la paridad, usar atlas de fallos. No optimizar un gráfico con fallos.

7. Perfilar antes de elegir optimizaciones

Leer evaluación comparativa. Separar:

• tiempo de prefill/codificador/frontend;
• tiempo de decodificación por token o por cuadro;
• tiempo de códec/vocoder/posprocesamiento;
• costo de compilación y primera ejecución;
• memoria máxima y en estado estable;
• movimiento de datos, sincronización y sobrecarga de Python.

Usar:

python3 scripts/benchmark_command.py --warmup 2 --runs 8 --output benchmark.json -- COMMAND ...

8. Aplicar la escalera de optimización

Consultar assets/optimization_guidance.yaml, assets/recommendation-taxonomy.yaml, assets/techniques.yaml y cargar solo las guías relevantes:

1. Tiempo de ejecución, compilación y kernels de MLX
2. Atención y caché KV
3. Decodificación y servicio
4. Cuantización
5. Entrenamiento y ajuste fino

Orden predeterminado:

1. eliminar evaluaciones no intencionadas y transferencias de host;
2. corregir dtype y diseño de tensores;
3. usar operaciones nativas fusionadas y SDPA rápida;
4. reducir asignaciones y hacer explícitas las actualizaciones de estado/caché;
5. compilar regiones estables y controlar la recompilación;
6. dividir prefill/frontends o transmitir donde la arquitectura lo permita;
7. optimizar la política de KV/caché y el lotes;
8. cuantizar pesos, luego KV/estado si está justificado;
9. agregar decodificación especulativa o de múltiples tokens solo para rutas autorregresivas compatibles;
10. escribir un kernel Metal personalizado solo después de que el perfilado demuestre un cuello de botella de kernel restante.

Para cada cambio, registrar hipótesis, diff, resultado de corrección, resultado de evaluación comparativa, resultado de memoria, resultado de calidad y decisión de mantener/revertir.

9. Empaquetar y publicar

Seguir empaquetado y publicación. Incluir:

• fuente fijada y comando de conversión;
• versiones compatibles de MLX y bibliotecas;
• archivos de modelo/configuración/tokenizador/procesador;
• receta de cuantización exacta y exclusiones;
• prueba de humo determinista;
• protocolo de evaluación comparativa y resultados sin procesar;
• limitaciones conocidas;
• licencia original y atribución;
• sin redacción no soportada de “más rápido” o “sin pérdidas”.

10. Devolver un informe de ingeniería

La respuesta final debe resumir:

• arquitectura y manual seleccionado;
• revisión fuente y referencias de evidencia;
• estado de implementación y mapa de pesos;
• matriz de paridad;
• métricas de línea base y optimizadas;
• optimizaciones aceptadas y rechazadas;
• riesgos restantes y rutas no soportadas;
• comandos reproducibles y ubicaciones de artefactos.

Usar las plantillas en assets/.

Cuándo detenerse

Detenerse e informar en lugar de improvisar cuando:

• la licencia fuente prohíbe la distribución solicitada;
• el comportamiento fuente requerido solo existe en código remoto no revisado;
• la paridad no se puede localizar después de las comprobaciones escalonadas;
• una optimización mejora un microbenchmark pero empeora la latencia de extremo a extremo, la memoria o la calidad;
• un kernel personalizado carece de una alternativa portátil o pruebas adecuadas;
• hardware, versión de MLX, carga de trabajo o objetivo de calidad faltan en una afirmación de rendimiento.

Mantenimiento

Para revisión de fuente, seguridad y el pipeline diario de candidatos, leer mantenimiento y procedencia. Ejecutar scripts/audit_skill.py y scripts/validate_sources.py antes de distribuir esta habilidad.