Optimiser la charge de travail réelle de Snake

Le chemin optionnel adopté combine la compilation MLX, des buckets de longueur de 16 tokens et un cache borné de préfixes de questions tokenisées. Il ne change pas les poids, ne quantize pas le modèle, ne met pas en cache les prédictions et ne réutilise pas les états cachés de l’encodeur bidirectionnel entre les questions.

Dans un test apparié sur la boucle complète, le chemin optimisé livré a atteint 75.40 coups/seconde sur 2,400 coups, contre 70.82 coups/seconde pour l’inférence eager dans le même test : 1.065×, soit environ 6.5%. Les deux avaient zéro mort, 2 interventions de sécurité et des actions exécutées identiques sur 2,400/2,400 pas. La croissance combinée de la mémoire MLX active après warmup et nettoyage du cache était de 0 octet.

L’activer

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

L’API générale expose les mêmes contrôles optionnels :

import laya_mlx as laya

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

Les trois options sont désactivées par défaut, ce qui préserve le comportement eager existant et la configuration du benchmark. La compilation se spécialise pour une forme d’entrée ; la première utilisation et les nouvelles formes peuvent entraîner un coût de compilation. Construis un nouvel Agent après avoir changé les poids ou la structure des modules. Le padding arrondit la longueur de séquence au multiple demandé sans dépasser la limite de contexte configurée. Les masques excluent les tokens de padding.

cache_prompts=True garde au plus 128 préfixes immuables PreparedQuestion par Agent, y compris les positions des marqueurs. Les clés de cache incluent l’identité du tokenizer, les tokens spéciaux, le type de question, les options rendues ordonnées, les instructions et le budget de préfixe. L’état est assaini et tokenisé une fois par appel prepare, puis concaténé indépendamment avec chaque préfixe de question. Les changements de question créent ou sélectionnent le préfixe approprié. Chaque question reçoit encore une passe avant complète du modèle.

Ablation de forme et de préparation

Cette démo pose trois questions par coup : direction, estimation de route sûre et estimation de joignabilité de la nourriture. Sa taille de lot est donc de 3, plutôt que 1. Sur 32 plateaux réels enregistrés échantillonnés :

• Les longueurs de séquence multilingues étaient 59, 61, 63 et 64. Un multiple de 16 tokens les place toutes dans un bucket de 64 tokens.
• Les longueurs de séquence anglaises étaient 66, 68, 69 et 70, correspondant à un bucket de 80 tokens.
• Faire passer l’entrée multilingue de 64 à 96 par padding ajoute 50% de travail par token ; ce n’est pas le petit ajustement de 93 à 96 suggéré par le jeu de test distinct de l’API pour textes courts.

Les candidats ont tourné dans un ordre rotatif au sein de chaque état identique, après avoir visité une fois chaque forme mesurée. Le tableau contient la latence Agent.predict synchronisée, incluant la tokenisation et la conversion de sortie, et exclut le travail du planificateur/de l’interface et le warmup initial de forme.

Les sept candidats correspondaient aux directions proposées et exécutées en eager sur 32/32 plateaux par checkpoint. La différence maximale dans les probabilités et estimations affichées à quatre décimales était de 0 sur ces échantillons. C’est un accord sur des sorties arrondies à échantillon fini, non une revendication de tenseurs à virgule flottante internes identiques bit à bit.

L’ablation a utilisé des wrappers bornés de préparation de préfixe pour présélectionner des conceptions. Le test sur la boucle complète ci-dessous utilise l’implémentation réelle livrée de l’API compile, pad_to_multiple et cache_prompts. Ses tests de correction comparent en outre les IDs préparés et les marqueurs sous troncature d’état, critères changeants, assainissement des masques et éviction du cache.

Le chemin optimisé livré a aussi passé la matrice complète de validation sur checkpoint réel : 63/63 d’accord sur la réponse sélectionnée pour chacun des trois checkpoints en FP32 et FP16 (378/378 au total). Les erreurs de probabilité calibrée sont restées dans les tolérances existantes. Chaque configuration a passé 10 appels répétés finis et déterministes supplémentaires avec 0 octet de croissance mesurée de la mémoire active. Données de validation optimisées. Les résultats à 100 répétitions par configuration du chemin eager d’origine restent dans le rapport de benchmark d’origine.

Pour l’anglais, la compilation avec la longueur de séquence réelle était meilleure que de forcer le bucket plus grand dans cet échantillon. La démo est par défaut en multilingue ; les utilisateurs de l’API générale peuvent laisser pad_to_multiple=None tout en activant la compilation et la réutilisation de préfixe.

Test apparié sur la boucle complète

Quatre graines, 600 coups chacune, avec un ordre de candidats alternant par graine. Le rendu inclut la composition Rich en truecolor et la sérialisation ANSI, et exclut le rendu de l’émulateur de terminal. Chaque coup effectue une prédiction neuve. Les résultats proviennent d’une seule exécution appariée locale.

Le chemin optimisé était plus lent sur une graine. Par conséquent, 6.5% est l’amélioration combinée dans cette exécution mesurée, plutôt qu’une amélioration garantie pour chaque épisode ou machine. Le large balayage de vitesse antérieur et ce test apparié plus tardif sont des exécutions différentes ; leurs fréquences absolues ne doivent pas être soustraites pour revendiquer un gain de vitesse. Données de boucle complète.

Choisir le modèle d’après le jeu autant que la latence

Les deux checkpoints ont tourné 20 graines appariées × 300 coups, avec l’ordre des checkpoints alternant à chaque graine. Chaque épisode a utilisé le même état initial, la même graine RNG de nourriture, les mêmes descriptions compactes de caractéristiques et le même bouclier de cycle. L’horizon est fixe ; ce sont des scores après 300 coups, pas des parties complètes se terminant par une mort ou un plateau rempli. Aucun rendu de terminal n’a été inclus dans cette comparaison de modèles.

Le multilingue a progressé davantage vers la nourriture et était plus rapide sur cette charge, il reste donc le checkpoint de démo par défaut. Le résultat évalue cette politique assistée par des caractéristiques, pas la qualité de raisonnement générale ni un modèle Snake sans assistance. Chaque épisode et inférence.

Reproduire

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

Exécute les mesures GPU séquentiellement. Télécharge d’abord les deux répertoires de modèles locaux. L’enregistrement source versionné fournit les états de plateau échantillonnés exacts. Ablations brutes : multilingue, anglais, pilote initial à 96 tokens.

La comparaison antérieure entre prompt compact et détaillé alternait l’ordre des prompts sur 64 états et a trouvé une médiane de 11.80 → 9.29 ms après avoir écarté les 8 premières itérations de warmup. Son enregistrement ad hoc d’origine ne stockait pas les instantanés de plateau, il s’agit donc d’une preuve à l’appui plutôt que de l’ablation reproductible principale. python -m benchmarks.snake_prompt fournit une version reproductible qui stocke les états, la graine, les décisions complètes et la méthode.

L’implémentation suit le guide officiel de compilation de MLX : utilise un callable compilé à longue durée de vie et une spécialisation de forme normale. Elle n’utilise pas shapeless=True sur du code modèle Python dépendant de la forme. La documentation actuelle a été vérifiée via le site officiel après que les requêtes du CLI Context7 ont échoué avec des erreurs réseau.