Пишем ядра на Triton

Зачем это нужно

В уроке 5.7.1 torch.compile сам сливал поэлементные операции в ядра на Triton, в уроке 5.7.2 вы написали FlashAttention на NumPy. Но NumPy не умеет класть блоки в SRAM, а torch.compile не придумает online softmax за вас: он сливает то, что видит в графе, и не меняет алгоритм. Когда нужен новый алгоритм — своё внимание, слитый слой с квантованием, особый вид разреженности, — ядро пишут руками.

Раньше это значило CUDA C++: сотни строк про потоки, варпы, shared memory и синхронизацию. TritonTritonEN Triton (language and compiler)Язык и компилятор GPU-ядер на Python (OpenAI, открытый под MIT). Программист описывает работу одного блока данных, а распределение по потокам, shared memory и тензорные ядра берёт на себя компилятор. На нём написаны ядра torch.compile, flash-linear-attention, многие ядра vLLM. (Филипп Тилле, статья 2019 года; с 2021-го развивается в OpenAI, лицензия MIT) позволяет писать ядра на Python-подобном языке, думая блоками данных, а не отдельными потоками. На Triton написаны ядра, которые генерирует torch.compile, библиотека flash-linear-attention из урока 5.5.3, ядро NSA из урока 5.5.2, многие ядра vLLM. Triton поддерживает видеокарты NVIDIA с compute capability 8.0 и выше, то есть и вашу Ada (8.9). Ставить отдельно ничего не нужно: пакет triton приходит вместе с PyTorch для Linux.

Модель программирования

Ядро на Triton — это функция с декоратором @triton.jit, которая описывает работу одной программы. Программа обрабатывает целый блок данных: вектор из 1024 чисел или тайл $64 \times 64$. При запуске вы задаёте сетку — сколько программ запустить; каждая программа узнаёт свой номер через tl.program_id и по нему вычисляет, какой кусок данных её.

	CUDA C++	Triton
Пишем код для	одного потока	одной программы (блока данных)
Единица данных	скаляр	вектор или тайл (tl.arange, двумерные блоки)
Shared memory	выделяем и синхронизируем сами	распределяет компилятор
Тензорные ядра	особые инструкции (WMMA, MMA)	tl.dot(a, b)
Объединение доступов к памяти	следим сами	делает компилятор
Контроль	полный	меньше, но обычно хватает

Пять функций, на которых держится почти всё:

• tl.program_id(axis) — номер программы в сетке (аналог номера блока потоковблок потоковEN thread block (CTA)До 1024 потоков, которые выполняются на одном SM и делят его shared memory. Ядро запускается сеткой (grid) из многих блоков. из урока 5.7.1);
• tl.arange(0, BLOCK) — вектор индексов $0, 1, \dots, \text{BLOCK} - 1$; BLOCK должен быть степенью двойки и константой компиляции (tl.constexpr);
• tl.load(ptr + offs, mask=..., other=...) — прочитать блок из видеопамяти; ptr — адрес начала тензора, offs — смещения;
• tl.store(ptr + offs, value, mask=...) — записать блок;
• маска mask — булев вектор: где False, чтения и записи нет. Без неё последний неполный блок вылезет за конец массива.

Ядро 1: сложение векторов

Начнём с ядра, которое не делает ничего умного, зато показывает всю механику: сетку, номера программ, маски, запуск и проверку.

"""vector_add.py — первое ядро на Triton: сложение векторов.  python vector_add.py"""
import os

import torch
import triton
import triton.language as tl

# без GPU: TRITON_INTERPRET=1 python vector_add.py — ядро выполнит интерпретатор на CPU (медленно)
DEVICE = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"


@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)                      # номер программы (блока) в сетке: 0, 1, 2, …
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # индексы чисел, за которые отвечает этот блок
    mask = offs < n                                  # последний блок может вылезти за конец массива
    x = tl.load(x_ptr + offs, mask=mask)             # загрузка из видеопамяти (лишнее не читаем)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)       # запись обратно (лишнее не пишем)


def add(x, y):
    out = torch.empty_like(x)
    n = out.numel()
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)   # сколько блоков запустить
    add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)
    return out


torch.manual_seed(0)
x = torch.rand(98_432, device=DEVICE)                # 98 432 не делится на 1024 — маска пригодится
y = torch.rand(98_432, device=DEVICE)
torch.testing.assert_close(add(x, y), x + y)          # упадёт с подробным сообщением, если ответ неверен
print("add: совпадает с x + y")

if DEVICE == "cuda":
    for n in [2**20, 2**24, 2**27]:
        x = torch.rand(n, device=DEVICE)
        y = torch.rand(n, device=DEVICE)
        for name, fn in [("triton", lambda: add(x, y)), ("torch ", lambda: x + y)]:
            ms = triton.testing.do_bench(fn)          # среднее время в миллисекундах, с прогревом
            gbs = 3 * n * x.element_size() / (ms * 1e-3) / 1e9   # прочитали 2 вектора, записали 1
            print(f"n = 2^{n.bit_length() - 1}, {name}: {ms:.3f} мс, {gbs:6.0f} ГБ/с")

Разберём запуск. grid — функция, которая по параметрам запуска возвращает размер сетки: $\lceil 98\,432 / 1024 \rceil = 97$ программ. Программа 96 отвечает за индексы $98\,304 \dots 99\,327$, из которых существуют только 128 — маска отрежет остальные. add_kernel[grid](...) компилирует ядро при первом вызове (несколько секунд) и кладёт результат в кэш на диске. Тензоры PyTorch передаются в ядро как указатели на их первый элемент.

torch.testing.assert_close — главный инструмент проверки: сравнивает с допуском, подобранным под тип данных, и при ошибке печатает, сколько элементов разошлось и где самое большое расхождение.

Без видеокарты. Переменная TRITON_INTERPRET=1 включает интерпретатор: ядро выполняется на CPU поверх NumPy. Медленно, но для проверки корректности хватает, и внутри ядра работает обычный print. Так проверялись все ядра этого урока.

Чего ждать на RTX 4080 (оценка, а не замер). Сложение — memory-bound с интенсивностью $1/12$ во FP32 (1 FLOP на 12 байт). На больших векторах и Triton, и x + y должны выйти на 80–90% от паспортных 716.8 ГБ/с; на векторе в $2^{20}$ чисел (12 МБ трафика) — заметно меньше: ядро заканчивается за десятки микросекунд, и накладные расходы запуска становятся заметны. Обогнать PyTorch на сложении нельзя — оба упираются в одну и ту же память. Выигрыш начинается там, где можно слить операции.

Ядро 2: слитый softmax

Softmax по строкам — типичная операция-«клей» (во внимании, в выходном слое). Наивная версия на PyTorch — 5 ядер:

Шаг	Читает	Пишет
x.max(dim=1)	$x$	вектор максимумов
x - x_max	$x$	$z$
torch.exp(z)	$z$	num
num.sum(dim=1)	num	вектор сумм
num / sum	num	ответ

Пять чтений и три записи матрицы — 8 проходов (векторы длины $M$ не в счёт). Слитое ядро читает строку один раз, считает всё в регистрах и пишет один раз: 2 прохода, в 4 раза меньше трафика. Ядро ниже — одна программа на строку, вся строка — один блок; если строка короче блока, хвост читается как $-\infty$ и не попадает ни в максимум, ни в сумму ($e^{-\infty} = 0$).

К прямому проходу добавим обратный. Для $y = \mathrm{softmax}(x)$ и градиента $dy$ по выходу (якобиан softmax из урока 1.3.2):

$$dx = y \odot \big(dy - \langle dy, y \rangle\big),$$

где $\odot$ — поэлементное умножение, а $\langle dy, y \rangle = \sum_j dy_j\, y_j$ — одно число на строку. Это тоже редукция по строке, то есть то же ядро другой формулы. Чтобы autograd знал о наших ядрах, оборачиваем их в torch.autograd.Function:

"""softmax_triton.py — слитый softmax по строкам: прямой и обратный проход, autograd.Function, бенчмарк."""
import os

import torch
import triton
import triton.language as tl

DEVICE = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"


@triton.jit
def softmax_fwd_kernel(x_ptr, y_ptr, n_cols, stride, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)                           # одна программа — одна строка
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols                             # BLOCK_SIZE — степень двойки, строка может быть короче
    x = tl.load(x_ptr + row * stride + cols, mask=mask, other=-float("inf")).to(tl.float32)
    x = x - tl.max(x, axis=0)                        # устойчивость (урок 1.3.2)
    num = tl.exp(x)                                  # exp(-inf) = 0: хвост не попадает в сумму
    y = num / tl.sum(num, axis=0)
    tl.store(y_ptr + row * stride + cols, y.to(y_ptr.dtype.element_ty), mask=mask)


@triton.jit
def softmax_bwd_kernel(y_ptr, dy_ptr, dx_ptr, n_cols, stride, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    y = tl.load(y_ptr + row * stride + cols, mask=mask, other=0.0).to(tl.float32)
    dy = tl.load(dy_ptr + row * stride + cols, mask=mask, other=0.0).to(tl.float32)
    dx = y * (dy - tl.sum(dy * y, axis=0))           # градиент softmax: y ⊙ (dy − ⟨dy, y⟩)
    tl.store(dx_ptr + row * stride + cols, dx.to(dx_ptr.dtype.element_ty), mask=mask)


def _launch(kernel, *tensors):
    n_rows, n_cols = tensors[0].shape
    out = torch.empty_like(tensors[0])
    kernel[(n_rows,)](*tensors, out, n_cols, out.stride(0), BLOCK_SIZE=triton.next_power_of_2(n_cols))
    return out


class TritonSoftmax(torch.autograd.Function):
    """Своё ядро как обычная операция PyTorch: autograd вызовет backward сам."""

    @staticmethod
    def forward(ctx, x):
        y = _launch(softmax_fwd_kernel, x.contiguous())
        ctx.save_for_backward(y)                     # для градиента достаточно выхода
        return y

    @staticmethod
    def backward(ctx, dy):
        (y,) = ctx.saved_tensors
        return _launch(softmax_bwd_kernel, y, dy.contiguous())


triton_softmax = TritonSoftmax.apply


def naive_softmax(x):                                # «как написал бы человек»: 5 ядер в eager-режиме
    x_max = x.max(dim=1, keepdim=True).values
    z = x - x_max
    num = torch.exp(z)
    return num / num.sum(dim=1, keepdim=True)


torch.manual_seed(0)
x = torch.randn(1823, 781, device=DEVICE, requires_grad=True)   # «неудобные» размеры: проверяем маски
y = triton_softmax(x)
torch.testing.assert_close(y, torch.softmax(x, dim=1))
dy = torch.randn_like(y)
(gx,) = torch.autograd.grad(y, x, dy)
(gx_ref,) = torch.autograd.grad(torch.softmax(x, dim=1), x, dy)
torch.testing.assert_close(gx, gx_ref)
print("softmax: прямой и обратный проход совпадают с torch.softmax")

if DEVICE == "cuda":
    compiled = torch.compile(naive_softmax)
    for n_cols in [256, 1024, 4096, 16384]:
        x = torch.randn(4096, n_cols, device=DEVICE, dtype=torch.bfloat16)
        line = []
        for name, fn in [("naive", naive_softmax), ("torch", lambda t: torch.softmax(t, 1)),
                         ("compile", compiled), ("triton", triton_softmax)]:
            ms = triton.testing.do_bench(lambda: fn(x))
            line.append(f"{name} {2 * x.numel() * x.element_size() / (ms * 1e-3) / 1e9:5.0f}")
        print(f"n_cols = {n_cols:5d}, ГБ/с: " + "  ".join(line))

Обратите внимание на .to(tl.float32) после загрузки: ядро принимает тензоры BF16, но считает экспоненты и суммы во FP32 и только при записи приводит обратно. Это та же логика, что у autocast (урок 5.4.1): хранить узко, считать широко.

Чего ждать (оценка). naive — в несколько раз медленнее остальных. torch (встроенное ядро PyTorch), compile (torch.compile сам сгенерирует почти такое же ядро на Triton) и triton — близко друг к другу и к пределу памяти. При n_cols = 16384 строка в BF16 занимает 32 КБ, и ядру с одной программой на строку может не хватить регистров — производительность просядет. Официальный туториал Triton разбирает, как это лечить.

Сделайте сами: слитый RMSNorm (forward)

RMSNorm из урока 4.5.3: $y = x / \sqrt{\mathrm{mean}(x^2) + \varepsilon} \cdot w$. По образцу softmax напишите ядро «одна программа — одна строка». Подсказки: хвост читайте как 0 (он не меняет сумму квадратов), веса $w$ — одни на все строки, считайте во FP32. Решение:

import os
import torch
import triton
import triton.language as tl

DEVICE = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"


@triton.jit
def rmsnorm_kernel(x_ptr, w_ptr, y_ptr, n_cols, stride, eps, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(x_ptr + row * stride + cols, mask=mask, other=0.0).to(tl.float32)   # считаем во FP32
    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    rms = tl.sqrt(tl.sum(x * x, axis=0) / n_cols + eps)   # нули в хвосте сумму не меняют
    tl.store(y_ptr + row * stride + cols, (x / rms * w).to(y_ptr.dtype.element_ty), mask=mask)


def rmsnorm(x, w, eps=1e-6):
    y = torch.empty_like(x)
    rmsnorm_kernel[(x.shape[0],)](x, w, y, x.shape[1], x.stride(0), eps,
                                  BLOCK_SIZE=triton.next_power_of_2(x.shape[1]))
    return y


x = torch.randn(512, 896, device=DEVICE)                  # 896 — ширина Qwen2.5-0.5B
w = torch.randn(896, device=DEVICE)
ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w
torch.testing.assert_close(rmsnorm(x, w), ref)
print("rmsnorm: совпадает")

В PyTorch без компиляции RMSNorm — это pow, mean, add, rsqrt, два умножения: шесть ядер. Сравните ГБ/с своего ядра, наивной версии и torch.nn.functional.rms_norm.

Ядро 3: тайловый матмул

Теперь самое интересное — матмул из практики урока 5.7.1, но на настоящей SRAM и тензорных ядрах. Каждая программа считает тайл $C$ размером $\text{BLOCK\_M} \times \text{BLOCK\_N}$: идёт по $K$ тайлами, загружает тайл $A$ и тайл $B$ и накапливает tl.dot(a, b, acc) во FP32. Двумерные блоки адресов строятся «бродкастингом» (урок 0.3.2): offs_m[:, None] — столбец, offs_k[None, :] — строка, их сумма — таблица адресов тайла.

Две новые вещи:

• порядок программ группами (GROUP_M). Если нумеровать тайлы $C$ построчно, подряд работающие программы читают одну полосу $A$ и разные полосы $B$ — по всей ширине матрицы. При группировке по 8 строк тайлов подряд идут программы с 8 полосами $A$ и несколькими полосами $B$, и все эти данные успевают переиспользоваться из L2. Так преодолевается предел «тайл в shared memory даёт 64 FLOP/байт» из урока 5.7.1;
• автоподбор (@triton.autotune). Лучшие размеры тайлов, число варпов и стадий конвейера загрузки (num_stages — сколько тайлов грузится заранее, пока считаются предыдущие) зависят от карты и размеров. Triton прогонит все конфигурации на первом вызове с новыми (M, N, K) и запомнит лучшую.

"""matmul_triton.py — тайловый матмул на Triton с автоподбором тайлов и сравнение с cuBLAS."""
import os

import torch
import triton
import triton.language as tl

DEVICE = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"

CONFIGS = [
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=3, num_warps=8),
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=4, num_warps=4),
    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_stages=4, num_warps=4),
    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4),
]
if DEVICE == "cpu":            # интерпретатор не умеет замерять время, а автоподбору нужны замеры:
    CONFIGS = CONFIGS[-1:]     # с одной конфигурацией Triton пропускает перебор


@triton.autotune(configs=CONFIGS, key=["M", "N", "K"])   # перебрать конфигурации и запомнить лучшую для (M, N, K)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr):
    # 1. Какой тайл C считает эта программа. Порядок «группами по GROUP_M строк тайлов»:
    #    соседние программы берут одни и те же полосы A и B, и те приходят из L2, а не из видеопамяти
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    first_pid_m = (pid // num_pid_in_group) * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # 2. Указатели на первые тайлы A (BLOCK_M × BLOCK_K) и B (BLOCK_K × BLOCK_N)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak   # двумерная сетка адресов
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn

    # 3. Идём по K тайлами, накапливая сумму во FP32 (в регистрах)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        k_left = K - k * BLOCK_K                     # сколько столбцов K осталось: маска для хвоста
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_left), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_left) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)                      # матмул тайлов на тензорных ядрах: acc += a @ b
        a_ptrs += BLOCK_K * stride_ak                # сдвиг к следующему тайлу по K
        b_ptrs += BLOCK_K * stride_bk

    # 4. Записываем тайл C, приведя к типу выхода
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


def matmul(a, b):
    M, K = a.shape
    K2, N = b.shape
    assert K == K2, "внутренние размеры не совпадают"
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
    matmul_kernel[grid](a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                        c.stride(0), c.stride(1), GROUP_M=8)
    return c


torch.manual_seed(0)
dtype = torch.bfloat16 if DEVICE == "cuda" else torch.float32
a = torch.randn(1000, 700, device=DEVICE, dtype=dtype)    # размеры не кратны тайлам
b = torch.randn(700, 900, device=DEVICE, dtype=dtype)
torch.testing.assert_close(matmul(a, b), a @ b, atol=0.05, rtol=0.02)   # допуск под BF16
print("matmul: совпадает с a @ b")

if DEVICE == "cuda":
    for n in [1024, 2048, 4096, 8192]:
        a = torch.randn(n, n, device=DEVICE, dtype=dtype)
        b = torch.randn(n, n, device=DEVICE, dtype=dtype)
        line = []
        for name, fn in [("cuBLAS", lambda: a @ b), ("triton", lambda: matmul(a, b))]:
            ms = triton.testing.do_bench(fn)
            line.append(f"{name} {2 * n**3 / (ms * 1e-3) / 1e12:5.1f} TFLOPS")
        print(f"n = {n}: " + "  ".join(line))

Проверим, что конфигурации помещаются в shared memory Ada (до 99 КБ на блок): самая большая — тайлы $A$ $128 \times 32$ и $B$ $32 \times 128$ в BF16, по 8 КБ, три стадии конвейера: $3 \cdot 16 = 48$ КБ.

Чего ждать (оценка). На $n = 4096 \dots 8192$ cuBLAS в BF16 — порядка 70–85% от паспортных 97.5 TFLOPS; Triton — от 80% до почти 100% от cuBLAS (официальный туториал утверждает, что результат «на уровне cuBLAS»). На $n = 1024$ оба заметно ниже: программ ($8 \cdot 8 = 64$ при тайле 128) меньше, чем SM (76), и часть карты простаивает. Попробуйте добавить в список свои конфигурации и посмотрите, какую выберет автоподбор (matmul_kernel.best_config).

Когда Triton, а когда нет

Инструмент	Что это	Когда брать
torch.compiletorch.compileEN torch.compileКомпилятор PyTorch 2: TorchDynamo перехватывает граф операций модели, TorchInductor генерирует слитые ядра (на GPU — на Triton). Ускоряет обучение и инференс без изменения кода модели; первый вызов долгий из-за компиляции.	автоматическое слияние, генерирует Triton	всегда сначала: бесплатно и часто достаточно
Triton	ядра на Python, блоками	новый алгоритм: своё внимание, слитый слой, квантованный матмул
CUDA C++	полный контроль над потоками и памятью	последние проценты скорости, особые инструкции железа, библиотеки
CUTLASS / CuTe (NVIDIA)	шаблоны C++ (и CuTe DSL на Python) для матмулов и внимания	FlashAttention-3, ядра под Hopper и Blackwell
TileLang	язык тайловых ядер на Python (Пекинский университет)	ядра DeepSeek, переносимость на разные ускорители

Порядок действий на практике: профайлер (урок 5.7.1) → torch.compile → если узкое место осталось и понятно, какой алгоритм нужен, — Triton → и только если и этого мало, CUDA. Ядро без теста против эталонной реализации на PyTorch не пишут вообще: ошибка в маске даёт не падение, а тихо неверные числа.

Практика

Видеокарты в браузере нет, но модель программирования Triton можно потрогать и на NumPy. В задании — «мини-Triton»: load и store с масками и launch, который вызывает программы сетки по очереди. Напишите на нём три ядра урока; тест на трафик сверит матмул со счётом из урока 5.7.1.

Ядра в стиле Triton на NumPynumpy · нужен интернет6 тестов

import numpy as np

# ── Мини-Triton на NumPy: готовые помощники ──────────────────────────────────
# Настоящий Triton компилирует ядро для GPU и запускает все программы сетки параллельно.
# Здесь launch вызывает программы по очереди, а load/store ведут себя как tl.load/tl.store.
STATS = {"loads": 0}                                  # сколько чисел прочитано через load


def arange(n):
    """Как tl.arange(0, n)."""
    return np.arange(n)


def load(arr, offs, mask, other=0.0):
    """Как tl.load: читает плоский массив arr по индексам offs там, где mask, иначе — other.
    Чтение мимо массива без маски — ошибка IndexError, как сбой памяти на GPU."""
    flat = arr.reshape(-1)
    out = np.full(offs.shape, other, dtype=float)
    out[mask] = flat[offs[mask]]
    STATS["loads"] += int(np.count_nonzero(mask))
    return out


def store(arr, offs, values, mask):
    """Как tl.store: пишет values в плоский массив arr по индексам offs там, где mask."""
    flat = arr.reshape(-1)                            # вид на тот же массив, без копии
    flat[offs[mask]] = values[mask]


def launch(kernel, grid, *args):
    """Как kernel[grid](*args): вызывает kernel(pid, *args) для каждого pid сетки; pid — кортеж."""
    for pid in np.ndindex(*grid):
        kernel(pid, *args)


def cdiv(a, b):
    """Как triton.cdiv: деление с округлением вверх."""
    return (a + b - 1) // b


# ── Ваша часть ───────────────────────────────────────────────────────────────
def add_kernel(pid, x, y, out, n, BLOCK):
    """Программа номер pid[0] складывает свой кусок из BLOCK чисел: out = x + y.
    Последний блок может вылезти за конец массивов — нужна маска offs < n."""
    # TODO: offs = pid[0] * BLOCK + arange(BLOCK); mask; load, load, store
    ...


def softmax_kernel(pid, x, out, n_cols, BLOCK):
    """Программа номер pid[0] считает устойчивый softmax строки pid[0] матрицы x (n_rows × n_cols).
    BLOCK ≥ n_cols — степень двойки; хвост блока читается как -inf, чтобы не попасть в сумму."""
    # TODO
    ...


def matmul_kernel(pid, A, B, C, M, N, K, BM, BN, BK):
    """Программа (pid_m, pid_n) считает тайл C[pid_m·BM : …, pid_n·BN : …] = A @ B.
    A: (M, K), B: (K, N), C: (M, N) — плоские индексы строка·ширина + столбец.
    Цикл по K тайлами по BK: загрузить тайл A (BM × BK) и тайл B (BK × BN) с масками, acc += a @ b."""
    pid_m, pid_n = pid
    offs_m = pid_m * BM + arange(BM)
    offs_n = pid_n * BN + arange(BN)
    acc = np.zeros((BM, BN))
    # TODO: цикл по k0 in range(0, K, BK), маски a_mask и b_mask, load, acc += a @ b
    ...
    # TODO: store тайла C с маской


def triton_matmul(A, B, BM=16, BN=16, BK=16):
    """Запуск matmul_kernel на сетке (cdiv(M, BM), cdiv(N, BN)). Возвращает C."""
    M, K = A.shape
    N = B.shape[1]
    C = np.zeros((M, N))
    # TODO: launch(...)
    return C


if __name__ == "__main__":
    x, y = np.arange(10.0), np.ones(10)
    out = np.zeros(10)
    launch(add_kernel, (cdiv(10, 4),), x, y, out, 10, 4)
    print("add:", out)   # ждём [1, 2, ..., 10]

ЗапуститьПроверитьПодсказкаСбросить

🖥 На вашей машине (4×GPU)

1. Запуск. Сохраните три скрипта урока (vector_add.py, softmax_triton.py, matmul_triton.py) и запустите на одной карте:

python -c "import torch, triton; print(torch.__version__, triton.__version__, torch.cuda.get_device_capability())"
CUDA_VISIBLE_DEVICES=0 python vector_add.py
CUDA_VISIBLE_DEVICES=0 python softmax_triton.py
CUDA_VISIBLE_DEVICES=0 python matmul_triton.py     # автоподбор займёт до минуты

Ожидаем версию Triton 3.x и (8, 9). Код урока проверен на Triton 3.8 в режиме интерпретатора; если в вашей версии API отличается (например, у triton.testing.do_bench поменялись аргументы), сверьтесь с туториалами из источников. Запишите в журнал ГБ/с и TFLOPS и сравните с roofline из урока 5.7.1: сколько процентов паспорта получили вы, а сколько — PyTorch?

2. Профиль и сгенерированный код. Профайлер из урока 5.7.1 покажет ваши ядра по именам (softmax_fwd_kernel, matmul_kernel). Запустите TORCH_LOGS=output_code на torch.compile(naive_softmax) и сравните сгенерированное ядро со своим — оно будет очень похожим.

3. Отладка. Неверный результат — сначала TRITON_INTERPRET=1 и print внутри ядра на маленьком примере (строка из 5 чисел, блок 8). Падение с «illegal memory access» почти всегда означает забытую маску.

4. Задача Boss: forward FlashAttention на Triton. По мотивам задания 2 Stanford CS336. Перенесите flash_attention из практики урока 5.7.2 в ядро Triton: сетка — (блоки запросов) × (батч · головы), внутри — цикл по блокам $K$, $V$ с online softmax и tl.dot, пропуск блоков правее диагонали при причинной маске. Ядро должно вернуть и $O$, и $L$ (logsumexp). Тесты:

"""test_flash.py — проверка forward FlashAttention на Triton.  pytest -q test_flash.py"""
import math
import os

import pytest
import torch
import torch.nn.functional as F

from flash_triton import flash_attention

DEVICE = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"
DTYPE = torch.float32 if DEVICE == "cpu" else torch.bfloat16
TOL = 1e-4 if DEVICE == "cpu" else 2e-2              # BF16 даёт расхождение порядка 1e-2


@pytest.mark.parametrize("T", [64, 100, 257])        # 100 и 257 не делятся на блок — проверка масок
@pytest.mark.parametrize("causal", [False, True])
def test_matches_sdpa(T, causal):
    torch.manual_seed(0)
    q, k, v = (torch.randn(2, 3, T, 64, device=DEVICE, dtype=DTYPE) for _ in range(3))
    o, _ = flash_attention(q, k, v, causal=causal)
    ref = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
    torch.testing.assert_close(o, ref, atol=TOL, rtol=TOL)


def test_logsumexp():
    torch.manual_seed(1)
    q, k, v = (torch.randn(1, 2, 100, 32, device=DEVICE, dtype=DTYPE) for _ in range(3))
    _, L = flash_attention(q, k, v, causal=True)
    s = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(32)
    s = s.masked_fill(~torch.tril(torch.ones(100, 100, dtype=torch.bool, device=DEVICE)), float("-inf"))
    torch.testing.assert_close(L, torch.logsumexp(s, dim=-1), atol=TOL, rtol=TOL)

pytest -q test_flash.py                              # на GPU
TRITON_INTERPRET=1 pytest -q test_flash.py           # без GPU, на CPU (около минуты)

Уровни: Strong — проходят тесты без маски; Boss — все тесты, плюс замер против F.scaled_dot_product_attention (FlashAttention-2) при $T = 1024 \dots 16\,384$: ваше ядро медленнее, но память у обоих растёт линейно. Хотите ещё — обратный проход через пересчёт по формулам урока 5.7.2 и torch.autograd.Function.

Эталонное решение (откройте после своей попытки)

"""flash_triton.py — прямой проход FlashAttention на Triton (по мотивам CS336, задание 2)."""
import math

import torch
import triton
import triton.language as tl


@triton.jit
def flash_fwd_kernel(q_ptr, k_ptr, v_ptr, o_ptr, l_ptr, T, scale, stride_bh, stride_t,
                     D: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, CAUSAL: tl.constexpr):
    pid_m = tl.program_id(0)                         # какой блок запросов
    bh = tl.program_id(1)                            # какая пара (пример батча, голова)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, D)
    base = bh * stride_bh
    q = tl.load(q_ptr + base + offs_m[:, None] * stride_t + offs_d[None, :], mask=offs_m[:, None] < T, other=0.0)

    m_i = tl.full((BLOCK_M,), -float("inf"), dtype=tl.float32)   # бегущий максимум строк
    l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)                 # бегущая сумма экспонент
    acc = tl.zeros((BLOCK_M, D), dtype=tl.float32)               # ненормированный выход
    if CAUSAL:
        hi = tl.minimum((pid_m + 1) * BLOCK_M, T)    # блоки K правее диагонали не нужны
    else:
        hi = T
    for start_n in range(0, hi, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        kv_mask = offs_n[:, None] < T
        k = tl.load(k_ptr + base + offs_n[:, None] * stride_t + offs_d[None, :], mask=kv_mask, other=0.0)
        v = tl.load(v_ptr + base + offs_n[:, None] * stride_t + offs_d[None, :], mask=kv_mask, other=0.0)
        s = tl.dot(q, tl.trans(k)) * scale           # оценки блока (BLOCK_M, BLOCK_N), во FP32
        valid = offs_n[None, :] < T
        if CAUSAL:
            valid = valid & (offs_n[None, :] <= offs_m[:, None])
        s = tl.where(valid, s, -float("inf"))
        m_new = tl.maximum(m_i, tl.max(s, axis=1))   # online softmax, как в уроке 5.7.2
        p = tl.exp(s - m_new[:, None])
        alpha = tl.exp(m_i - m_new)
        l_i = l_i * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m_i = m_new
    o = acc / l_i[:, None]
    tl.store(o_ptr + base + offs_m[:, None] * stride_t + offs_d[None, :], o.to(o_ptr.dtype.element_ty),
             mask=offs_m[:, None] < T)
    tl.store(l_ptr + bh * T + offs_m, m_i + tl.log(l_i), mask=offs_m < T)   # logsumexp для backward


def flash_attention(q, k, v, causal=True, BLOCK_M=64, BLOCK_N=64):
    """q, k, v: (B, H, T, D), D — степень двойки от 16. Возвращает (O, L): (B, H, T, D) и (B, H, T)."""
    B, H, T, D = q.shape
    q, k, v = (t.reshape(B * H, T, D).contiguous() for t in (q, k, v))
    o = torch.empty_like(q)
    L = torch.empty((B * H, T), device=q.device, dtype=torch.float32)
    grid = (triton.cdiv(T, BLOCK_M), B * H)          # блоки запросов × (батч · головы)
    flash_fwd_kernel[grid](q, k, v, o, L, T, 1.0 / math.sqrt(D), q.stride(0), q.stride(1),
                           D=D, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, CAUSAL=causal)
    return o.view(B, H, T, D), L.view(B, H, T)

Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Китайский угол: TileLang и ядра DeepSeek

TileLang (github.com/tile-ai/tilelang) — язык тайловых ядер на Python, созданный в Пекинском университете (группа профессора Чжи Яна, часть работы — во время стажировки в Microsoft Research). Идея близка к Triton, но программист управляет раскладкой тайлов по памяти явнее. В README заявлена поддержка NVIDIA от SM70 до SM120 (то есть и вашей Ada), AMD, Apple Metal и Huawei Ascend, а среди примеров — FlashAttention, декодирование MLA и ядра для DeepSeek V3.2 и V4.

DeepSeek пишет на TileLang библиотеку TileKernels (MIT): десятки ядер для обучения и инференса — маршрутизация MoE, приведение к FP8 и FP4, слитые нормализации, RoPE. Требует SM90 или SM100 (Hopper, Blackwell), а 30 сентября 2026 года в репозиторий добавили поддержку Huawei Ascend 950. Для китайских лабораторий переносимость ядер между NVIDIA и отечественными ускорителями — вопрос не удобства, а выживания под экспортными ограничениями (урок china/02).

flash-linear-attention — ядра линейного внимания на Triton, на которых стоят слои Qwen3-Next и Kimi Linear (урок 5.5.3). Это лучший открытый пример большой библиотеки на Triton: после этого урока её код можно читать.

Где учиться по-китайски. llm-algo-leetcode (Datawhale) — задачи «условие → решение → тесты» по цепочке PyTorch → Triton → CUDA. diy-llm — китайская адаптация CS336 с заданиями на Triton.

Термины. Отдельную операцию на GPU по-китайски обычно называют «оператором» — буквально «считающим элементом», а само ядро GPU — «ядерной функцией». Тайлинг — «разбиение на блоки», автоподбор — «автоматическая настройка», маска — «маскирующий код».

Проверь себя

Проверь себя6 вопросов·

1 Вектор из 10 000 чисел складывается ядром с BLOCK_SIZE = 1024. Сколько программ нужно запустить?

2 Что произойдёт, если убрать маску из tl.load в ядре сложения при $n$, не кратном размеру блока?

Ничего: Triton сам обрежет лишнееПоследний блок прочитает и запишет за пределами массивов — мусор в результате или падение с illegal memory accessЯдро не скомпилируетсяРезультат будет верным, но медленнее

3 Почему хвост строки в ядре softmax читается с other=-float("inf"), а в ядре RMSNorm — с other=0.0?

Так требует синтаксис Triton$-\infty$ не влияет на максимум и даёт $e^{-\infty} = 0$ в сумме; 0 не меняет сумму квадратовЭто ускоряет загрузкуРазницы нет, можно наоборот

4 Наивный softmax на PyTorch делает 8 проходов по матрице, слитое ядро — 2. Во сколько раз меньше трафик памяти у слитого ядра?

5 Какие утверждения о torch.autograd.Function верны? Выберите все.Можно выбрать несколько ответов

В forward можно вызвать своё ядро TritonЧерез ctx.save_for_backward сохраняют тензоры, нужные для градиентаbackward возвращает градиенты по входам forwardAutograd сам выведет градиент ядра Triton, backward писать не нужно

6 Матмул $1024 \times 1024$ с тайлом $C$ $128 \times 128$. Сколько программ в сетке? (Сравните с числом SM у RTX 4080 — 76.)

Проверить ответы

Итоги

• Ядро на Triton описывает одну программу, которая обрабатывает блок данных; программы запускаются сеткой и узнают свой номер через tl.program_id. Маски в tl.load/tl.store защищают последний неполный блок.
• Каждое ядро сверяют с PyTorch через torch.testing.assert_close и замеряют triton.testing.do_bench в ГБ/с или TFLOPS — с оглядкой на roofline. Без GPU помогает TRITON_INTERPRET=1.
• Выигрыш дают слияние (softmax, RMSNorm: в разы меньше трафика) и тайлы на тензорных ядрах (tl.dot) с порядком обхода, который бережёт L2. В обучение ядро встраивается через torch.autograd.Function.
• Порядок действий: профайлер → torch.compile → Triton → CUDA. Альтернативы для особых случаев — CUTLASS/CuTe и TileLang, на котором пишет ядра DeepSeek.

Дальше: Четыре видеокарты: DDP и FSDP на практике — от одной карты ко всем четырём: обучение GPT через torchrun, замер масштабирования и узкое место PCIe.