Amal-David/mlx-porting-skill

Portabilidade e otimização de modelos MLX

Missão

Produzir uma implementação MLX correta, reproduzível e consciente da arquitetura. A correção vem antes da velocidade. Toda alegação de velocidade ou memória deve nomear o hardware, as versões de software, a carga de trabalho, a linha de base e o critério de qualidade.

Regras não negociáveis

1. Não executar código de modelo não confiável durante a ingestão. Inspecione JSON, cabeçalhos de safetensors, arquivos fonte e licenças estaticamente. Trate auto_map, módulos personalizados, hooks de instalação e trust_remote_code como portões de revisão.
2. Fixar a fonte. Registre repositório, revisão, arquivos de modelo, revisão do tokenizador/processador, licença e checksum ou manifesto de artefato.
3. Construir um oráculo de fonte antes de portar. Congele fixtures determinísticas e capture tensores intermediários em limites significativos.
4. Portar primeiro o menor caminho eager. Nenhuma quantização, compilação, kernels personalizados, batching ou decodificação especulativa até que a paridade básica seja aprovada.
5. Alterar uma dimensão de otimização por vez. Mantenha um registro de medição e reversão.
6. Preferir operações nativas do MLX. Tente operações fundidas nativas, mudanças de layout, design de cache e mx.compile antes de um kernel Metal personalizado.
7. Não traduzir folclore CUDA mecanicamente. Uma técnica CUDA é apenas um candidato a pesquisa até que seu gargalo Metal/MLX e implementação sejam demonstrados.
8. Nunca esconder regressões de qualidade atrás da taxa de transferência. Para modelos de áudio, linguagem, visão e generativos, use verificações de qualidade específicas da tarefa, além de tolerâncias de tensor.
9. Não publicar pesos convertidos sem verificações de licença e proveniência. Preserve o cartão de modelo original, atribuição, configuração de geração, arquivos de tokenizador/processador e receita de conversão.
10. A automação de pesquisa diária é apenas para revisão. Ela pode coletar e classificar candidatos, mas não deve reescrever silenciosamente runbooks ou mesclar recomendações.

Fluxo de trabalho

1. Inspecionar e classificar

Execute:

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

Leia intake e roteamento. Confirme:

• tarefa e domínio do modelo;
• família de arquitetura e estado recorrente/cache;
• framework fonte e operações personalizadas;
• contagem de parâmetros, dtypes, pesos compartilhados/amarrados, fragmentos e adaptadores;
• pré-processamento, tokenização, amostragem e pós-processamento;
• licença e risco de código remoto;
• Mac alvo, orçamento de memória, objetivo de latência/taxa de transferência e objetivo de qualidade.

Não inicie a implementação se a arquitetura, revisão da fonte ou alvo de avaliação permanecer ambíguo. Registre incertezas em PORT_PLAN.md em vez de adivinhar.

2. Selecionar a referência MLX comprovada mais próxima

Consulte mapa de suporte de modelos e assets/architectures.yaml.

Use esta ordem:

1. implementação oficial MLX / MLX-LM;
2. implementação ativa MLX-VLM ou MLX-Audio com testes;
3. implementação MLX de terceiros ativa;
4. implementação fonte upstream mais artigo da arquitetura;
5. protótipo de pesquisa que requer uma nova implementação MLX.

Reutilize padrões arquitetônicos, não suposições copiadas. Verifique semântica de configuração e layouts de tensor em relação à fonte fixada.

3. Estabelecer o oráculo da fonte

Siga paridade e testes:

• defina sementes determinísticas e modo de inferência;
• salve entradas exatas após o pré-processamento;
• capture forma/dtype/estatísticas e tensores selecionados após embeddings/frontends, cada grupo de bloco, gargalo, atualização de cache, logits/latentes e saída do decodificador/vocoder;
• salve saídas fonte em fixtures portáteis .npz, .json ou .wav;
• registre tolerâncias por dtype e operação;
• inclua pelo menos um caso mínimo, um comum e um limite.

4. Implementar o grafo MLX eager mínimo

Leia método de portabilidade principal e o runbook de arquitetura correspondente:

• Transformer decodificador denso
• Transformer Mixture-of-Experts
• Transformer codificador
• Transformer codificador-decodificador
• SSM, recorrente e híbrido
• Difusão e fluxo
• Visão-linguagem e omni
• Codec de áudio neural
• LM de áudio autoregressivo / TTS
• TTS de fluxo ou difusão
• Vocoder
• ASR
• Fala em streaming
• Separação e aprimoramento

Restrições de implementação inicial:

• apenas pesos de ponto flutuante;
• tamanho de lote um, a menos que batching seja intrínseco;
• sem decorador compile;
• sem kernels personalizados;
• objetos de estado e cache explícitos;
• asserções de forma em torno de cada transformação não trivial;
• manifesto de mapa de pesos reversível.

5. Converter pesos deterministicamente

Crie um mapa de pesos com chave fonte, chave alvo, forma fonte, forma alvo, transformação, dtype e regra de amarração/compartilhamento. Valide-o:

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

Nunca confie em strict=False no carregamento para ocultar pesos ausentes ou extras. Categorize toda exceção como intencionalmente ignorada, gerada, compartilhada ou não suportada.

6. Passar na escada de paridade

Use scripts/compare_tensors.py e passe, nesta ordem:

1. paridade de configuração e pré-processamento;
2. cobertura de pesos e paridade de forma transformada;
3. paridade de primitivo/bloco único;
4. paridade intermediária escalonada;
5. paridade determinística ponta a ponta;
6. paridade de qualidade da tarefa;
7. paridade de cache/estado e geração incremental;
8. paridade de limite, entrada longa, streaming e lote.

Quando a paridade falhar, use atlas de falhas. Não otimize um grafo com falha.

7. Perfilar antes de escolher otimizações

Leia benchmarking. Separe:

• tempo de prefill/encoder/frontend;
• tempo de decodificação por token ou por frame;
• tempo de codec/vocoder/pós-processamento;
• custo de compilação e primeira execução;
• memória de pico e estado estacionário;
• movimento de dados, sincronização e sobrecarga Python.

Use:

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

8. Aplicar a escada de otimização

Consulte assets/optimization_guidance.yaml, assets/recommendation-taxonomy.yaml, assets/techniques.yaml e carregue apenas guias relevantes:

1. Runtime, compilação e kernels MLX
2. Atenção e cache KV
3. Decodificação e serviço
4. Quantização
5. Treinamento e fine-tuning

Ordem padrão:

1. Remova avaliações não intencionais e transferências de host;
2. Corrija dtype e layout de tensor;
3. Use operações fundidas nativas e SDPA rápido;
4. Reduza alocações e torne as atualizações de estado/cache explícitas;
5. Compile regiões estáveis e controle a recompilação;
6. Faça chunk de prefill/frontends ou stream onde a arquitetura permitir;
7. Otimize política de cache/KV e batching;
8. Quantize pesos, depois KV/estado se justificado;
9. Adicione decodificação especulativa ou multi-token apenas para caminhos autoregressivos compatíveis;
10. Escreva um kernel Metal personalizado somente após a criação de perfil provar um gargalo de kernel remanescente.

Para cada alteração, registre hipótese, diff, resultado de correção, resultado de benchmark, resultado de memória, resultado de qualidade e decisão de manter/reverter.

9. Empacotar e publicar

Siga empacotamento e publicação. Inclua:

• fonte fixada e comando de conversão;
• versões compatíveis de MLX e bibliotecas;
• arquivos de modelo/config/tokenizador/processador;
• receita de quantização exata e exclusões;
• teste de fumaça determinístico;
• protocolo de benchmark e resultados brutos;
• limitações conhecidas;
• licença original e atribuição;
• nenhuma redação não suportada de "mais rápido" ou "sem perdas".

10. Retornar um relatório de engenharia

A resposta final deve resumir:

• arquitetura e runbook selecionado;
• revisão da fonte e referências de evidência;
• status de implementação e mapa de pesos;
• matriz de paridade;
• métricas de linha de base e otimizadas;
• otimizações aceitas e rejeitadas;
• riscos remanescentes e caminhos não suportados;
• comandos reproduzíveis e locais de artefatos.

Use os templates em assets/.

Quando parar

Pare e relate em vez de improvisar quando:

• a licença da fonte proibir a distribuição solicitada;
• o comportamento exigido da fonte existir apenas em código remoto não revisado;
• a paridade não puder ser localizada após as verificações escalonadas;
• uma otimização melhorar um microbenchmark, mas piorar a latência ponta a ponta, memória ou qualidade;
• um kernel personalizado carecer de uma alternativa portátil ou testes adequados;
• hardware, versão MLX, carga de trabalho ou alvo de qualidade estiver faltando em uma alegação de desempenho.

Manutenção

Para revisão de fonte, segurança e pipeline de candidatos diários, leia manutenção e proveniência. Execute scripts/audit_skill.py e scripts/validate_sources.py antes de distribuir esta habilidade.