Optimizar la carga de trabajo real de Snake

La ruta opt-in adoptada combina la compilación de MLX, buckets de longitud de 16 tokens y una caché acotada de prefijos de preguntas tokenizadas. No cambia los pesos, no cuantiza el modelo, no cachea predicciones ni reutiliza estados ocultos del codificador bidireccional entre preguntas.

En una prueba emparejada de bucle completo, la ruta optimizada incluida alcanzó 75.40 movimientos/segundo en 2,400 movimientos, frente a 70.82 movimientos/segundo de la inferencia eager en la misma prueba: 1.065×, o alrededor de 6.5%. Ambas tuvieron cero muertes, 2 intervenciones de seguridad y acciones ejecutadas idénticas en 2,400/2,400 pasos. El crecimiento combinado de memoria MLX activa tras el warmup y la limpieza de caché fue de 0 bytes.

Cómo habilitarla

laya-snake --optimize
laya-snake --optimize --max-speed

La API general expone los mismos controles opt-in:

import laya_mlx as laya

agent = laya.load(
    "aac6fef/laya-multilingual-mlx",
    compile=True,
    pad_to_multiple=16,
    cache_prompts=True,
)

Las tres opciones están desactivadas por defecto, lo que preserva el comportamiento eager existente y la configuración del benchmark. La compilación se especializa por forma de entrada; el primer uso y las formas nuevas pueden acarrear un coste de compilación. Construye un Agent nuevo tras cambiar los pesos o la estructura de módulos. El padding redondea la longitud de secuencia hacia arriba hasta el múltiplo solicitado sin superar el límite de contexto configurado. Las máscaras excluyen los tokens de padding.

cache_prompts=True conserva como máximo 128 prefijos PreparedQuestion inmutables por Agent, incluidas las posiciones de los marcadores. Las claves de caché incluyen la identidad del tokenizador, los tokens especiales, el tipo de pregunta, las opciones renderizadas en orden, las instrucciones y el presupuesto de prefijo. El estado se sanea y se tokeniza una vez por llamada a prepare, y luego se concatena de forma independiente con cada prefijo de pregunta. Los cambios de pregunta crean o seleccionan el prefijo adecuado. Cada pregunta sigue recibiendo una pasada completa del modelo.

Ablación de forma y preparación

Esta demo hace tres preguntas por movimiento: dirección, estimación de ruta segura y estimación de alcanzabilidad de la comida. Por tanto, su tamaño de lote es 3, no 1. En 32 tableros reales grabados y muestreados:

• Las longitudes de secuencia multilingües fueron 59, 61, 63 y 64. Un múltiplo de 16 tokens las coloca todas en un bucket de 64 tokens.
• Las longitudes de secuencia en inglés fueron 66, 68, 69 y 70, que van a un bucket de 80 tokens.
• Rellenar la entrada multilingüe de 64 a 96 añade 50% al trabajo por token; no es el pequeño ajuste de 93 a 96 sugerido por el fixture independiente de la API de texto breve.

Los candidatos se ejecutaron en orden rotatorio dentro de cada estado idéntico, después de visitar una vez cada forma medida. La tabla contiene la latencia sincronizada de Agent.predict, incluida la tokenización y la conversión de salida, y excluye el trabajo del planificador/UI y el warmup inicial de formas.

Los siete candidatos coincidieron con las direcciones propuestas y ejecutadas del modo eager en 32/32 tableros por checkpoint. La diferencia máxima en las probabilidades y estimaciones mostradas, con cuatro decimales, fue 0 en estas muestras. Esto es una coincidencia de salida redondeada en una muestra finita, no una afirmación de tensores internos de coma flotante idénticos bit a bit.

La ablación usó envoltorios acotados de preparación de prefijos para cribar diseños. La prueba de bucle completo de abajo usa la implementación real de la API compile, pad_to_multiple y cache_prompts incluida. Sus pruebas de corrección comparan además los IDs y marcadores preparados bajo truncamiento del estado, criterios cambiantes, saneado de máscaras y desalojo de caché.

La ruta optimizada incluida también superó la matriz completa de validación con checkpoints reales: 63/63 de coincidencia de respuesta seleccionada para cada uno de los tres checkpoints en FP32 y FP16 (378/378 en total). Los errores de probabilidad calibrada se mantuvieron dentro de las tolerancias existentes. Cada configuración superó 10 llamadas adicionales, finitas y deterministas repetidas con 0 bytes de crecimiento medido de memoria activa. Datos de validación de la ruta optimizada. Los resultados de 100 repeticiones por configuración de la ruta eager original permanecen en el informe de benchmark original.

Para el inglés, la compilación con la longitud de secuencia real fue mejor que forzar el bucket más grande en esta muestra. La demo usa el multilingüe por defecto; los usuarios de la API general pueden dejar pad_to_multiple=None mientras habilitan la compilación y la reutilización de prefijos.

Prueba emparejada de bucle completo

Cuatro semillas, 600 movimientos cada una, con el orden de los candidatos alternándose por semilla. La renderización incluye composición con Rich en color verdadero y serialización ANSI, y excluye el pintado del emulador de terminal. Cada movimiento realiza una predicción nueva. Los resultados provienen de una única ejecución emparejada local.

La ruta optimizada fue más lenta en una semilla. En consecuencia, 6.5% es la mejora combinada en esta ejecución medida, no una mejora garantizada para todos los episodios ni todas las máquinas. El barrido de velocidad amplio anterior y esta prueba emparejada posterior son ejecuciones distintas; sus tasas absolutas no deben restarse para afirmar una mejora. Datos completos del bucle.

Elegir el modelo usando tanto el juego como la latencia

Ambos checkpoints ejecutaron 20 semillas emparejadas × 300 movimientos, con el orden de los checkpoints alternándose en cada semilla. Cada episodio usó el mismo estado inicial, la misma semilla RNG de comida, las mismas descripciones compactas de características y el mismo escudo de ciclo. El horizonte es fijo; son puntuaciones tras 300 movimientos, no partidas completas que terminan en muerte o con el tablero lleno. En esta comparación de modelos no se incluyó renderización de terminal.

El multilingüe logró más progreso hacia la comida y fue más rápido en esta carga de trabajo, así que sigue siendo el checkpoint por defecto de la demo. El resultado evalúa esta política asistida por características, no la calidad general de razonamiento ni un modelo de Snake sin ayuda. Cada episodio e inferencia.

Reproducir

uv run --extra demo python -m experiments.snake_runtime \
  --output artifacts/snake/runtime-multilingual.json
uv run --extra demo python -m experiments.snake_runtime \
  --model models/hub/laya-mlx --bucket 80 \
  --output artifacts/snake/runtime-english.json
uv run --extra demo python -m benchmarks.snake_optimized \
  --output artifacts/snake/optimized-paired.json
uv run --extra demo python -m benchmarks.snake_models \
  --episodes 20 --steps 300 --output artifacts/snake/models.json

Ejecuta las mediciones de GPU de forma secuencial. Descarga primero los dos directorios de modelos locales. La grabación de origen incluida en el repositorio proporciona los estados de tablero muestreados exactos. Ablaciones en bruto: multilingüe, inglés, piloto inicial de 96 tokens.

La comparación anterior entre prompt compacto y detallado alternó el orden de los prompts en 64 estados y encontró una mediana de 11.80 → 9.29 ms tras descartar las primeras 8 iteraciones de warmup. Su registro ad hoc original no almacenaba instantáneas del tablero, así que es evidencia de apoyo más que la ablación reproducible principal. python -m benchmarks.snake_prompt proporciona una versión reproducible que almacena estados, semilla, decisiones completas y método.

La implementación sigue la guía oficial de compilación de MLX: usa un callable compilado de larga vida y una especialización de formas normal. No usa shapeless=True en código de modelo Python dependiente de la forma. La documentación actual se comprobó a través del sitio oficial después de que las solicitudes de la CLI de Context7 fallaran con errores de red.