NVIDIA hat den Universal Sparse Tensor (UST) in nvmath-python v0.9.0 integriert, um Sparse-Anwendungen in wissenschaftlichem Rechnen und Deep Learning zu beschleunigen. Der UST trennt die Sparsity eines Tensors von seinem Speicherlayout und soll so mehr Flexibilität bei Formaten und Ausführungspfaden ermöglichen.
Die Implementierung umfasst laut NVIDIA eine weitgehend kostenfreie Interoperabilität mit PyTorch, SciPy, CuPy und NumPy, benutzerdefinierte Sparse-Formate, polymorphe Operationen wie matmul() und solve(), die Nutzung in bestehenden PyTorch-Modellen sowie transparentes Caching, um erneute JIT/LTO-Kompilierung und erneute Planung bei wiederholter Ausführung derselben Operation zu vermeiden.
Grundlage ist eine einfache Domain-Specific Language (DSL), mit der sich gängige und weniger verbreitete Sparse-Speicherformate wie COO, CSR, CSC und BSR beschreiben lassen. Im Beitrag wird CSC als Zuordnung von Dimensionen und Level-Formaten gezeigt. NVIDIA stellt dem eine C++-Variante in MatX gegenüber, in der die DSL rein über Typen und Templates umgesetzt ist und formatabhängige Verzweigungen bereits zur Compile-Zeit aufgelöst werden können. In nvmath-python wird stattdessen eine stärker Python-orientierte Darstellung verwendet, bei der Formate als Objekte zur Laufzeit konstruiert werden können, auch durch Parsen von Strings. Formatabhängige Entscheidungen erfordern dadurch Laufzeitinspektion; NVIDIA beschreibt dies als Tausch von etwas Leistung bei der Formatinspektion gegen mehr Allgemeinheit und Laufzeitflexibilität.
Die Konvertierung zwischen UST und Tensoren oder Sparse-Matrizen aus PyTorch, SciPy, CuPy und NumPy erfolgt für dichte Tensoren sowie Formate wie COO, CSR, CSC, BSR, BSC und DIA nach Angaben des Beitrags meist ohne Datenbewegung oder Kopien. Stattdessen referenziert das UST-Objekt die Speicherpuffer der ursprünglichen Datenstruktur. Gezeigt wird dies unter anderem an einer Konvertierung einer SciPy-COO-Matrix in ein UST-Objekt und zurück in das Ursprungsformat.
Eigene Sparsity-Schemata lassen sich ebenfalls über die DSL definieren. Als Beispiel zeigt NVIDIA die Umwandlung eines dichten PyTorch-Tensors in ein 2-Bit-delta-komprimiertes Format, das PyTorch nativ nicht unterstützt. Dabei werden Koordinaten nicht direkt gespeichert, sondern als Abstand zum jeweils nächsten gespeicherten Element. Im gezeigten Beispiel ist Padding erforderlich, wenn der Abstand den mit 2 Bit darstellbaren Maximalwert von 3 überschreiten würde.
Für Debugging und Analyse stellt nvmath-python mehrere Hilfsfunktionen bereit. Die Methode __repr__ gibt eine eindeutige Textdarstellung des UST mit Datentypen, Gerät, Dimensionen, gespeicherten Arrays, Speicherbedarf und Sparsity aus. Zusätzlich erzeugen draw_storage() und draw() Visualisierungen des Speichers beziehungsweise der Tensorinhalte, während draw_raw() für größere 2D- oder 3D-Tensoren nur die Nonzero-Struktur darstellt.
Bei Operationen setzt die UST-Implementierung auf sparsity-agnostische Aufrufe wie matmul() und solve(). Das System inspiziert die Formate der Operanden und entscheidet dann zwischen einem optimierten, handgeschriebenen Kernel oder Bibliothekspfad für verbreitete Formate und automatisch erzeugtem Sparse-Code, wenn keine optimierte Lösung vorhanden ist. Dadurch sollen neue Sparse-Formate ohne manuelle Implementierung eines vollständigen Operationssatzes nutzbar werden, während für Standardformate vorhandene Bibliothekslösungen verwendet werden.
Die Ausführung ist in eine Planungsphase und eine Ausführungsphase getrennt. In der Planungsphase wird unter anderem zwischen Dispatch und JIT-LTO-Codegenerierung entschieden und die Konfiguration für spätere Wiederverwendung zwischengespeichert. Die eigentliche Berechnung erfolgt in der Ausführungsphase. NVIDIA verweist darauf, dass sich die anfänglichen Planungskosten bei wiederholter Ausführung gleicher matmul()-Aufrufe amortisieren, etwa bei iterativen Sparse-Solvern oder geprunten linearen Schichten im Deep Learning.
Für PyTorch bietet nvmath-python eine Möglichkeit, UST in bestehende Modelle einzubringen, ohne den Modellcode vollständig umzuschreiben. Dazu wird der UST in eine Unterklasse von torch.Tensor verpackt und mit Wrappern für polymorphe UST-Operationen versehen. Genannt werden unter anderem Wrapper für dot, mv, mm und linear. Zusätzlich zu dem bereits vorhandenen JIT/LTO-Caching gibt es dabei eine weitere Caching-Schicht für die Wiederverwendung geplanter Matmul-Instanzen bei ähnlichen Operationen.
Außerdem ist eine Methode reformat_model() vorhanden, die über die Gewichte aller linearen Schichten iteriert und eine benutzerdefinierte Funktion aufruft. Diese Funktion kann Gewichte inspizieren, gegebenenfalls prunen und sparsifizieren oder unverändert lassen. Wird None zurückgegeben, bleibt die jeweilige Gewichtsmatrix unverändert; andernfalls kann etwa ein UST auf Basis von to_sparse_csr() zurückgegeben werden. Bei der Inferenz werden dann die geänderten linearen Schichten mit UST verwendet.
Im ersten Performance-Vergleich untersucht NVIDIA die SpMV-Leistung auf der 1.489.752 × 1.489.752 großen Matrix atmosmodl aus Matrix Market. Verglichen werden CuPy über den @-Operator für COO, CSR und DIA, PyTorch mv() für COO und CSR sowie UST-matmul(). Für UST wird nur die Laufzeit von execute() gemessen, was laut Beitrag für wiederholte Multiplikationen ein fairer Vergleich sei, da CuPy und PyTorch keine entsprechende Planungsphase bereitstellen. NVIDIA berichtet in diesem Szenario von Beschleunigungen zwischen 1,1 und 444. Für COO wird eine schwache Leistung der CuPy- und PyTorch-Implementierungen hervorgehoben. Bei CSR nutzen alle Varianten cuSPARSE, wobei UST von der wiederverwendeten Planungsphase profitiert. Beim DIA-Format führt NVIDIA den Vorteil darauf zurück, dass UST einen spezialisierten Kernel einsetzt, während CuPy zunächst nach CSR konvertiert und PyTorch DIA nicht unterstützt.
Das zweite Experiment knüpft an die Veröffentlichung MACKO: Sparse Matrix-Vector Multiplication for Low Sparsity an. Dort wird ein effizientes SpMV-Verfahren für geringe Sparsity beschrieben, das für Single-Token-Inferenz im Deep Learning relevant ist. Das MACKO-Format entspricht laut NVIDIA im Wesentlichen dem zuvor gezeigten delta-komprimierten Format. Die Autoren berichten über eine CUDA-Implementierung, die gegenüber dichten Implementierungen bereits ab 50 Prozent unstrukturierter Sparsity Beschleunigungen erzielt.
NVIDIA hat diese Implementierung als Backend-Kernel für das delta-komprimierte Format in den UST-Lookup integriert und die Laufzeit mit einer dichten GEMV-Implementierung, cuSPARSE-SPMV, MACKO UST und der ursprünglichen MACKO-Implementierung auf 8192 × 8192 großen, gleichverteilten zufälligen Sparse-Matrizen verglichen. Die Sparsity variiert dabei von 0 Prozent bis 100 Prozent. Laut Beitrag arbeitet MACKO UST leicht schneller, weil Padding am Ende jeder Zeile stoppt; diese Optimierung könne aber auch die ursprüngliche MACKO-Implementierung übernehmen.
Für weitergehende Details verweist NVIDIA auf die Online-Dokumentation von nvmath-python.
