Prellenado superficial, decodificación profunda: inferencia eficiente de contexto largo mediante visibilidad asimétrica de KV entre capas

Shallow Prefill, Deep Decoding: Efficient Long-Context Inference via Layer-Asymmetric KV Visibility

May 7, 2026

Autores: Jungsuk Oh, Hyeseo Jeon, Hyunjune Ji, Kyongmin Kong, Jay-Yoon Lee

cs.AI

Resumen

La inferencia de contexto largo en modelos de lenguaje solo con decodificador es costosa porque las indicaciones largas se procesan durante el Prellenado, se almacenan en caché en cada capa y se atienden repetidamente durante el Decodificado autorregresivo. Presentamos Shallow Prefill, dEEp Decode (SPEED), una política de visibilidad KV asimétrica en fase que materializa los estados KV de los tokens de indicación no ancla solo en las capas inferiores, mientras mantiene los tokens de la fase de Decodificado con profundidad completa. A diferencia de enfoques anteriores que hacen que los estados KV de las indicaciones en las capas superiores sean más baratos de almacenar o construir, SPEED elimina por completo los tokens de prellenado del conjunto de visibilidad de Decodificado en las capas superiores. Con un ancla BoS mínima, este cambio simple preserva la calidad general en los puntos de referencia mientras reduce el costo de contexto largo. En un estudio controlado de ajuste por instrucciones de Llama-3.1-8B, SPEED, que utiliza solo el 75% de las capas para los tokens de prellenado, alcanza una puntuación promedio de 51.2 en puntos de referencia estilo OLMES, en comparación con 51.4 para la línea base de profundidad completa, mientras mejora el TTFT en un 33%, el TPOT en un 22% y reduce la memoria KV activa en un 25.0% en un contexto de 128K. Los diagnósticos por capa sugieren que este corte retiene las principales regiones de selección de indicaciones y estabilización de representación del modelo de profundidad completa. Estos resultados muestran que los tokens de indicación de contexto largo no necesitan persistir siempre como objetos de caché KV de profundidad completa cuando los tokens de la fase de Decodificado permanecen con profundidad completa.

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.