Abstract

Moderne KI-Arbeitslasten stützen sich stark auf optimierte Rechenkerne sowohl für das Training als auch für die Inferenz. Diese KI-Kerne folgen klar definierten Datenflussmustern, wie beispielsweise dem Verschieben von Tiles zwischen DRAM und SRAM sowie der Ausführung einer Folge von Berechnungen auf diesen Tiles. Dennoch bleibt die Entwicklung hochleistungsfähiger Kerne trotz dieser Musterkomplexität herausfordernd. Um die Spitzenleistung zu erreichen, sind sorgfältige, hardwarezentrierte Optimierungen notwendig, um moderne Beschleuniger voll auszuschöpfen. Obwohl domänenspezifische Compilerausgaben versuchen, die Belastung bei der Entwicklung hochleistungsfähiger Kerne zu verringern, leiden sie häufig unter Usability- und Ausdruckskraftlücken. In diesem Paper stellen wir TileLang vor – ein verallgemeinertes, geteiltes Programmiermodell zur effizienteren Entwicklung von KI-Kernen. TileLang trennt den Scheduling-Raum (Thread-Bindung, Layout, Tensorisierung und Pipeline) vom Datenfluss und kapselt diese als Satz von Anpassungsannotierungen und Primitiven. Dieser Ansatz ermöglicht es den Nutzern, sich auf den eigentlichen Datenfluss des Kernels zu konzentrieren, während die meisten anderen Optimierungen den Compilern überlassen werden. Wir führen umfassende Experimente auf gängigen Geräten durch und zeigen in zahlreichen Tests, dass TileLang bei zentralen Kernen Spitzenleistung erreicht. Dies belegt, dass das einheitliche Block-und-Thread-Paradigma sowie die transparenten Scheduling-Fähigkeiten sowohl die Leistungsfähigkeit als auch die Flexibilität bieten, die moderne KI-Systementwicklung erfordert.