Jenseits des KV-Cache: Training und Serving von i14s Laplace 3B Linear-Attention-Modell auf AWS Trainium2 mit nativem PyTorch

Die meisten großen Sprachmodelle (LLMs) verwenden Transformer-Architekturen mit Softmax-Attention. Während der Inferenz wächst ihr Key-Value-(KV-)Cache mit dem Konversations- oder Dokumentenkontext. Dies kann die Parallelität reduzieren, die Serving-Kosten erhöhen und agentische Anwendungen zwingen, nützlichen Verlauf zu verwerfen.

i14 hat Laplace 3B als Linear-Attention-Alternative entworfen. Sein rekurrenter Zustand bleibt bei 29.5 MB pro Sequenz, unabhängig von der Kontextlänge. Gemeinsam mit AWS nutzte i14 native PyTorch-Unterstützung durch TorchNeuron Native um diese benutzerdefinierte Architektur auf AWS Trainium2zu bringen. Dies bewahrte sein Gated Linear Attention (GLA)-Design und den Muon-Optimizer.

Auf einer trn2.48xlarge-Instanz wurde Laplace 3B mit etwa 120,000 Tokens pro Sekunde trainiert. Das Modell lieferte 11,544 Ausgabe-Tokens pro Sekunde über 1,280 gleichzeitige Sequenzen. Dieser Beitrag erläutert das Modell design, die Änderungen, die i14 für Trainium2 vorgenommen hat, sowie die Trainings- und Inferenzergebnisse.

Über i14

i14 ist ein KI-Forschungs- und Technologieunternehmen, das Modellarchitekturen, Komponenten und Runtimes entwickelt. Seine Arbeit stützt sich auf KI, computergestützte Kognitionswissenschaft und Forschung dazu, wie biologische Systeme im Laufe der Zeit lernen. i14 wendet diese Ideen auf Modelle an, die langfristigen Kontext behalten können, ohne dass sich der Serving-Zustand mit jedem Token vergrößert.

Warum Linear-Attention-Modelle?

Transformer-Softmax-Attention vergleicht Tokens über den gesamten Kontext. Während des Trainings wächst diese Berechnung quadratisch mit der Sequenzlänge. Während der Inferenz wächst der KV-Cache, während das Modell mehr Tokens verarbeitet.

Linear Attention reorganisiert die Berechnung um einen rekurrenten Zustand, der für jedes Token aktualisiert wird. Laplace 3B verwendet GLA, um den vorherigen Kontext in einen Matrixzustand fester Größe zu komprimieren. Dies verändert die Trainings-Attention von O(T2) zu O(T) und hält den Inferenzzustand konstant bei 29.5 MB pro Sequenz.

Lineare Attention komprimiert den Verlauf, anstatt jeden Schlüssel und Wert zu behalten. Sie müssen daher die Abrufqualität für Ihre Workloads bewerten. Der feste Zustand kann langlebigen Agenten, der Analyse großer Dokumente und Anwendungen zugutekommen, die persistenten Speicher mit vorhersehbaren Serving-Anforderungen benötigen.

Wie i14 Laplace 3B für AWS Trainium2 angepasst hat

Die ursprüngliche Implementierung von Laplace 3B verwendete GPU-optimierte Triton-Kernels. Das AWS Generative AI Innovation Center (GenAIIC)-Team schrieb zunächst den GLA-Kernel mit standardmäßigen PyTorch-Operatoren für Trainium2 um. Nach dem Vergleich beider Ansätze entschied sich das Team für die Betaversion von TorchNeuron Native, dem nativen PyTorch-Backend für AWS Trainium, weil es eine bessere Leistung erbrachte.

Die Implementierung erforderte drei Änderungen:

1. Verwendung von nativem PyTorch für den GLA-Attention-Kernel. Das Team ersetzte den GPU-spezifischen fused_chunk_simple_gla-Triton-Kernel durch standardmäßige PyTorch-Operatoren und verwendete 16 Chunks mit jeweils 128 Token, um die 128 Partitionen des On-Chip-Puffers des NeuronCore abzubilden.
2. Neugestaltung der Kommunikation von Muon. Zwei all_to_all-Kollektive in MuonFSDP2 reduzierten den Datenverkehr pro Rank um das 24-Fache und steigerten den Durchsatz um das 1,63-Fache.
3. Kompilierung des sequenziellen Backward Pass. Das Team entrollte die 16 Chunk-Schritte der GLA-Rekurrenz statisch, um die Graphgrößen- und On-Chip-Speicherlimits des Compilers einzuhalten.

Der Abschnitt zu den technischen Herausforderungen beschreibt jede Änderung im Detail.

Warum AWS Trainium-Beschleuniger?

AWS Trainium ist eine Familie von Machine-Learning-Beschleunigern (ML-Beschleunigern) von AWS, die für Modelltraining und Inferenz entwickelt wurden. Eine trn2.48xlarge-Instanz kombiniert 16 Trainium2-Chips, 64 logische NeuronCores, Hochbandbreitenspeicher und NeuronLink-v3-Konnektivität. Diese Ressourcen unterstützen die in diesem Beitrag beschriebenen verteilten Trainings- und Inferenzkonfigurationen.

Drei Eigenschaften von Trainium2 prägten die Implementierung: softwareverwalteter On-Chip-Speicher (Herausforderung 1), eine schnelle Chip-zu-Chip-Verbindung (Herausforderung 2) und ein statischer Compiler (Herausforderung 3). Der Compiler zeichnet die GLA-Rekurrenz als festen Graphen auf.

Softwareverwalteter On-Chip-Speicher. Jeder NeuronCore verfügt über einen 28 MiB State Buffer (SBUF) mit 128 Partitionen und einen 2 MiB Partial-Sum-Puffer (PSUM). Der PSUM ist mit einer 128×128-systolischen Tensor Engine verdrahtet, daneben befinden sich separate Vektor- und Skalareinheiten. Der SBUF ist kein Cache: Der Compiler plant jede HBM-to-SBUF Übertragung im Voraus, sodass die Daten ohne Evictions oder Misses an Ort und Stelle sind (Abbildung 1).

Abbildung 1: Interne Architektur von NeuronCore-v3. Softwareverwalteter SBUF (28 MiB, 128 Partitionen), Tensor Engine (128×128 systolisches Array), Vector Engine, Scalar Engine und GPSIMD-Kerne.

i14 und die Laplace-3B-Architektur

Laplace 3B ist der 24-Schicht-i14-Decoder, bei dem jede Schicht Softmax-Attention durch Simple GLA ersetzt. Abbildung 2 verfolgt eine Schicht vom Eingabe- x_t oben bis zur Ausgabe- x_t + 1 unten. Dieser Abschnitt erläutert jedes Modul: Normalisierung und Faltung, die QKVG-Projektion, den GLA-Mixer und den Feed-Forward-Kanalmixer.

Normalisierung und Faltung. Der Residual-Stream-Vektor jedes Tokens x_t ∈ ℝ^D (D = 3, 072) durchläuft einen LayerNorm und eine depth-wise kausale Conv1d der Breite k = 4 mit SiLU. Dies verleiht jeder Position eine 4-Gram-Zusammenfassung vor dem Mixer. Der GLA-Mixer führt dann eine inhaltsbasierte Suche über die gesamte Historie durch seinen abklingenden Zustand durch.

QKVG-Projektion. Die Post-Conv-Ausgabe geht durch eine fusionierte Matrix W_qkvg ∈ ℝ^{9, 216 × 3, 072}, wobei 9, 216 = H × (2d_k + 2d_v). Er gliedert sich in vier Ströme pro Token: Queries q_t, Keys k_t, Values v_t, und ein Gate g_t (H = 16, d_k = 96, d_v = 192). Das Gate wird nach dem Mixer verwendet.

Abbildung 2: Architektur einer Laplace-Decoder-Schicht. Grün = fest (keine Parameter), grau = gelernt, gelb = rekurrenter Zustand, blau = Projektion.

GLA-Mixer. Jeder Kopf h hält einen Matrixzustand S_t^(h) der Größe d_k × d_v (96 × 192), einen komprimierten, abklingenden Datensatz vergangener Schlüssel-Wert-Paare. Bei jedem Token aktualisiert er sich wie folgt:

GLA-Rekurrenzzustands-Aktualisierungsgleichung

GLA-Zustands-Readout-Gleichung

GLA-Self-Augmentation-Gleichung

Abklingplan (Decay schedule). Die Abklingrate pro Kopf λ_h steuert, wie schnell ein Kopf vergangene Tokens vergisst. Sie wird bei der Initialisierung festgelegt (nicht gelernt) und durch einen geometrischen Zeitplan gesetzt. Der Zeitplan gibt den 16 Köpfen eine Streuung von Speicherlängen und definiert die logarithmische Abklingrate jedes Kopfes γ_h:

GLA decay schedule equations

Channel-Mixer (FFN). Nach dem Attention-Block und seinem Residualverbindung speist ein zweiter LayerNorm ein GLU-(SiLU)-FFN mit innerer Dimension 8,192. Seine drei Matrizen sind zwei Up-Projektionen W₁, W₂ ∈ ℝ^{8, 192 × D} und eine Down-Projektion W₃ ∈ ℝ^{D × 8, 192}. Es folgt eine zweite Residualverbindung.

Der Zustand S^(h) ist festgelegt auf d_k × d_v = 96 × 192 pro Kopf, unabhängig von der Sequenzlänge. Dadurch ist die Inferenz O(1) pro Token, und das Training ist O(T) statt O(T²).

Engineering-Herausforderungen

Jede der drei Änderungen stellte ihre eigene Engineering-Herausforderung dar.

Herausforderung 1: Nutzung von nativem PyTorch auf Trainium für GLA

Auf GPU läuft GLA als fused-Triton-Kernel fused_chunk_simple_gla aus der flash-linear-attention-Bibliothek. Das GenAIIC-Team formulierte ihn mit standardmäßigen PyTorch-Operatoren neu, die über TorchNeuron Nativelaufen, als explizite Schleifen über Chunks fester Größe. Die Chunk-Anzahl wird zur Kompilierzeit festgelegt und bestimmt die Tiefe der Backward-Unrolling. Das Team verwendet fla_num_chunks = 16, daher enthält jeder Chunk C = 128 Tokens.

Innerhalb jedes Chunks cist die Intra-Chunk-Ausgabe maskierte lineare Aufmerksamkeit, und der laufende Zustand S fügt früheren Kontext hinzu:

Chunked-GLA-Intra-Chunk-Ausgabe

Der Zustand geht dann zum nächsten Chunk über:

Chunked-GLA-Zustandsaktualisierung

Hier λ^C ist der pro Kopf verzögerte Zerfall, der über den Chunk aufsummiert wird. Die 16 Intra-Chunk-Berechnungen laufen parallel, und nur die Zustandsaktualisierung ist sequenziell (Abbildung 3).

Abbildung 3: GLA-Chunked-Rekurrenz mit paralleler Intra-Chunk-Attention und einer sequenziellen Zustandskette.

Da die SBUF 128 Partitionen hat, bildet jeder 128-Token-Chunk genau ein Token pro Partition ab, ohne Padding oder verschiebung über Partitionsgrenzen hinweg. Diese Ausrichtung erhöhte den gemessenen Durchsatz um 14 Prozent (Abbildung 4).

Abbildung 4: Die Chunk-Größe C = 128 bildet eins zu eins auf die 128 SBUF-Partitionen ab.

Herausforderung 2: Verteilte Muon-Orthogonalisierung unter FSDP2

Muon orthogonalisiert jeden Gradienten vor der Aktualisierung. Für jede der 96 zweidimensionalen Gewichtsmatrizen (2,718 Milliarden Parameter) führt er fünf Newton–Schulz-Iterationen durch (Algorithmus 1).

Der Newton-Schulz-Algorithmus

Das verteilte Problem. Die Gram-Matrix A = XX^⊤ korreliert jede Zeile des Gewichts mit jeder anderen Zeile. Unter Fully Sharded Data Parallel (FSDP2) hält jeder der 64 Ränge nur 1/64 der Zeilen. Eine shard-lokale XX^⊤ liefert daher nur einen diagonalen Block und verwirft jeden kreuz-Shard-Term (Abbildung 5).

Abbildung 5: Newton–Schulz unter FSDP2: der shard-lokale Gram-Block sowie Gather/Broadcast im Vergleich mit all_to_all.

Naiver Ansatz: Gather und Broadcast. Jeder Rang sammelt den vollständigen Gradienten (5,44 GB), ein Rang führt Newton–Schulz aus und sendet das Ergebnis als Broadcast. Dabei werden 10,87 GB pro Schritt bewegt, während die anderen 63 Ränge warten.

Produktionsansatz: all_to_all-Transposition. In MuonFSDP2 verschafft ein all_to_all jedem Rang einige vollständige Matrizen, alle 64 Ränge führen Newton–Schulz parallel aus, und ein zweites all_to_all gibt die Ergebnisse zurück. Der Datenverkehr pro Rang sinkt auf 0.45 GB, also 24-mal weniger. Der Durchsatz steigt von 73,378 auf 119,934 Token pro Sekunde, also um das 1.63-Fache, mit einer bitidentischen Verlustkurve (Abbildung 6).

Abbildung 6: Kommunikationsvolumen pro Schritt: 10,87 GB mit all_gather im Vergleich zu 0,45 GB mit all_to_all.

Kommunikationskosten. Newton–Schulz macht nur etwa 2 Prozent der FLOPs in einem Trainingsschritt aus (1.95 × 10¹⁴ von 9.53 × 10¹⁵), daher war die Kommunikation der Engpass. Über NeuronLink-v3 mit 1,28 TB/s dauert der 0,45-GB-Austausch weniger als 0,4 ms des ungefähr 4,4 Sekunden langen Schritts.

Challenge 3: Backward-Pass-Unrolling für eine sequenzielle Rekurrenz

Der Backward-Pass differenziert durch alle 16 sequenziellen Zustandsaktualisierungen. Auf der GPU baut Autograd diesen Graphen zur Laufzeit auf. Auf Trainium2 muss der Compiler den gesamten unaufgerollten 16-Schritt-Graphen im Voraus sehen, was drei Einschränkungen erzeugt.

Die Graphentiefe wächst mit der Chunk-Anzahl: bei 32 Chunks (C = 64) überlief der Graph den Scheduler. Der Speicher pro Chunk wächst mit der Chunk-Größe: bei 8 Chunks (C = 256) liefen Zwischenergebnisse in den HBM über. Jede Form wird zur Trace-Zeit festgelegt, daher benötigt jedes Schleifen-Zwischenergebnis seine eigene statische Allokation.

Sechzehn Chunks (C = 128) erfüllen alle drei: das Unrolling passt in den Scheduler, Zwischenergebnisse bleiben auf dem Chip, und die Tiles richten sich an den 128 SBUF-Partitionen aus.

Trainingsergebnisse

Das GenAIIC-Team trainierte Laplace 3B auf einer trn2.48xlarge-Instanz mit allen 64 logischen NeuronCores und dem all_to_all-Muon-Pfad. Das komplette Rezept umfasst 190,735 Schritte (100 Milliarden FineWeb-Edu-Tokens). Die folgenden Ergebnisse sind bis Schritt 84,210, etwa 44 Milliarden Tokens.

Durchsatz. Nach der Kompilierungs-Aufwärmphase stabilisiert sich der Durchsatz bei etwa 120,000 Tokens pro Sekunde und hält dort bis Schritt 84,210.

Loss- und Gradientenstabilität. Der Loss (Abbildung 7a) sinkt sanft von etwa 11.4 auf 2.3729 bei Schritt 84,210. Die Gradientennorm (Abbildung 7b) zeigt den erwarteten frühen Anstieg und dann einen stetigen Rückgang auf etwa 0.05, ohne Spitzen nach der Aufwärmphase, Divergenz oder NaN/Inf-Werte.

Abbildung 7: Trainingsdynamik. (a) Loss 11.4 → 2.3729. (b) Gradientennorm aus TensorBoard-Daten.

Inferenzoptimierung auf Trainium2

Über eine unveränderte vLLM-Engine bediente das GenAIIC-Team 11,544 Output-Tokens pro Sekunde bei 1,280 gleichzeitigen Sequenzen auf einer trn2.48xlarge-Instanz.

Konstanter Serving-Zustand. Der rekurrente Zustand von 29.5 MB pro Sequenz ist etwa 82-mal kleiner als ein vergleichbarer KV-Cache bei 8,192 Tokens. Eine Erhöhung des Prompts von 2,048 auf 8,192 Tokens fügte den Kosten pro Output-Token nur 0.5% hinzu.

Unverändertes vLLM. Ein Wrapper ersetzt pro Decoder-Layer ein Attribut, den Token-Mixer, um den rekurrenten Zustand im vLLM-Pool zu lesen und zu schreiben. Da sich das Modell als aufmerksamkeitsfrei deklariert, allokiert vLLM eine feste Page pro Sequenz. Die Outputs stimmten Token für Token mit der i14-Referenz überein (48/48 Schritte bei 3/3 Prompts in BF16).

Auffinden des Flaschenhalses. Eine lineare Anpassung der Decode-Zeit an die Anzahl der Decoder-Layer L ergibt:

Skalierung auf eine volle Instanz. Ohne Tensor-Parallelität skaliert der Durchsatz über Replikas, mit einer Engine und 20 Sequenzen pro logischem NeuronCore. Der Eager-Durchsatz wuchs von 133 Output-Tokens pro Sekunde auf einem Replik auf 6,772 bei 64, wobei die Host-CPU-Auslastung 82 Prozent erreichte. Das Kompilieren des Decode-Schritts entfernte die Operator-Dispatches pro Operator und erhöhte den Durchsatz um das 1.7-fache auf 11,544 (Abbildung 8).

Abbildung 8: Inferenz-Durchsatz der gesamten Instanz. (a) Replika-Skalierung bei eager und kompiliertem Decode. (b) Kompiliertes Decode hebt den Durchsatz von 6,772 auf 11,544 Ausgabe-Tokens pro Sekunde.

Was Fixed-Batch-Benchmarks verbergen. Unter echtem Datenverkehr ist jede neue Decode-Batch-Breite eine separat kompilierte Shape, deren erste Verwendung etwa 11 Sekunden kostet. Diese Schritte machten nur 1,1 Prozent der Gesamtzeit aus, verbrauchten aber 47,5 Prozent der Decode-Zeit. Padding auf feste Breiten beseitigte diesen Aufwand bei bit-exakten Ergebnissen. Kompiliertes Decode verbesserte den Durchsatz pro Replika anschließend um das 3,52-Fache bei gestaffelten Ankünften, verglichen mit dem 1,46-Fachen bei festen Batches. Da die BF16-Reassoziierung dazu führt, dass die kompilierte Ausgabe leicht vom eager-Ergebnis abweicht, verwenden Sie eager Decode, wenn exakte Reproduzierbarkeit wichtig ist.

Wichtigste Erkenntnisse

Das Projekt lieferte Erkenntnisse sowohl für Training als auch für Inferenz:

• Natives PyTorch kann benutzerdefinierte Architekturen auf Trainium2 unterstützen. i14 behielt die GLA-Architektur bei und nutzte TorchNeuron Native für die finale Implementierung.
• Kommunikationsdesign ist bei benutzerdefinierten Optimizern entscheidend. Das Muon-all_to_all-Design reduzierte den Datenverkehr pro Rang um das 24-Fache und steigerte den Trainingsdurchsatz um das 1,63-Fache.
• Fester Zustand hält den Inferenzspeicher vorhersagbar. Jede Sequenz verwendet 29.5 MB, unabhängig von der Kontextlänge.
• Profilen vor dem Kompilieren. Operator-Dispatch und nicht Arithmetik oder HBM-Bandbreite machten den Großteil der Decode-Zeit aus.
• Realistische Ankunftsmuster benchmarken. Feste Batches verbargen die Kosten für das Kompilieren neuer Shapes bei sich ändernden Batch-Breiten.

Fazit

Laplace 3B zeigt, dass AWS Trainium2 Modellarchitekturen über Standard-Transformer hinaus unterstützen kann. Unter Verwendung von nativem PyTorch über TorchNeuron Native und das AWS Neuron SDK trainierte und betrieb i14 sein GLA-basiertes Modell unter Beibehaltung eines festen rekurrenten Zustands.

Für agentische Systeme ändert dies eine grundlegende Serving-Einschränkung. Längere Verläufe erfordern keinen größeren Zustand pro Sequenz, was Teams helfen kann, mehr Kontext zu behalten und gleichzeitig die Speichernutzung vorhersagbar zu halten.

Wenn wachsender Kontext die Serving-Kosten erhöht, die Parallelität reduziert oder den Agentenspeicher einschränkt, testen Sie i14 oder wenden Sie sich an Ihr AWS-Kontoteam, um Laplace 3B auf AWS Trainium2 zu bewerten.

Ressourcen

• i14 ausprobieren
• AWS Trainium und Trainium2
• AWS Neuron SDK-Dokumentation
• Native PyTorch-Unterstützung auf AWS Trainium mit TorchNeuron Native
• Gated Linear Attention Transformers mit hardwareeffizientem Training
• Muon-Optimizer-Referenzimplementierung
• AWS Generative AI Innovation Center

Über die Autoren