ReplaySSM beschleunigt das Decoding von Hybridmodellen durch Caching der Eingaben statt des Zustands

NVIDIA und das Dao AI Lab an der Princeton University haben ReplaySSM vorgestellt, eine neue Methode zur Optimierung der Decoding-Phase von hybriden großen Sprachmodellen. Die im Oktober 2026 veröffentlichte Studie überdenkt, wie State Space Models (SSMs) ihren internen Zustand während der Inferenz speichern und aktualisieren. Durch das Caching kürzlicher Eingabetokens anstelle einer kontinuierlichen Überschreibung der vollständigen Zustandsmatrix reduziert die Technik den Bandbreitenverbrauch des Speichers erheblich und ermöglicht ein effizienteres spekulatives Decoding.

Was passiert ist

Hybride Modelle, die SSM-Schichten mit traditionellen Attention-Schichten abwechselnd verwenden, werden in modernen Architekturen wie NVIDIA Nemotron 3, Qwen3.5 und Kimi-K3 immer üblicher. Diese Modelle nutzen SSMs, um lange Kontexte effizient zu verarbeiten, und ersetzen den wachsenden Key-Value-Cache von Attention-Mechanismen durch einen fest definierten Zustand. Bei hohen Serving-Batch-Größen werden die SSM-Schichten jedoch häufig zum primären Latenzengpass, da sie häufige Lese- und Schreibzugriffe auf große Zustandsmatrizen aus dem High-Bandwidth-Speicher erfordern.

Das Forschungsteam stellte fest, dass konventionelles Decoding die Ausgabe-Berechnung und die Zustandsaktualisierung in einem einzigen Schritt fusioniert, was zu übermäßigem Speicherdatenverkehr führt. ReplaySSM entkoppelt diese Operationen, indem es den Zustand als Checkpoint plus einen Puffer für kürzliche Eingaben darstellt. Diese Änderung ermöglicht es dem System, Ausgaben zu berechnen, ohne ständig den vollständigen Zustand zurück in den Speicher schreiben zu müssen; der Checkpoint wird nur aktualisiert, wenn der Eingabepuffer voll ist. Die Methode ist nun in wichtige Inferenz-Frameworks wie TensorRT-LLM, vLLM und FlashInfer integriert.

Wie es funktioniert

Beim standardmäßigen SSM-Decoding lädt jeder Token-Generierungsschritt den vorherigen Zustand, wendet eine Abschwächung und Aktualisierung an, berechnet die Ausgabe und speichert den neuen Zustand. Dieser Prozess ist speicherlimitiert, da die arithmetische Intensität im Verhältnis zum Datentransfer gering ist. ReplaySSM ändert dies, indem es die Rekursion über ein kleines Fenster von Tokens entrollt. Anstatt den Zustand bei jedem Schritt zu aktualisieren, fügt es die Eingangsvektoren des aktuellen Tokens einem Puffer hinzu. Die Ausgabe wird unter Verwendung des gespeicherten Checkpoints und der gepufferten Eingaben berechnet.

Wenn der Puffer eine vordefinierte Größe erreicht, faltet eine Flush-Operation alle gepufferten Eingaben in den Checkpoint zusammen und schreibt ihn einmalig in den Speicher. Dies reduziert die Anzahl der Zustands-Schreibvorgänge um etwa die Hälfte. Für das spekulative Decoding, bei dem mehrere Entwurf-Tokens generiert und verifiziert werden, vereinfacht ReplaySSM den Rollback. Anstatt Snapshots des gesamten Zustands für jedes Entwurf-Token zu speichern, fügt das System die Entwurfs-Eingaben einfach dem Puffer hinzu. Wird ein Entwurf verworfen, setzt sich der Pufferzeiger zurück und verwirft die nicht übernommenen Eingaben ohne kostspielige Wiederherstellung des Zustands.

Wichtige Details

• ReplaySSM reduziert den Speicherdatenverkehr für den Zustand um approximately 50 %, indem Checkpoints erst geschrieben werden, nachdem ein Puffer für kürzliche Eingaben gefüllt ist.
• Auf NVIDIA Blackwell Ultra B300 GPUs bei einer Batch-Größe von 256 laufen SSM-Decoding-Kernels 1,46x bis 1,80x schneller als konventionelle Methoden.
• Der Durchsatz beim spekulativen Decoding steigt um 1,85x bis 2,10x für Modelle der 120B-Klasse, da der Rollback kein Speichern und Wiederherstellen von Zustands-Snapshots mehr erfordert.
• Die Technik unterstützt Low-Precision-Zustandsspeicherung, indem ein High-Precision-Checkpoint für die Rekursion beibehalten wird, während Low-Precision-Kopien für Lesevorgänge verwendet werden, was die Fehlerakkumulation begrenzt.
• Ausgewertete Modelle umfassen Varianten von Nemotron-3 und Qwen3.5, die konsistente Gewinne über verschiedene SSM-Architekturen wie Mamba-2 und Gated DeltaNet hinweg zeigen.
• Die Methode ist bereits in Produktionssystemen wie GLM-5.3-Flash und Kimi-K3 eingesetzt.

Warum es wichtig ist

Für Ingenieure, die große hybride Modelle bauen und bereitstellen, ist die Speicherbandbreite oft der limitierende Faktor für den Durchsatz. Indem ReplaySSM den dominierenden Speicherdatenverkehr in SSM-Schichten reduziert, ermöglicht es höhere Batch-Größen ohne Erhöhung der Latenz pro Token. Dies ist entscheidend für kosteneffektives Serving, da es eine bessere Hardwareauslastung auf teuren GPUs wie der H100 und B300 ermöglicht. Der Leistungsgewinn skaliert mit der Batch-Größe, was ihn besonders wertvoll für Produktionsumgebungen mit hoher Nebenläufigkeit macht.

Darüber hinaus macht ReplaySSM das spekulative Decoding für SSM-basierte Modelle praktikabel. Zuvor negierten die Kosten für das Speichern von Zustands-Snapshots für den Rollback die Vorteile der Spekulation in diesen Architekturen. Indem der Eingabepuffer wie ein Key-Value-Cache-Zeiger behandelt wird, stellt die Methode die Effizienz der Multi-Token-Vorhersage wieder her. Dies ermöglicht Entwicklern, aggressive Spekulationsstrategien zu nutzen, um die Latenz weiter zu reduzieren und die SSM-Leistung näher an die theoretischen Grenzen der Hardware zu bringen.

Was Sie tun können

• Lesen Sie das vollständige ReplaySSM-Paper auf arXiv, um die mathematische Herleitung und detaillierte Evaluationsmetriken zu verstehen.
• Erkunden Sie die Open-Source-Implementierung auf GitHub, um den Caching-Mechanismus in Ihren eigenen Projekten zu experimentieren.
• Führen Sie ein Upgrade auf die neuesten Versionen von TensorRT-LLM oder vLLM durch, um ReplaySSM-Unterstützung für kompatible hybride Modelle zu aktivieren.
• Benchmarken Sie Ihre aktuellen SSM-Workloads bei höheren Batch-Größen, um potenzielle Durchsatzverbesserungen durch reduzierten Speicherdatenverkehr zu quantifizieren.
• Untersuchen Sie Konfigurationen für spekulatives Decoding mit größeren Draft-Längen, da ReplaySSM die früheren Speicherstrafen im Zusammenhang mit dem Rollback eliminiert.
• Überwachen Sie Quantisierungsfehler, wenn Sie Low-Precision-Zustände verwenden, wobei die Trennung von Checkpoint und Puffer hilft, die Genauigkeit nahe am FP32-Niveau zu halten.