Kimi Linear: Eine hardwarebewusste Architektur für effiziente KI-Inferenz mit langen Kontexten

Moonshot AI hat erneut eine bemerkenswerte Veröffentlichung vorgestellt. Nachdem Kimi-K2 und der dazugehörige Post-Training-Ansatz bereits einen starken Eindruck hinterlassen haben, präsentiert das Team nun zusätzlich Kimi Linear sowie Kimi-K2-Thinking, das ebenfalls einen Blick wert ist. Kimi Linear ist eine hybride Architektur für lineare Attention und führt mit Kimi Delta Attention, kurz KDA, einen neuen Attention-Mechanismus ein.

Die Veröffentlichung umfasst einen quelloffenen KDA-Kernel, der in Triton geschrieben wurde, vLLM-Implementierungen sowie vortrainierte und instruction-tuned Modell-Checkpoints. Das Modell verfügt über insgesamt 48 Milliarden Parameter, davon 3 Milliarden aktivierte Parameter, und unterstützt eine Kontextlänge von 1 Million Tokens.

In diesem Artikel werden die wichtigsten Erkenntnisse aus dem Kimi-Linear-Paper zusammengefasst. Außerdem wird erklärt, wie das Modell in einer allgemeinen GPU-basierten Cloud-Server-Umgebung betrieben werden kann.

Wichtige Erkenntnisse

• Kimi Delta Attention, kurz KDA, ist ein linearer Attention-Mechanismus mit feingranularem, kanalweisem Gating. Dadurch werden Speicherverwaltung und Hardwareeffizienz im Vergleich zu früheren Ansätzen wie Gated DeltaNet, auch GDN genannt, und Mamba2 verbessert.
• KDA nutzt eine spezialisierte Variante von Diagonal-Plus-Low-Rank-Übergangsmatrizen, kurz DPLR, um die Auslastung der Tensor Cores zu erhöhen.
• Kimi Linear setzt auf eine hybride Architektur, bei der drei KDA-Schichten mit einer MLA-Schicht kombiniert werden. Dieses Design reduziert die KV-Cache-Nutzung um 75 % und kann bei einer Kontextlänge von 1 Million Tokens einen bis zu 6-fach höheren Decoding-Durchsatz erreichen.

In früheren Erklärungen zu FlashAttention wurden bereits der Attention-Mechanismus und die Bedeutung hardwarebewusster Algorithmen behandelt. Die Entwicklung von Attention-Mechanismen, die stärker auf Hardwareeigenschaften abgestimmt und zugleich speichereffizienter sind, bleibt ein aktives Forschungsfeld. Daher sind in diesem Bereich weiterhin Fortschritte zu erwarten. Kimi Delta Attention ist eine Variante der linearen Attention, die zusätzlich Gating einsetzt. Um verständlich zu machen, warum die Entwickler von Kimi überhaupt einen neuen Attention-Ansatz eingeführt haben, geben die folgenden Abschnitte zunächst einen Überblick über lineare Attention und erklären anschließend, welche Rolle Gating-Mechanismen für Speichereffizienz und numerische Stabilität spielen.

Grundlagen der linearen Attention

Klassische Attention-Verfahren berechnen Attention-Scores, indem Softmax auf eine Ähnlichkeitsmatrix angewendet wird. Dieser Ansatz besitzt eine quadratische Zeit- und Speicherkomplexität in Bezug auf die Sequenzlänge, häufig als O(n²) beschrieben. Verschiedene Varianten der linearen Attention versuchen, diese quadratische Komplexität der Standard-Attention zu reduzieren. Lineare Attention ist dabei ein approximatives Attention-Verfahren, das mehr Effizienz ermöglicht, dafür jedoch gewisse Genauigkeitskompromisse in Kauf nimmt.

Bei linearer Attention wird die Softmax-Operation der klassischen Attention durch eine positive Feature Map ersetzt. Diese Feature Map ist so gestaltet, dass der daraus entstehende Kernel, also das Skalarprodukt, immer positiv bleibt. Dadurch wird der Effekt von Softmax nachgebildet, ohne dass ein expliziter Normalisierungsschritt erforderlich ist.

Lineare Attention hat jedoch Schwierigkeiten beim Abruf von Informationen über sehr lange Kontexte hinweg. Diese Einschränkung ist einer der Gründe, warum Kimi Linear eine hybride Architektur verwendet, die später näher erläutert wird.

Die Rolle von Gating

Informationen über lange Sequenzen hinweg effizient zu verwalten, ist eine anhaltende Herausforderung im Design von KI-Modellen. Gating-Mechanismen sollen die Speichereffizienz verbessern, indem sie einen selektiven Vergessensfaktor in den Attention-Prozess integrieren. Ein ähnliches Prinzip kommt auch in rekurrenten neuronalen Netzwerken wie LSTMs zum Einsatz.

Gating bei linearer Attention

Wenn Gating mit linearer Attention kombiniert wird, wird der quadratische und stetig wachsende Key-Value-Cache der klassischen Attention durch einen festen, matrixwertigen Zustand und lernbare Gates ersetzt. Während die Sequenz verarbeitet wird, kann das Modell dadurch gezielter entscheiden, welche Informationen gespeichert und welche verworfen werden sollen.

Gating wird häufig mit einer Delta-Update-Regel kombiniert, wie es bei KDA und Gated DeltaNet der Fall ist. Dadurch werden präzisere Änderungen am Speicherzustand möglich. Mit Delta-Update-Regel ist hier gemeint, dass die Differenz, also das Delta, zwischen neuen und vorhergesagten Werten berechnet wird, um den verborgenen Zustand zu aktualisieren, der als Speicherzustand dient.

Hardwarebewusste Algorithmen entwickeln

Hardwarebewusste Algorithmen zu entwickeln bedeutet, die Eigenschaften der Hardware zu verstehen und die zugrunde liegende Mathematik entsprechend neu zu betrachten. Moderne GPUs entfalten ihre Stärken besonders dann, wenn Arbeitslasten gut parallelisierbar sind und überwiegend aus Matrixmultiplikationen bestehen. Weniger effizient sind sie dagegen bei sequenziellen Abhängigkeiten und Operationen, die nicht auf Matrixmultiplikationen beruhen. Genau das stellt lineare Attention-Modelle wie KDA vor eine Herausforderung, da sie von Natur aus rekurrent sind.

Bei einer rekurrenten Methode wie KDA besteht das Ziel daher darin, die Berechnung so aufzubauen, dass sie möglichst gut in Blöcke aufgeteilt und dadurch einfacher parallelisiert werden kann. Gleichzeitig sollen unnötige Nicht-Matrixmultiplikationsoperationen entfernt werden, sofern sie das Gesamtergebnis nicht verändern. Die Reduzierung von Nicht-Matmul-FLOPs war auch ein wichtiger Gedanke hinter FlashAttention-2, da dadurch die Nutzung der Tensor Cores maximiert werden kann. Tensor Cores sind speziell darauf ausgelegt, den Durchsatz bei Matrixmultiplikationen zu beschleunigen.

Der folgende Abschnitt erklärt, wie das Kimi-Team diese Überlegungen bei der Entwicklung von Kimi Linear angewendet hat.

Wie bereits erwähnt, ist das KDA-Update rekurrent. Es handelt sich um einen sequenziellen, autoregressiven Prozess. Anders gesagt: Um den Zustand zum Zeitpunkt St zu berechnen, muss der Zustand zum Zeitpunkt St-1 bereits bekannt sein. Würde man dies auf einfache Weise direkt implementieren, müsste die GPU Tokens nacheinander verarbeiten. Dadurch bliebe ein großer Teil der verfügbaren Rechenleistung ungenutzt.

Um diese Ineffizienz zu reduzieren, wurde die Gleichung für Sr[t] in Blöcke aufgeteilt.

Dieser Schritt ist deshalb wichtig, weil er eine Berechnung, die sonst sequenziell ausgeführt werden müsste, mathematisch in eine Berechnung überführt, die parallel verarbeitet werden kann. Dadurch lassen sich mehrere Blöcke gleichzeitig berechnen.

Die Forschenden nutzten die WY-Darstellung, um eine Reihe von Rang-1-Updates in eine kompakte Form zu überführen. Dadurch werden teure Matrixinversionen während der Berechnung vermieden. Zusätzlich wurde eine UT-Transformation, also eine Upper-Triangular-Transformation, angewendet, um die Anzahl der Nicht-Matmul-FLOPs weiter zu reduzieren.

Die Architektur von Kimi Linear

Das Kimi-Team steht außerdem hinter Moonlight, das auf dem Erfolg von Muon aufbaut, einem Optimierer. Neben dem Nachweis, dass Muon für das Training großer Sprachmodelle skalieren kann, bildet Moonlight auch das Rückgrat der Kimi-Linear-Modellarchitektur.

Die Architektur von Kimi Linear kombiniert mehrere KDA-Schichten mit klassischen Full-Attention-Schichten in einem Verhältnis von 3:1.

Hybridisierung

Es stellt sich die Frage, warum Kimi Linear einen hybriden Ansatz nutzt. Anders formuliert: Warum kombiniert die Architektur globale Attention, also vollständige MLA, mit KDA? Der Grund liegt in der Schwäche linearer Attention beim Abruf von Informationen über lange Kontexte hinweg. Globale Attention ist rechenintensiver und benötigt mehr Speicher, weil alle Token-Paare verarbeitet werden. Das kann die Inferenz verlangsamen. Gleichzeitig ist globale Attention jedoch besser darin, den vollständigen Kontext und weitreichende Abhängigkeiten zu erfassen.

Positionskodierungen? NoPE

Der klassische Transformer-Attention-Mechanismus erkennt die Reihenfolge der Eingabeelemente nicht von selbst. Deshalb werden normalerweise explizite Positionskodierungen benötigt, um Informationen über die Sequenzreihenfolge einzubringen. RoPE, also Rotary Position Embeddings, gehört zu den verbreitetsten Formen solcher Positionskodierungen.

Bei Kimi Linear entschieden sich die Forschenden dafür, auf Positionskodierungen zu verzichten und stattdessen NoPE, also No Position Encoding, zu verwenden. NoPE ermöglicht es, diese Modelle bei der Inferenz in das rechnerisch effizientere reine Multi-Query-Attention-Format, kurz MQA, umzuwandeln. Außerdem vereinfacht NoPE das Training auf langen Kontexten, da keine Anpassung von RoPE-Parametern notwendig ist, etwa durch Änderung der Frequenzbasis oder durch Verfahren wie YaRN.

Der technische Bericht zu Kimi Linear verweist außerdem auf mehrere Arbeiten, die die Wirksamkeit des Verzichts auf Positionskodierungen mit NoPE gezeigt haben.

Implementierung

Zu Beginn wird ein GPU-basierter Cloud-Server vorbereitet. Für diese Einrichtung wird ein für Inferenz optimiertes Image ausgewählt. Um dieses Modell auszuführen, wird ein 4XH100-Cluster verwendet.

Die Verbindung zum Server erfolgt per SSH über die bevorzugte Entwicklungsumgebung und die öffentliche IPv4-Adresse. In diesem Beispiel wird Cursor verwendet.

ssh root@your_server_ip

weiter:

apt install python3.10-venv
# PyTorch mit CUDA-12.1-Unterstützung installieren, häufig genutzt für H100-Systeme
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Hugging Face Transformers, Accelerate und Kimi-Abhängigkeiten installieren
pip install transformers accelerate bitsandbytes sentencepiece protobuf tiktoken

# Den für Kimi benötigten Flash-Linear-Attention-Kern installieren
pip install vllm
pip install -U fla-core
vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct \

 --port 8000 \

 --tensor-parallel-size 8 \

 --max-model-len 1048576 \

 --trust-remote-code

Fazit

Kimi Linear ist eine hardwarebewusste Architektur. Die zentrale Innovation ist Kimi Delta Attention, kurz KDA, das feingranulares, kanalweises Gating nutzt und mathematische Anpassungen wie das Aufteilen des rekurrenten Updates in Blöcke einsetzt. Diese Verbesserungen unterstützen eine effizientere Speicherverwaltung und eine stärkere Auslastung der Tensor Cores. Die hybride Architektur mit einem Verhältnis von 3:1 zwischen KDA und vollständiger MLA ist darauf ausgelegt, die Schwäche linearer Attention beim Abruf langer Kontexte zu verringern, ohne die Effizienzvorteile aufzugeben. Der Einsatz von NoPE, also No Position Encoding, vereinfacht zudem das Training und erleichtert die Optimierung des Modells für effiziente Multi-Query-Attention-Inferenz. Das Ergebnis ist ein Modell, das die KV-Cache-Nutzung um 75 % reduziert und bei einer Kontextlänge von 1 Million Tokens einen bis zu 6-fach höheren Decoding-Durchsatz erreicht. Dadurch wird die Bereitstellung großer Sprachmodelle skalierbarer und kosteneffizienter.