Flaches Prefill, tiefes Decoding: Effiziente Inferenz über lange Kontexte durch schichtasymmetrische KV-Sichtbarkeit

Zusammenfassung

Langkontext-Inferenz in Decoder-only-Sprachmodellen ist kostspielig, da lange Prompts während des Prefills verarbeitet, in jeder Schicht zwischengespeichert und während des autoregressiven Decodens wiederholt beachtet werden. Wir stellen Shallow Prefill, dEEp Decode (SPEED) vor, eine phasen-asymmetrische KV-Sichtbarkeitsstrategie, die Nicht-Anker-Prompt-Token-KV-Zustände nur in unteren Schichten materialisiert, während Token der Decode-Phase in voller Tiefe bleiben. Im Gegensatz zu früheren Ansätzen, die Prompt-KV-Zustände in oberen Schichten günstiger speichern oder konstruieren, entfernt SPEED Prefill-Tokens vollständig aus dem Decode-Sichtbarkeitsset der oberen Schichten. Mit einem minimalen BoS-Anker bewahrt diese einfache Änderung die breite Benchmark-Qualität und reduziert gleichzeitig die Langkontextkosten. In einer kontrollierten Llama-3.1-8B-Instruction-Tuning-Studie erreicht SPEED, das nur 75 % der Schichten für Prefill-Tokens verwendet, eine durchschnittliche Punktzahl von 51,2 bei OLMES-artigen Benchmarks, verglichen mit 51,4 für die Volltiefe-Baseline, während TTFT um 33 %, TPOT um 22 % verbessert und der aktive KV-Speicher bei 128K Kontext um 25,0 % reduziert wird. Schichtweise Diagnosen legen nahe, dass dieser Cutoff die Haupt-Prompt-Selektions- und Repräsentations-Stabilisierungsregionen des Volltiefe-Modells beibehält. Diese Ergebnisse zeigen, dass Langkontext-Prompt-Tokens nicht immer als KV-Cache-Objekte in voller Tiefe bestehen bleiben müssen, wenn Token der Decode-Phase in voller Tiefe bleiben.

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.