compile=True und der TileLang-Fast-Path: Engineering-Notizen

Diese Notizen beschreiben, wie sich compile=True und fast=True über das hinaus verhalten, was das README sagt. Sie stammen aus Messungen, die während der Arbeit an #472, #576 und #718 auf einer RTX 4070 Ti SUPER mit torch 2.11 und tilelang 0.1.14 entstanden sind. Sie stehen hier, damit die nächste Person sie nicht erneut messen muss.

Backend-Auswahl

Agent(..., backend="auto") und laya.load(..., backend="auto") entscheiden sich für die Backend-Klassenschicht. Der Standard bleibt eager. Ein explizites backend= hat Vorrang vor compile und fast; lässt man es weg, bleibt das bestehende Verhalten beider Flags erhalten.

• eager: der Standard-PyTorch-Forward, auf jedem unterstützten Gerät.
• compile: CUDA-only torch.compile mit dynamischen Formen und reduce-overhead-Modus, Bucket-Padding, persistentem Inductor-Cache und Warmup bei der Installation. Es verwendet denselben Unabhängige-Dimensionen-Scope wie compile=True, das seinen bestehenden Standardmodus und CPU-Support behält. Setze LAYA_COMPILE_WARMUP=0, um das Backend-Warmup zu verschieben, und LAYA_INDUCTOR_CACHE_DIR, um sein Cache-Verzeichnis zu wählen (Standard ~/.cache/laya/inductor).
• tilelang: ein Adapter um den aktuellen Fast-Path, der den bf16- oder fp16-dtype des Agenten verwendet.
• auto: TileLang auf CUDA mit einem unterstützten ModernBERT-Encoder und dtype, wenn TileLang installiert ist, sonst compile auf CUDA; eager auf anderen Geräten.
• onnx: laya.load(..., backend="onnx", onnx_path="model.onnx") gibt den bestehenden ONNXAgent zurück. Ohne onnx_path verwendet es laya.onnx.

Ein nicht verfügbares Backend gibt eine RuntimeWarning aus, die das aufgelöste Backend nennt, und fällt auf eager zurück. Um ein Backend zu erzwingen, nutze agent.set_backend("tilelang", strict=True). Ein Wechsel wartet auf aktive Inferenz; agent.backend meldet den aktiven Namen und agent.backend_object legt das installierte Objekt offen. agent.set_backend("compile", warmup=False) verschiebt die Kompilierung bis zur Inferenz, sodass Kompilierfehler dann bei der Anfrage auftauchen. agent.warmup() bleibt verfügbar. agent.deaccelerate() entfernt ein über die Klassenschicht installiertes Backend.

Router leiten eine explizite Auswahl über Router(agent_kwargs={"backend": "auto"}) weiter. Sie übergeben standardmäßig kein Backend-Argument und bewahren so die Kompatibilität mit bestehenden Agent-ähnlichen Konstruktoren. Scoped CPU-OOM-Wiederholungen lösen das Backend ab und stellen es wieder her, wenn das Modell auf sein ursprüngliches Gerät zurückkehrt.

compile=True materialisiert die Attention-Maske

Eager-SDPA nimmt ModernBERTS (rows, 1, L, L)-Attention-Maske als Broadcast-View. Unter den dynamischen Formen, die compile=True verwendet, kann inductor nicht beweisen, dass die letzte Dimension ausgerichtet ist. Es expandiert die Maske auf jeden Head und padded sie in einen echten Puffer von rows x heads x L x L. In bf16 mit 12 Heads sind das 0,8 GB bei 32 Zeilen x 1024 Token.

• GPU mit Reserve. Der Puffer kostet Bandbreite, zig Millisekunden pro langem Batch.
• GPU fast voll. Der Caching-Allocator thrash, und derselbe Aufruf kann zig Sekunden dauern.

Wenn du mit langen Batches auf einer ausgelasteten GPU kompilierst, begrenze die Batch-Größe (predict_batch(..., batch_size=)) oder nutze fast=True. Die TileLang-Attention liest den gepackten QKV-Puffer und maskiert nach Sequenzlänge, hat also keinen solchen Puffer.

Kaltstart

• Erste Kompilierung. Sie dauert zig Sekunden pro Graph. compile=True braucht zwei Graphen: einen für Batches und einen für eine einzelne Zeile, die torch spezialisiert. compile=True ruft jetzt agent.warmup() während des Ladens auf. compile_warmup=False stellt die faule Kompilierung wieder her, und agent.warmup(shapes=...) bleibt manuell verfügbar. Eager- und TileLang-Ladevorgänge wärmen nicht automatisch. Diese Formen decken gängige Anfragen ab, nicht jede mögliche Shape Guard.
• Warm-up-Fehler. Das automatische Warm-up ist Best Effort: Ein Fehler gibt eine RuntimeWarning aus, die den Fehler nennt (einschließlich des zugrunde liegenden Compiler-Fehlers), und das Laden kehrt mit intaktem torch.compile-Wrapper und intakten Kompilierungseinstellungen zurück. Zum Beispiel kann Windows ohne MSVC mit compile=True laden, obwohl das Warm-up fehlschlägt. Spätere Anfragen verwenden trotzdem das kompilierte Modell und zeigen Kompilierungsfehler auf; Laya schaltet sie nicht auf eager Ausführung um. Explizite agent.warmup()-Aufrufe propagieren Fehler ebenfalls, auch nach einem fehlgeschlagenen automatischen Warm-up. Ein erfolgreiches Laden garantiert also nicht, dass die kompilierte Inferenz bereit ist.
• Laya-Cache-Opt-in. laya.load(..., compile=True, compile_cache=True) setzt die prozessweite TORCHINDUCTOR_CACHE_DIR nur, wenn sie fehlt, auf $XDG_CACHE_HOME/laya/torchinductor oder ~/.cache/laya/torchinductor, wenn XDG nicht gesetzt oder nicht absolut ist. Eine bestehende Einstellung, einschließlich einer von einem früheren PyTorch-Compile gesetzten, gewinnt. Das Verzeichnis wird beim Laden erstellt; Dateisystem- fehler propagieren. compile_cache=False (Standard), eager und TileLang-Ladevorgänge lassen die Umgebung unberührt. Dies verschiebt oder löscht keine alten Caches. Container brauchen weiterhin ein persistentes Home/Volume. Cache-Kompatibilität und -Invalidierung werden von PyTorch verwaltet; eine Änderung an GPU, torch, Compiler, Modell oder Input Guard kann eine erneute Kompilierung erfordern.
• Über Neustarts hinweg. Inductors FX-Graph-Cache hält kompilierte Graphen unter TORCHINDUCTOR_CACHE_DIR. Der Standard liegt unter /tmp, was einen Neustart oder einen Container-Neustart nicht überlebt. Setze ihn auf ein persistentes Verzeichnis, oder ein Volume in einem Container, und ein zweiter Prozess lädt die Graphen, statt sie zu kompilieren. In der Messung von #472 verkürzte das das Warm-up von etwa 120 s auf etwa 50 s.

Opt-in-CUDA-Graphen

agent = laya.load("convaiinnovations/laya", compile=True,
                  compile_cache=True, compile_mode="reduce-overhead")

compile_mode ist standardmäßig "default"; auf dem aktiven kompilierten Pfad sind nur "default" und "reduce-overhead" zulässig. Eager- und TileLang-Ladevorgänge ignorieren die Compile-Optionen. CPU-Kompilierung funktioniert weiterhin, aber CUDA-Graph-Aufzeichnung gilt nur auf CUDA. Der CUDA-Modus erfordert PyTorchs torch.compiler.cudagraph_mark_step_begin-API; ältere Builds ohne sie lösen einen expliziten Fehler aus.

Dynamische Dynamo-Graphen bedeuten keine shape-unabhängigen CUDA-Graphen: Neue konkrete Formen können erneut Warm-up und Aufzeichnung erfordern, ohne einen neuen Dynamo-Graphen. Die zwei synthetischen Standard-Warm-up-Formen zeichnen nicht jede Anfrageform vorab auf. Wiederholte Formen können profitieren, aber variierende Formen können zusätzliche Latenz zahlen und Graph-Pools behalten. PyTorch kann CUDA-Graphen für nicht unterstützte Operationen oder Konfigurationen überspringen; das Setzen dieses Modus ist keine Capture-Garantie.

Laya markiert jeden kompilierten CUDA-Forward als neuen Schritt, serialisiert diese Forwards über seine Agenten und klont beide Ausgabe-Tensoren außerhalb des kompilierten Graphen, bevor es die Sperre freigibt. Das hält behaltene Ausgaben über Replays hinweg gültig, auf Kosten von zwei Kopien und serialisierter Forward-Ausführung. Die Sperre koordiniert keine unabhängigen, anwendungseigenen kompilierten Modelle; Aufrufer, die CUDA-Graph-Iterationen teilen oder benutzerdefinierte Streams verwenden, müssen ihre eigene Koordination verwalten. Disk-Caches verwenden kompilierten Code wieder, nicht Live-CUDA-Graph-Aufzeichnungen oder deren Gerätespeicher, über Prozesse hinweg.

Reproduziere Kalt-/Neustart-Timings, Speicher und Cache-Zähler mit benchmarks/bench_compile_defaults.py; siehe die aufgezeichneten Messungen.

AOTInductor: noch nicht

Ein vorkompiliertes Artefakt pro Checkpoint und GPU-Architektur auszuliefern (torch._inductor.aoti_compile_and_package) würde die Kompilierung ganz entfernen. Auf torch 2.11 bleibt es beim Packaging hängen:

• Export funktioniert. torch.export von DecisionModel gelingt, in etwa 5 s, mit dynamischen Zeilen, Markern und Tokens. Tokens müssen als Vielfaches von 16 deklariert werden (16 * Dim(...)); ein einfacher Bereich scheitert am eigenen L % 8-Ausrichtungsguard des Exporters. Das ist dieselbe Maskenausrichtung wie oben.
• Packaging schlägt fehl. Wie es fehlschlägt, hängt davon ab, wie das Programm exportiert wurde:
  • Unter Autocast trägt das Programm dtype-Asserts, über die AOTI außerhalb von Autocast stolpert: Tensor dtype mismatch! Expected: torch.bfloat16, Got: torch.float32.
  • Aus einer bf16-Kopie ohne Autocast schlägt das Tracing im Forward fehl: mat1 and mat2 must have the same dtype. DecisionModel.forward castet den gepoolten Zustand und die Konfidenz-Merkmale vor dem Action-Head nach fp32 hoch, und Autocast gleicht das normalerweise aus.

Der Artefakt-Weg braucht daher einen dtype-expliziten Action-Head: entweder seinen Input auf den dtype des Heads casten, oder den Head in fp32 laufen lassen.

TileLang-Portabilität: die Kernel sind CUDA-only

tilelang registriert Targets für CUDA, HIP, Metal, WebGPU und ein C-Backend. Ohne AMD- oder Apple- Hardware war die beantwortbare Frage, ob laya/tl_kernels.py überhaupt für die CPU lowern kann. Geprüft mit tilelang.compile(kernel.prim_func, target=...) auf Linux x86-64:

Jeder Laya-Kernel scheitert auf "c", aus einem von drei Gründen:

Das C-Backend akzeptiert hingegen:

• fp32-elementweise Schleifen (T.Parallel);
• T.Pipelined, was zu einer einfachen Schleife lowert;
• T.gemm mit einem T.alloc_local-Akkumulator, was zu einer skalaren Dreifachschleife lowert.

Eine CPU-Version wäre daher ein zweiter Satz von Kernels, kein Target-Flag. Ihr GEMM wäre eine unblockierte skalare Schleife, und sie würde nicht mit dem MKL-/oneDNN-Pfad konkurrieren, den der stock Forward auf der CPU schon nutzt. Dieselben drei Konstrukte sind die, die man auf HIP und Metal zuerst prüfen sollte: Fragmente, T.gemm mit GemmWarpPolicy, und bf16/fp16-Unterstützung.

Um die erste Zeile der zweiten Tabelle zu reproduzieren: tilelang.compile(K.gemm_kernel(768, 768).prim_func, target="c").