ReplaySSM: Zustandsspeicherung neu denken für schnellere Dekodierung hybrider Modelle

ReplaySSM: State-Speicherung neu denken für schnellere Dekodierung hybrider Modelle

Hybride Sprachmodelle – Architekturen, die Attention-Schichten mit State-Space- oder linear-rekurrenten Schichten verschränken – versprachen das Beste aus zwei Welten: die exakte Erinnerungsfähigkeit von Attention und das Streaming-Verhalten mit konstantem Speicher von Rekurrenz. In der Praxis stellen viele Teams fest, dass das Versprechen während des Prefill hält und dann während des Decode still erodiert. Der Durchsatz pro Ausgabe-Token sinkt, die Speicherbilanz wird schwer nachvollziehbar, und Funktionen, die für reine Transformer Routine sind – Beam Search, spekulative Dekodierung, viele Samples pro Prompt – werden umständlich.

Der NVIDIA-Blogbeitrag „ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding“ (vom 12. Oktober 2026) untersucht diese Spannung und schlägt vor, State-Speicherung als Designanliegen erster Klasse statt als Implementierungsdetail zu behandeln. Dieser Artikel erklärt den Problemraum, führt durch ein praktisches Setup zum Experimentieren mit hybrider Dekodierung auf eigener Hardware und skizziert, was gemessen werden sollte. Die konzeptionelle Erklärung hier ist Interpretation; die Details des vorgeschlagenen Verfahrens sollten dem Originalbeitrag entnommen werden.

Warum hybride Modelle attraktiv – und umständlich zu dekodieren sind

Eine Transformer-Schicht hält einen Key/Value-Cache, der linear mit der Sequenzlänge wächst. Eine SSM-Schicht hält einen rekurrenten Zustand fester Größe, der von jedem Token neu geschrieben wird. Stapelt man beide zusammen, erhält man ein Modell, das präzise dort attenden kann, wo es darauf ankommt, und überall sonst günstig streamt.

Der Haken ist, dass die beiden Speichersysteme fast nichts gemeinsam haben:

• KV-Cache ist append-only, pro Token und inhaltsadressiert nach Position. Man kann ihn zerschneiden, paginieren, Präfixe teilen und ein Suffix verwerfen, ohne den Rest anzufassen.
• SSM-State ist eine laufende Zusammenfassung. Er hat keinen Positionsindex. Token t erzeugt einen neuen Zustand, der den alten überschreibt, und dieser Zustand ist eine komprimierte Funktion von allem davor.

Diese Asymmetrie ist die Wurzel des Dekodierungsproblems. Ein KV-Cache ist eine Datenstruktur, die man manipulieren kann. Ein rekurrenter Zustand ist ein Wert, den man entweder bewahren oder neu berechnen muss.

Warum Decode und nicht Prefill der schwierige Teil ist

Prefill verarbeitet einen ganzen Prompt in einem Durchgang. Jedes Token ist gleichzeitig verfügbar, die GPU ist mit Matmul-Arbeit gesättigt, und die Kosten für das Schreiben von State werden über eine lange Sequenz amortisiert. Decode ist das Gegenteil: ein Token pro Schritt, winzige Matmuls und ein speichergebundener Kernel, dessen Laufzeit vom Lesen und Aktualisieren von State dominiert wird.

Bei einem reinen Attention-Modell ist die Decode-Leistung weitgehend eine Frage der KV-Cache-Größe und davon, wie effizient man ihn lesen kann. Quantisierung, Paging und Prefix-Sharing setzen alle dort an. Bei einem hybriden Modell erbt man all diese Anliegen plus eine zweite, strukturell andere Familie von State, die nicht auf dieselbe Weise paginiert oder per Prefix geteilt werden kann.

Der State-Speicherungs-Engpass bei verzweigten Workloads

Das Problem wird akut, wann immer ein Präfix mehrere Fortsetzungen hervorbringt:

• Beam Search hält k Hypothesen am Leben, jede mit eigenem State.
• Spekulative Dekodierung schlägt mehrere Draft-Token vor, verifiziert sie parallel und muss die abgelehnten zurückrollen.
• Best-of-n-Sampling erzeugt n unabhängige Vervollständigungen aus einem gemeinsamen Prompt.
• Baumstrukturiertes Reasoning erkundet Zweige, die sich aufteilen und wieder zusammenlaufen.

Für Attention-Schichten sind all diese Dinge handhabbar. Das Zurückrollen eines abgelehnten Drafts ist eine Sache des Abschneidens des KV-Cache. Ein Prompt-Präfix zu teilen bedeutet, seine gecachten Keys und Values zu teilen. Für SSM-Schichten ist keine der beiden Operationen kostenlos. Ein Verzweigungspunkt bedeutet, dass der State für jeden Zweig unabhängig existieren muss, und ein Rollback bedeutet, dass der State auf einen früheren Wert zurückgesetzt werden muss, der nicht mehr existiert, es sei denn, man hat ihn gespeichert.

Es gibt nur drei Auswege, und jeder hat Kosten:

1. Aggressiv snapshotten. Behalten Sie eine Kopie des rekurrenten Zustands an jedem Punkt, zu dem Sie möglicherweise zurückkehren müssen. Der Speicher skaliert mit Verzweigungsfaktor mal Checkpoint-Frequenz.
2. Aus einem Checkpoint neu berechnen. Speichern Sie State selten und führen Sie die Schichten ab dem letzten Checkpoint erneut aus, wenn Sie einen Zwischenzustand benötigen. Speicher ist gering; Latenz ist hoch.
3. Die Speicherung selbst umstrukturieren. Ändern Sie, was gespeichert wird, wann und in welcher Form, sodass Rekonstruktion billig genug ist, um der Standardweg zu sein.

Was ein Replay-First-Design impliziert

Der Name des Ansatzes verweist auf die dritte Option. „Replay“ legt nahe, dass Rekonstruktion kein Notfall-Fallback mehr ist, sondern der normale Pfad, und dass die Speicherschicht darum herum entworfen ist, Replay billig zu machen.

So gelesen verschiebt sich die Designfrage von „Wie viel State behalten wir?“ zu „Was ist das Minimum, das wir behalten müssen, damit jeder benötigte State schnell neu aufgebaut werden kann?“ Diese Rahmung hat mehrere Konsequenzen, über die es sich nachzudenken lohnt:

• Speicherung wird zu einer abstimmbaren Achse statt zu fixen Kosten. Man kann Speicher gegen Neuberechnung entlang einer kontinuierlichen Skala eintauschen und pro Workload einen Punkt wählen.
• Die Decode-Schleife erhält eine neue primitive Operation: State auf Position p wiederherstellen. Ob dies durch Laden eines Snapshots oder durch Replay eines kurzen Suffix implementiert wird, ändert die Latenz, nicht die Semantik.
• Serving-Systeme können State-Budgets so zuweisen, wie sie bereits KV-Cache-Budgets zuweisen – als Ressource pro Anfrage mit Eviction-Richtlinien statt als unvermeidbare Zuweisung pro Sequenz.
• Korrektheit lässt sich leichter begründen, weil ein wieder abgespielter State numerisch identisch (oder nahezu identisch) mit dem gespeicherten sein sollte. Das ist eine testbare Eigenschaft, anders als die unscharfen Qualitätsabwägungen der Quantisierung.

Ob die veröffentlichte Methode dieser Beschreibung im Detail entspricht, beantwortet die Quelle direkt; die obige Argumentation ist ein Rahmen zum Lesen derselben, kein Ersatz dafür.

Voraussetzungen

Bevor Sie experimentieren, vergewissern Sie sich, dass Ihre Umgebung hybride Dekodierung überhaupt unterstützen kann. Die realistischen Voraussetzungen sind:

• Eine CUDA-fähige GPU mit genügend Gerätespeicher für die Gewichte plus den größten State-Footprint, den Sie zulassen möchten. Beginnen Sie mit einem Modell, das bequem passt, und verkleinern Sie dann das Budget.
• Ein aktueller Python-3.x-Interpreter, der von Ihrer gewählten Framework-Version unterstützt wird. Bleeding-Edge-Python mit gepinnten CUDA-Wheels zu mischen, ist die häufigste Ursache für Setup-Fehler.
• PyTorch, das für Ihren CUDA-Treiber gebaut ist. Eine Versionsabweichung zwischen Treiber und Wheel erzeugt verwirrende Laufzeitfehler statt sauberer Fehlschläge.
• Ein Framework, das hybride Architekturen versteht. Neuere transformers-Releases unterstützen hybride Cache-Objekte; Serving-Engines wie vLLM und TensorRT-LLM behandeln hybride Modelle mit unterschiedlichem Reifegrad. Prüfen Sie die Unterstützung für Ihr spezifisches Modell, bevor Sie benchmarken.
• Genügend Host-RAM und Speicherplatz, wenn Sie State checkpointen oder auslagern möchten. Ein State-Tensor, der auf der GPU klein ist, kann groß werden, wenn er über Sequenzen und Checkpoints multipliziert wird.

Genaue Versionsanforderungen für die ReplaySSM-Implementierung, falls eine veröffentlicht wird, gehören zum Originalbeitrag und den dazugehörigen Artefakten.

Schritt-für-Schritt-Installation

Die folgenden Befehle bauen eine generische Umgebung für Inferenz mit hybriden Modellen auf. Sie installieren kein spezifisches ReplaySSM-Paket – prüfen Sie die Quelle für den offiziellen Weg.

Vergewissern Sie sich zunächst, dass die GPU sichtbar und der Treiber gesund ist.

nvidia-smi

Erstellen Sie eine isolierte Umgebung, damit Framework-Versionen nicht mit anderem auf dem Rechner kollidieren.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip wheel

Installieren Sie PyTorch. Wählen Sie die Index-URL, die zu Ihrer CUDA-Runtime passt; das Standard-PyPI-Wheel kann auf eine andere CUDA-Version als Ihr Treiber abzielen.

pip install torch --index-url https://download.pytorch.org/whl/cu124

Bestätigen Sie, dass PyTorch die GPU sieht, bevor Sie etwas anderes installieren, denn jeder spätere Fehler lässt sich von einer bekanntermaßen funktionierenden Basis leichter diagnostizieren.

python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

Installieren Sie die Modellierungs- und Serving-Bibliotheken, die Sie vergleichen möchten.

pip install transformers accelerate
pip install vllm

Wenn Ihr hybrides Modell Mamba-artige Schichten verwendet, werden Kernel-Bibliotheken zur Beschleunigung linear-rekurrenter Schichten häufig zusammen mit dem Hauptframework verwendet. Installieren Sie sie nur, wenn das von Ihnen ausgeführte Modell sie als Abhängigkeit dokumentiert.

pip install flash-linear-attention causal-conv1d

Zeichnen Sie schließlich die genauen Versionen in einer Lockdatei auf, damit ein Benchmark-Ergebnis später reproduziert werden kann.

pip freeze > requirements-lock.txt

Konfiguration

Die Konfiguration für hybride Dekodierung fällt in zwei Kategorien: Schalter, die Ihr Framework bereits offenlegt, und etwaige Switches, die die ReplaySSM-Implementierung hinzufügt. Die erste Kategorie ist heute nutzbar.

Stellen Sie den Allocator so ein, dass Fragmentierung reduziert wird, wenn State-Tensoren unterschiedlicher Größe über Anfragen hinweg allokiert und freigegeben werden.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Binden Sie den Lauf an eine einzelne GPU, damit Messungen nicht durch Platzierungsentscheidungen verfälscht werden.

export CUDA_VISIBLE_DEVICES=0

Aktivieren Sie beim Serving Prefix-Caching und Chunked Prefill, wenn Ihre Engine sie unterstützt. Beide ändern, wie State geteilt und geplant wird, und beide interagieren mit jeder State-Speicherungsstrategie.

vllm serve YOUR_HYBRID_MODEL_ID \
  --enable-prefix-caching \
  --enable-chunked-prefill \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.90

Für die zweite Kategorie konsultieren Sie die Quelle. Jeder Konfigurationsblock für ReplaySSM sollte aus der offiziellen Dokumentation gelesen und nicht abgeleitet werden, weil Feldnamen und Standardwerte bestimmen, ob Sie das Feature überhaupt messen.

Nutzungsbeispiele

Beispiel 1: Baseline-Decode-Durchsatz

Etablieren Sie eine Baseline, bevor Sie etwas ändern. Dieses Skript dekodiert eine feste Anzahl von Token und meldet Token pro Sekunde, die Metrik, die am empfindlichsten auf State-Handling-Overhead reagiert.

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "YOUR_HYBRID_MODEL_ID"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="cuda"
).eval()

prompt = "Explain why recurrent state is expensive to branch."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

steps = 128
with torch.no_grad():
    out = model(**inputs, use_cache=True)
    next_id = out.logits[:, -1].argmax(-1, keepdim=True)

    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(steps):
        out = model(
            input_ids=next_id,
            past_key_values=out.past_key_values,
            use_cache=True,
        )
        next_id = out.logits[:, -1].argmax(-1, keepdim=True)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

print(f"{steps / elapsed:.1f} tokens/sec")

Beispiel 2: Den State-Footprint messen

Der Spitzenspeicher zeigt, wie viel Spielraum für größere Batches oder mehr Verzweigungen bleibt. Setzen Sie die Statistiken vor dem gemessenen Bereich zurück, sodass Sie nur erfassen, was die Decode-Schleife allokiert.

import torch

torch.cuda.reset_peak_memory_stats()
# ... führe die Decode-Schleife aus Beispiel 1 aus ...
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"peak allocated: {peak_gib:.2f} GiB")

Ziehen Sie die Gewichte von dieser Zahl ab, um den State-Footprint zu schätzen, und teilen Sie dann durch die Anzahl gleichzeitiger Sequenzen, um einen Wert pro Sequenz zu erhalten. Diese Zahl bestimmt Ihre maximale Batch-Größe unter einer State-Speicherungsstrategie.

Beispiel 3: Spekulative Dekodierung und Rollback

Dies ist der Workload, bei dem die Speicherrichtlinie am meisten zählt. Die folgende Struktur ist schematisch; passen Sie die Cache-Aufrufe an die API Ihres Frameworks an.

# 1. Eine Referenz auf den State nach dem akzeptierten Präfix speichern.
state_checkpoint = cache.snapshot()

# 2. k Draft-Token vorschlagen und sie in einem Durchgang durch das Modell laufen lassen.
draft_ids = drafter.generate(prefix_ids, k=4)
candidate_logits = target_model.forward_with_cache(draft_ids, cache)

# 3. Das längste Präfix akzeptieren, das zur Zielverteilung passt.
accepted = longest_matching_prefix(candidate_logits, draft_ids)

# 4. Wenn nichts akzeptiert wurde, wiederherstellen statt das ganze Präfix erneut auszuführen.
if accepted == 0:
    cache.restore(state_checkpoint)
else:
    cache.truncate_to(accepted)

Die Aufrufe snapshot und restore sind der Punkt, an dem die beiden Strategien divergieren. Eine snapshot-basierte Implementierung zahlt Speicher in Schritt 1. Eine replay-basierte Implementierung zahlt Neuberechnung in Schritt 4. Die richtige Wahl hängt von Ihrer Akzeptanzrate ab: Bei hoher Akzeptanzrate sind Rollbacks selten und billiges Replay gewinnt; bei niedriger Akzeptanzrate können sich Snapshots bezahlt machen.

Was vorher und nachher gemessen werden sollte

Eine Änderung der State-Speicherung lohnt sich nur, wenn sie eine Metrik bewegt, die Ihnen wichtig ist. Verfolgen Sie diese konsistent:

• Zeit pro Ausgabe-Token, nicht nur die gesamte Anfragelatenz. Prefill-Rauschen verdeckt Decode-Regressionen.
• Spitzen-Gerätespeicher bei fester Batch-Größe, der bestimmt, ob Sie die Batch vergrößern können.
• Maximal nachhaltige Nebenläufigkeit, gemessen, indem die Batch-Größe erhöht wird, bis Sie auf einen Out-of-Memory-Fehler oder ein Latenzziel stoßen.
• Akzeptanzrate, falls spekulative Dekodierung im Spiel ist. Sie steuert direkt, wie oft Replay ausgelöst wird.
• Rollback-Kosten, gemessen als Latenz, die ein abgelehnter Draft hinzufügt. Das ist die Zahl, die replay-basierte Designs zu reduzieren versuchen.

Führen Sie jede Konfiguration mindestens dreimal aus, wobei ein Warmup-Durchlauf verworfen wird. Messungen der hybriden Dekodierung sind verrauscht, und Einzellauf-Vergleiche führen regelmäßig zu irreführenden Schlussfolgerungen.

Grenzen und offene Fragen

Ein paar ehrliche Einschränkungen. Der Nutzen jeder State-Speicherungsoptimierung hängt vom Verhältnis von Attention-Schichten zu rekurrenten Schichten in Ihrem Modell ab; ein leicht hybridisiertes Modell dürfte kaum Veränderung sehen, während ein stark rekurrentes viel sehen könnte. Replay führt Neuberechnung ein, zahlt sich also nur aus, wenn Speicher oder Rollback-Latenz die bindende Einschränkung ist statt rohe FLOPs. Und die Trade-off-Kurve ist workload-spezifisch: Das optimale Checkpoint-Intervall für Best-of-n-Sampling ist unwahrscheinlich optimal für Beam Search.

Wo die veröffentlichte Methode auf dieser Kurve landet und welche Workloads sie adressiert, sind Fragen, die der Originalbeitrag direkt behandelt. Betrachten Sie die Rahmung hier als Gerüst zum Lesen desselben, nicht als Ersatz.

Fazit

Hybride Modelle sind nicht schwieriger zu betreiben als Transformer; sie sind schwieriger zu planen. Der Grund ist, dass sie zwei inkompatible Arten von Speicher mit sich führen: einen adressierbaren KV-Cache und einen nicht adressierbaren rekurrenten Zustand. Alles, was modernes Serving effizient macht – Paging, Prefix-Sharing, spekulatives Rollback –, wurde um die erste Art herum entworfen.

Der Beitrag von ReplaySSM, wie er durch seinen Titel gerahmt wird, besteht darin, die zweite Art nicht länger als unvermeidbare Kosten zu behandeln und sie stattdessen als Designvariable zu betrachten. Wenn State billig genug rekonstruiert werden kann, wird Speicherung zu einem Regler statt zu einer Wand, und die verzweigten Workloads, die hybride Modelle derzeit schlecht bewältigen, werden wieder gewöhnlich.

Der praktische Weg ist unkompliziert: Bauen Sie eine saubere Umgebung auf, messen Sie Decode-Durchsatz und Spitzenspeicher, bevor Sie etwas ändern, isolieren Sie die Rollback-Kosten in einer spekulativen Dekodierungsschleife, und entscheiden Sie erst dann, wo auf der Speicher-gegen-Neuberechnung-Kurve Ihr Workload sitzen möchte.