Préremplissage superficiel, décodage profond : inférence efficace pour longs contextes via une visibilité KV asymétrique selon les couches
Shallow Prefill, Deep Decoding: Efficient Long-Context Inference via Layer-Asymmetric KV Visibility
May 7, 2026
Auteurs: Jungsuk Oh, Hyeseo Jeon, Hyunjune Ji, Kyongmin Kong, Jay-Yoon Lee
cs.AI
Résumé
L'inférence à long contexte dans les modèles de langage à décodeur seul est coûteuse car de longues instructions sont traitées lors du Préremplissage, mises en cache à chaque couche et sollicitées de façon répétée pendant le Décodage autorégressif. Nous introduisons Shallow Prefill, dEEp Decode (SPEED), une politique de visibilité KV asymétrique selon la phase qui matérialise les états KV des jetons d'instruction non ancrés uniquement dans les couches inférieures, tout en maintenant les jetons de la phase de Décodage en profondeur intégrale. Contrairement aux approches antérieures qui rendent les états KV des instructions dans les couches supérieures moins coûteux à stocker ou à construire, SPEED retire totalement les jetons de préremplissage de l'ensemble de visibilité du Décodage dans les couches supérieures. Avec un ancrage BoS minimal, cette simple modification préserve une qualité comparable sur des références larges tout en réduisant le coût du long contexte. Dans une étude contrôlée sur Llama-3.1-8B ajusté par instructions, SPEED utilisant seulement 75 % des couches pour les jetons de préremplissage atteint un score moyen de 51,2 sur les références de type OLMES, contre 51,4 pour la ligne de base en profondeur intégrale, tout en améliorant le TTFT de 33 %, le TPOT de 22 % et en réduisant la mémoire KV active de 25,0 % pour un contexte de 128K. Les diagnostics par couche suggèrent que cette coupure préserve les régions principales de sélection d'instruction et de stabilisation de représentation du modèle en profondeur intégrale. Ces résultats montrent que les jetons d'instruction en long contexte n'ont pas besoin de persister en tant qu'objets de cache KV en profondeur intégrale lorsque les jetons de la phase de Décodage restent en profondeur intégrale.
English
Long-context inference in decoder-only language models is costly because long prompts are processed during Prefill, cached at every layer, and repeatedly attended to during autoregressive Decode. We introduce Shallow Prefill, dEEp Decode (SPEED), a phase-asymmetric KV-visibility policy that materializes non-anchor prompt-token KV states only in lower layers while keeping Decode-phase tokens full-depth. Unlike previous approaches that make upper-layer prompt KV states cheaper to store or construct, SPEED removes prefill tokens from the upper-layer Decode visibility set altogether. With a minimal BoS anchor, this simple change preserves broad benchmark quality while reducing long-context cost. In a controlled Llama-3.1-8B instruction-tuning study, SPEED using only 75\% of layers for prefill tokens reaches 51.2 average score on OLMES-style benchmarks, compared with 51.4 for the full-depth baseline, while improving TTFT by 33\%, TPOT by 22\%, and reducing active KV memory by 25.0\% at 128K context. Layer-wise diagnostics suggest that this cutoff retains the main prompt-selection and representation-stabilization regions of the full-depth model. These results show that long-context prompt tokens need not always persist as full-depth KV-cache objects when Decode-phase tokens remain full-depth.