Titre original : LoopVQ : Quantification du cache KV en 1 bit pour les Transformers en boucle
Les Transformers en boucle réutilisent le même ensemble de couches, permettant d'égaler des modèles standards plusieurs fois plus grands avec moins de paramètres, mais chaque étape de boucle doit stocker son propre cache KV, qui croît linéairement avec le nombre d'étapes, devenant ainsi le principal goulot d'étranglement mémoire. Les auteurs ont observé une forte redondance du KV le long de l'axe de la boucle, ainsi qu'une sensibilité très variable à la quantification selon les étapes de boucle. Ils proposent donc LoopVQ : regrouper les canaux de plusieurs étapes de boucle consécutives en un sous-vecteur partageant un seul index de codebook (quantification vectorielle sur l'axe de la boucle), puis affiner le codebook k-means pondéré par Fisher pour que la sortie du modèle quantifié se rapproche de celle du modèle en précision complète (chaque étape de boucle lit le cache quantifié pendant l'entraînement). Cette méthode est orthogonale aux approches destinées aux LLM standards et peut être combinée avec la transformation des clés de TaSQ.
J'aime 8 · Étoiles du code 3
Raisons
- Les auteurs affirment qu'avec 0.5 bits par élément, la méthode surpasse les méthodes existantes de quantification vectorielle du cache KV destinées aux LLM standards sur les tâches de raisonnement général, de chaîne de pensée longue et de récupération de contexte long pour les modèles Ouro et Huginn ; combinée à TaSQ, elle conserve la majeure partie de la précision BF16 même à 0.25 bits.
- Les gains système sont concrets : après implémentation dans SGLang par les auteurs, la latence de décodage est presque divisée par deux pour un contexte de 32K, et le débit de service maximal pour Ouro 1.4B / 2.6B est respectivement multiplié par 4.3 et 6.1 par rapport à la base BF16.
- Le code est sous licence Apache-2.0 : le processus de calibration, la branche SGLang avec stockage d'index packés et noyaux de décodage fusionnés, ainsi que les scripts de reproduction pour chaque tableau du papier sont inclus dans le dépôt.
Utilisation
- ./setup_envs.sh crée trois environnements (loopvq-serve pour le moteur, loopvq-eval pour l'évaluation, loopvq-build pour la calibration), exportez MODELS_DIR vers Ouro-1.4B, Ouro-2.6B, Ouro-*-Thinking et huginn-0125 ; pour reproduire le tableau général : MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.sh, puis MODEL=ouro_1p4b bash experiments/loopvq_general/run_all.sh, et enfin python3 experiments/loopvq_general/make_table.py.
- Démarrage du service : dans l'environnement loopvq-serve, exécutez PREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <dossier du codebook> --port 30000 --disable-piecewise-cuda-graph, ce qui fournit une interface compatible OpenAI via SGLang ; les tokens de préfixe et récents restent en bf16.
Risques
- Applicable uniquement aux Transformers en boucle (le papier utilise Ouro 1.4B / 2.6B et Huginn-3.5B) ; ne peut pas être appliqué directement aux LLM standards ; la calibration n'utilise que 16 séquences GPQA.
- Les chiffres de latence et de débit proviennent de l'environnement des auteurs ; les likes 8 et les étoiles du code 3 (représentant uniquement l'attention).
Ouvrir · Code · Page du papier