Triton einen Kernel nach dem anderen lernen: Softmax

Haftungsausschluss: Sofern nicht anders angegeben, stammen alle Illustrationen und Animationen vom Autor.

Definition

Der Softmax ist wie folgt definiert:

Die Normalisierung stellt sicher, dass sich der Vektor summiert 1sodass sie als gültige Wahrscheinlichkeitsverteilung interpretiert werden kann.

Beachten Sie, dass diese Formulierung des Softmax sehr empfindlich reagiert numerischer Überlauf. Denken Sie daran, dass der Maximalwert ein Commonplace ist float16 darstellen kann ist 65 504was ungefähr ist exp(11). Dies bedeutet, dass jeder Eingabewert größer als ~11 zu führt exp(z_i) Überschreiten des darstellbaren Bereichs, was dazu führt Überlauf.

Ein gängiger Trick, um dieses Drawback zu entschärfen, besteht darin, den Maximalwert des Eingabevektors von jedem Factor zu subtrahieren, sodass sich das neue Most ergibt 0 vor Potenzierung und 1 nach.

Naive Implementierung

Wie Sie sehen, ist die Berechnung des Softmax erforderlich zwei ReduktionsoperationenA max und a Summe. Ein naiver Algorithmus erfordert drei separate Durchgänge über den Eingabevektor. Zuerst wird das Most berechnet, dann die Summe und schließlich die normalisierten Ausgaben.

So sieht eine naive Numpy-Implementierung aus:

Ein wiederkehrendes Thema in dieser Triton-Serie ist die Minimierung hoher Latenzzeiten globaler Speicherzugriff. Unsere aktuelle Numpy-Implementierung erfordert drei separate Speicherlesevorgänge des vollständigen Eingabevektors, was äußerst ineffizient ist.

On-line-Softmax

Glücklicherweise können wir einen cleveren Trick anwenden, der sogenannte On-line-Softmaxum die zu verschmelzen max Und sum Schritte, wodurch die Anzahl der Speicherlesevorgänge reduziert wird 2.

Zunächst definieren wir die Summe der Exponentialfunktionen rekursiv. In der folgenden Reihe von Gleichungen gilt: m_i bezieht sich auf das Most über x bis zum ich-ter Index.

Diese Gleichheit ermöglicht es uns, die Summe der Exponentialfunktionen zu berechnen iterativ unter Verwendung des Maximalwerts bis jetzt. Wir können es nutzen, um die erste und zweite Schleife in der naiven Implementierung zu verschmelzen und das Most und die Summe der Exponentialfunktionen iterativ zu berechnen.

Unser Algorithmus wird:

Dies lässt sich leicht in Numpy übersetzen:

Nachdem wir nun die Hauptprinzipien hinter dem Softmax verstanden haben, werden wir es in Triton implementieren, angefangen bei der einfachen Einzelblockversion bis hin zur On-line-Mehrblockformulierung. Letztendlich möchten wir, dass sich unser Kernel wie ein PyTorch-Modul verhält und mit diesem kompatibel ist autograd.

Leider verhalten sich Triton-Kernel aus Sicht von PyTorch wie Black Packing containers: Die von ihnen ausgeführten Operationen werden nicht nachverfolgt autograd. Dies erfordert, dass wir den Rückwärtsdurchlauf selbst implementieren und explizit angeben, wie Gradienten berechnet werden sollen. Lassen Sie uns unsere geliebte Kettenregel auffrischen und den Softmax-Gradienten ableiten.

Gradient

Da die Ausgaben des Softmax streng positiv sind, können wir das verwenden logarithmische Ableitung um die Ableitung des Gradienten zu erleichtern. Hier nehmen wir die Ableitung von Protokoll der Ausgabe und wenden Sie die Kettenregel an:

Von dort aus ordnen wir die Begriffe neu und folgen diesen Schritten:

Nehmen wir nun an, dass wir einen Upstream-Gradienten haben, der beispielsweise durch eine Verlustfunktion erzeugt wird L (zB ein Kreuzentropieverlust). Wir erhalten den folgenden Ausdruck des Gradienten:

Die Vereinfachung des linken Begriffs in (9) liegt daran, dass δ_ij wird nur gleich sein 1 für die ich-tes Factor, wodurch die Summe zusammenfällt J auf einen einzigen Begriff.

Triton-Implementierung

Einzelblock Softmax

Nachdem wir nun die Ableitung des Gradienten durchgearbeitet haben, können wir die Vorwärts- und Rückwärts-Softmax-Kernel schreiben. Konzentrieren wir uns zunächst auf den PyTorch-Wrapper, um zu verstehen, wie die Einzelblockimplementierung auf hoher Ebene funktioniert. Bei einem 2D-Eingabetensor verarbeiten die Vorwärts- und Rückwärtskerne alle Zeilen parallel.

Der Einfachheit halber definieren wir die BLOCK_SIZE groß genug sein, um alle Spalten gleichzeitig zu verarbeiten. Konkret legen wir es als die nächste Potenz von 2 fest, die über der Anzahl der Spalten liegt, wie von Triton gefordert.

Dann definieren wir unser „Raster“ als die Anzahl der Zeilen (es könnte möglicherweise auch eine Batch-Dimension verarbeiten).

Der PyTorch-Wrapper für unsere SoftmaxSingleBlock ist eine Klasse, von der erbt torch.autograd.Perform das umsetzt ahead Und backward. Beide Methoden dauern eine ctx Argument, das wir verwenden werden, um die Softmax-Ausgaben während des Vorwärtsdurchlaufs zwischenzuspeichern und sie während des Rückwärtsdurchlaufs wiederzuverwenden.

Beide Kernel sind ziemlich unkompliziert. Wir beginnen mit dem Laden der Zeileneingaben mit derselben Syntax wie in meinem vorherigen Vektoraddition Artikel. Beachten Sie das BLOCK_SIZE Und num_warps werden mit a berechnet calculate_settings Funktion. Diese Funktion stammt aus dem Unfault Bibliothek und wurde in anderen Kernel-Bibliotheken wiederverwendet, z LigerKernel (auf dem die Kernel in diesem Artikel lose basieren) bietet es Heuristiken zum Optimieren beider Variablen:

def calculate_settings(n: int) -> tuple(int, int):
 MAX_FUSED_SIZE = 65536 # most grid dimension on Nvidia GPUs
    BLOCK_SIZE = next_power_of_2(n)
    if BLOCK_SIZE > MAX_FUSED_SIZE:
        # we take away this assertion on this article
        increase RuntimeError(
            f"Can not launch Triton kernel since n = {n} exceeds "
            f"the utmost CUDA blocksize = {MAX_FUSED_SIZE}."
        )
    num_warps = 4
    if BLOCK_SIZE >= 32768:
        num_warps = 32
    elif BLOCK_SIZE >= 8192:
        num_warps = 16
    elif BLOCK_SIZE >= 2048:
        num_warps = 8
    return BLOCK_SIZE, num_warps

Dann implementieren wir den regulären Softmax für den Vorwärtsdurchlauf und die Gleichung (10) für den Rückwärtspass. Die einzige Neuerung im Vergleich zu früheren Artikeln ist die Verwendung von Cache-Modifikatoren, die dem Compiler mitteilen, wie Daten zwischengespeichert und entfernt werden sollen. Im Second konzentrieren wir uns nur auf drei Cache-Modifikatoren:

• .ca (Cache auf allen Ebenen): Weist den Compiler an, die Daten sowohl in den L1- als auch in den L2-Cache zu laden, was darauf hindeutet, dass sie bald wiederverwendet werden könnten. Dieser Modifikator sollte verwendet werden, wenn die Daten klein genug sind, um in L1 zu passen (~128–192 KB professional SM auf einem A100) und wahrscheinlich wiederholt darauf zugegriffen wird.
• .cs (Streaming): Daten behandeln als Streamingwird es einmal verwendet und dann verworfen, um Platz in L1 freizugeben.
• .wb (Zurückschreiben): Normales zwischengespeichertes Schreiben, die Daten bleiben in der Cache-Hierarchie, intestine, wenn die Ausgabe wiederverwendet werden kann.

In den folgenden Kerneln verwenden wir die .ca Modifikator für Ladungen, da wir mehrere Operationen an den geladenen Daten durchführen. Zum Speichern verwenden wir .cs im Vorwärtsdurchlauf, da die Ausgaben nicht sofort wiederverwendet werden und .wb im Rückwärtsgang da im Kontext von autograd (dh die Kettenregel), Gradientenausgaben werden von nachgeschalteten Kerneln verbraucht.

Multiblock-Softmax

Werfen wir nun einen Blick auf die On-line-Formulierung des Softmax. In diesem Abschnitt implementieren wir eine Multiblock-Variante des vorherigen Kernels. Diese Model wird verwendet BLOCK_SIZE < n_colsmit anderen Worten, wir laden nur eine Kachel mit BLOCK_SIZE Elemente nacheinander, ähnlich wie wir mit gekacheltem GEMM in der umgegangen sind letztes Tutorial. Nun fragen Sie sich vielleicht: „Wie wählen wir die Blockgröße aus?“

Dies ist eine großartige Gelegenheit, Triton vorzustellen autotune Dienstprogramm. Ausgestattet mit einer Liste der Konfigurationen, autotune führt eine Rastersuche durch, um die beste Konfiguration für eine bestimmte Eingabeform zu ermitteln und zwischenzuspeichern. Dieser Vorgang wird jedes Mal wiederholt, wenn eine neue Eingabeform an den Kernel übergeben wird.

Hier führen wir mit der folgenden Hilfsfunktion eine Rastersuche über die Blockgröße und die Anzahl der Warps durch:

from itertools import product

# --- Multi Block Tuning ---
BLOCK_SIZES = (256, 512, 1024, 2048, 4096, 8192)
NUM_WARPS = (2, 4, 8, 16)

def get_autotune_config(
    block_sizes: listing(int), num_warps: listing(int)
) -> listing(triton.Config):
    return (
        triton.Config(kwargs={"BLOCK_SIZE": bs}, num_warps=nw)
        for (bs, nw) in listing(product(block_sizes, num_warps))
    )

Wir können jetzt unsere Multi-Block-Kernel damit dekorieren autotune und übergeben Sie die Liste der Konfigurationen, key=”n_cols” gibt an, dass die optimale Konfiguration von der Anzahl der Spalten der Eingabe abhängt.

Die Implementierung dieser Kernel ähnelt konzeptionell sehr dem zuvor behandelten On-line-Softmax. Der Hauptunterschied besteht darin, dass wir über Kacheln iterieren (nicht über einzelne Elemente wie in Numpy), was einige Anpassungen erfordert. Zum Beispiel addieren wir eine Summe über die Kachel im d Replace und der Rückwärtskernel erfordern nun ebenfalls zwei Iterationen.

Hinweis: Der PyTorch-Wrapper ist genau derselbe, außer dass wir die Zeile wo löschen BLOCK_SIZE Und num_warps werden deklariert (da sie ausgewählt werden von autotune).

Testen und Benchmarking

Wir können nun einen Vorwärts- und Rückwärtsdurchlauf mit beiden Kerneln ausführen und sicherstellen, dass sie mit den PyTorch-Baselines übereinstimmen:

def validate_kernel(kernel_fn: callable) -> None:
    system = "cuda:0" if torch.cuda.is_available() else "cpu"
    torch.random.manual_seed(0)

    # Generate inputs
    x = torch.randn((256, 512), system=system) # triton enter
    x.requires_grad = True
    xt = deepcopy(x) # torch enter

    triton_output = kernel_fn(x)
    torch_output = torch.softmax(xt, dim=1)
    torch.testing.assert_close(triton_output, torch_output) # take a look at fwd kernel

    # Setup pretend labels
    y = torch.zeros_like(x)
    inds = (torch.arange(0, y.form(0)), torch.randint(0, 3, (y.form(0),)))
    y(inds) = 1

    # Outline loss and run backward go
    loss_fn = torch.nn.CrossEntropyLoss()
    loss = loss_fn(torch_output, y)
    loss.backward()

    # Save gradient tensor for later
    torch_xgrad = xt.grad.detach().clone()
    triton_loss = loss_fn(triton_output, y)
    triton_loss.backward()
    torch.testing.assert_close(x.grad, torch_xgrad) # take a look at grad outputs

validate_kernel(softmax_sb)
validate_kernel(softmax_mb)

Abschließend vergleichen wir unsere Implementierung mit der PyTorch-Basislinie mithilfe des folgenden Snippets:

# --- Supply: Triton softmax tutorial ---
@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=("N"),  # argument names to make use of as an x-axis for the plot
        x_vals=(
            128 * i for i in vary(2, 100)
        ),  # totally different doable values for `x_name`
        line_arg="supplier",  # argument title whose worth corresponds to a unique line within the plot
        line_vals=(
            "triton_single_block",
            "triton_multi_block",
            "torch",
        ),  # doable values for `line_arg``
        line_names=(
            "Triton_single_block",
            "Triton_multi_block",
            "Torch",
        ),  # label title for the traces
        kinds=(("blue", "-"), ("inexperienced", "-"), ("crimson", "-")),
        ylabel="GB/s",  # label title for the y-axis
        plot_name="softmax-performance",  # title for the plot. Used additionally as a file title for saving the plot.
        args={"M": 4096},  # values for perform arguments not in `x_names` and `y_name`
    )
)
def benchmark(M, N, supplier):
    x = torch.randn(M, N, system=DEVICE, dtype=torch.float32)
    stream = getattr(torch, DEVICE.sort).Stream()
    getattr(torch, DEVICE.sort).set_stream(stream)
    if supplier == "torch":
        ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))
    if supplier == "triton_single_block":
        torch.cuda.synchronize()
        ms = triton.testing.do_bench(lambda: softmax_sb(x))
        torch.cuda.synchronize()
    if supplier == "triton_multi_block":
        torch.cuda.synchronize()
        ms = triton.testing.do_bench(lambda: softmax_mb(x))
        torch.cuda.synchronize()
    gbps = lambda ms: 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
    return gbps(ms)

benchmark.run(show_plots=True, print_data=True)

Gute Nachrichten! Unser Single-Block-Kernel übertrifft durchweg die PyTorch-Basislinie, während die Multi-Block-Variante bei Eingaben mit mehr als 6.000 Spalten abfällt:

Betrachtet man größere Eingaben, können wir mehrere Beobachtungen machen:

1. Der Multiblock-Kernel stabilisiert schließlich einen Durchsatz von etwa 900 GB/s und übertrifft damit die PyTorch-Basislinie für Eingaben mit mehr als 30.000 Spalten.
2. Interessanterweise scheint die Multiblock-Variante bei Eingaben mit mehr als 60.000 Spalten dominant zu sein.
3. Auch wenn wir mit der Single-Block-Variante die maximale Blockgröße überschreiten, läuft der Kernel aus irgendeinem Grund immer noch reibungslos. Tatsächlich verwaltet Triton die Blockgröße automatisch unter der Haube.
   Wann n_cols größer als das Hardwarelimit ist, zerlegt Triton die Eingabe und iteriert darüber. Dies scheint jedoch langsamer zu sein als der Multiblock-Ansatz.

Um noch weiter zu gehen, könnten wir beide Ansätze in einem einzigen Kernel kombinieren, der explizit den optimalen Kernel basierend auf der Eingabegröße auswählt. Auf diese Weise würden wir von der hohen Leistung des Single-Block-Kernels für kleine Eingaben und dem höheren Durchsatz der Multi-Block-Variante für Eingaben mit mehr als 60.000 Spalten profitieren.

Damit ist die dritte Folge dieser Triton-Serie abgeschlossen. Nochmals vielen Dank für Ihre Unterstützung!

Im nächsten Artikel werden wir die On-line-Softmax-Formulierung im Kontext von nutzen Blitzaufmerksamkeit.

Bis zum nächsten Mal! 👋

Ressourcen: