Prefill Raso, Decodificação Profunda: Inferência Eficiente de Contexto Longo via Visibilidade KV Assimétrica por Camadas

Resumo

A inferência de contexto longo em modelos de linguagem apenas com decodificador é custosa porque prompts longos são processados durante o Pré-preenchimento, armazenados em cache em todas as camadas e repetidamente atendidos durante a Decodificação autoregressiva. Apresentamos o Shallow Prefill, dEEp Decode (SPEED), uma política de visibilidade KV assimétrica por fase que materializa estados KV de tokens de prompt não âncora apenas nas camadas inferiores, enquanto mantém os tokens da fase de Decodificação em profundidade total. Diferentemente de abordagens anteriores que tornam os estados KV de prompt das camadas superiores mais baratos para armazenar ou construir, o SPEED remove completamente os tokens de pré-preenchimento do conjunto de visibilidade da Decodificação nas camadas superiores. Com uma âncora BoS mínima, essa mudança simples preserva a ampla qualidade dos benchmarks, ao mesmo tempo que reduz o custo de contexto longo. Em um estudo controlado de ajuste de instruções com Llama-3.1-8B, o SPEED utilizando apenas 75% das camadas para tokens de pré-preenchimento alcançou uma pontuação média de 51,2 nos benchmarks no estilo OLMES, em comparação com 51,4 da linha de base em profundidade total, enquanto melhorou o TTFT em 33%, o TPOT em 22% e reduziu a memória KV ativa em 25,0% com contexto de 128K. Diagnósticos por camada sugerem que esse ponto de corte retém as principais regiões de seleção de prompt e estabilização de representação do modelo em profundidade total. Esses resultados mostram que tokens de prompt de contexto longo não precisam persistir sempre como objetos de cache KV em profundidade total quando os tokens da fase de Decodificação permanecem em profundidade total.

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.