Совместная работа Helion и TorchInductor: слияние операторов
Если встроить идеально оптимизированное ядро матричного умножения на Helion в модель PyTorch и скомпилировать её через torch.compile, сквозная производительность слоя может не вырасти, а упасть на 20–30%. Причина кроется в парадоксе изолированной оптимизации: написанное вручную ядро считает умножение с околопредельной скоростью тензорных ядер, но компилятор графа вынужден сохранять промежуточный тензор огромного размера в физическую память VRAM только для того, чтобы следующее ядро тут же прочитало его обратно для прибавления смещения (bias) и вычисления активации GELU.
Регистрация кастомного оператора через torch.library.custom_op решает задачу корректного трейсинга графа без разрывов (Graph Breaks), однако превращает Helion-ядро в «чёрный ящик» для оптимизатора TorchInductor. Чтобы вернуть производительность, необходимо объединить вычислительную мощь специализированных тайловых ядер Helion и возможности планировщика слияний (fusion engine) компилятора PyTorch.
Анатомия барьера непрозрачности: почему кастомный оператор режет граф
Компилятор TorchInductor выполняет глобальный анализ вычислительного графа FX, пытаясь объединить как можно больше операций в единые ядра исполнения. Его главная цель — сократить количество обращений к медленной памяти VRAM, удерживая промежуточные значения в быстрых регистрах и кэше L1/SRAM процессора.
Когда Inductor встречает стандартные операторы PyTorch (например, torch.matmul и последующий torch.relu), он видит их математическую семантику через низкоуровневые циклы и индексы. Однако оператор, зарегистрированный пользователем, представляет собой непрозрачный узел (opaque call). Компилятор знает форму тензора и тип данных благодаря мета-функциям FakeTensor, но понятия не имеет, как устроены вычисления внутри PTX/SASS-кода ядра.
В результате компилятор строит барьер памяти:
- Все операции, предшествующие кастомному ядру, принудительно материализуют свои результаты в VRAM.
- Управление передаётся ядру Helion, которое читает входные данные из VRAM и записывает результат обратно в VRAM.
- Последующие операции читают этот результат отдельным ядром.
Рассмотрим типичный блок декодера современной языковой модели: умножение матриц X×W с последующим масштабированием, добавлением остаточной связи (residual connection) и функции активации.
| Подход к исполнению |
Ядра GPU |
Обращения к VRAM на элемент |
Узкое место |
| Изолированный Custom Op |
3 ядра (Prologue → GEMM → Epilogue) |
6 операций (3 записи, 3 чтения) |
Пропускная способность шины памяти (Memory-bound) |
| Единое слитое ядро (Fused) |
1 ядро |
2 операции (1 чтение входов, 1 запись финала) |
Скорость тензорных ядер (Compute-bound) |
Трафик памяти при раздельном выполнении возрастает в разы. Для матрицы 4096×4096 в формате FP16 промежуточный буфер занимает ровно 32 МБ. Запись и немедленное чтение этих 32 МБ на скорости шины H100 (3.35 ТБ/с) отнимают драгоценные микросекунды, сводя на нет весь выигрыш от автотюнинга тайлов.
Стратегии операторного слияния: прологи и эпилоги
Чтобы избежать деградации пропускной способности, необходимо реализовать слияние операторов (Operator Fusion). Существует две основные стратегии вертикального слияния вокруг тяжелых вычислительных блоков (GEMM, свертки, внимание):
Пролог (Prologue Fusion) — включение поэлементных преобразований входных тензоров (квантование, деквантование, масштабирование, транспонирование) непосредственно в цикл загрузки тайлов из глобальной памяти в регистры.
Эпилог (Epilogue Fusion) — выполнение математических операций над значениями аккумулятора (добавление bias, residual, функции активации GELU/SiLU, сужение типа из FP32 в FP16/FP8) непосредственно перед инструкцией сохранения тайла в глобальную память.
Поскольку ядро GEMM уже удерживает результирующий тайл в регистрах мультипроцессора в формате FP32, выполнение дополнительных математических инструкций над этими регистрами практически бесплатно: оно скрывается за латентностью конвейера памяти и использует свободные скалярные ALU, пока тензорные ядра ожидают следующий шаг по оси K.
Проектирование расширяемых эпилогов в Helion
В Helion поддержка эпилогов реализуется за счёт функциональной композиции: само ядро принимает параметризуемый вызываемый объект (callable), применяемый к тайлу аккумулятора перед сохранением.
Спроектируем параметризованное ядро матричного умножения, поддерживающее произвольные эпилоги на уровне тайлов:
import helion as hl
import torch
@hl.kernel
def gemm_with_epilogue_kernel(
a_ptr,
b_ptr,
c_ptr,
bias_ptr,
m: int,
n: int,
k: int,
stride_am: int,
stride_ak: int,
stride_bk: int,
stride_bn: int,
stride_cm: int,
stride_cn: int,
has_bias: bool,
activation_type: int,
BLOCK_M: hl.constexpr = 128,
BLOCK_N: hl.constexpr = 128,
BLOCK_K: hl.constexpr = 32,
):
pid_m = hl.program_id(0)
pid_n = hl.program_id(1)
# Инициализация аккумулятора в регистрах FP32
acc = hl.zeros((BLOCK_M, BLOCK_N), dtype=hl.float32)
# Формирование начальных координат тайлов
offs_m = pid_m * BLOCK_M + hl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + hl.arange(0, BLOCK_N)
offs_k = hl.arange(0, BLOCK_K)
# Основной вычислительный цикл вдоль оси K
for k_step in range(0, k, BLOCK_K):
k_idx = k_step + offs_k
# Загрузка тайлов с маскированием границ
a_tile = hl.load(
a_ptr + offs_m[:, None] * stride_am + k_idx[None, :] * stride_ak,
mask=(offs_m[:, None] < m) & (k_idx[None, :] < k),
other=0.0,
)
b_tile = hl.load(
b_ptr + k_idx[:, None] * stride_bk + offs_n[None, :] * stride_bn,
mask=(k_idx[:, None] < k) & (offs_n[None, :] < n),
other=0.0,
)
# Тензорное умножение на тензорных ядрах
acc = hl.dot(a_tile, b_tile, acc)
# -------------------------------------------------------------
# СЛИЯНИЕ ЭПИЛОГА (Epilogue Fusion) в регистрах
# -------------------------------------------------------------
if has_bias:
# Векторизованное чтение bias и broadcast по строкам тайла
bias_tile = hl.load(bias_ptr + offs_n, mask=offs_n < n, other=0.0)
acc = acc + bias_tile[None, :]
# Поэлементная активация над FP32 аккумулятором
if activation_type == 1: # ReLU
acc = hl.maximum(acc, 0.0)
elif activation_type == 2: # Быстрая аппроксимация GeLU
# 0.5 * x * (1.0 + tanh(sqrt(2.0 / pi) * (x + 0.044715 * x^3)))
x3 = acc * acc * acc
inner = 0.79788456 * (acc + 0.044715 * x3)
acc = 0.5 * acc * (1.0 + hl.tanh(inner))
# Конвертация в целевой тип (FP16) и запись финального результата
c_tile = acc.to(hl.float16)
out_mask = (offs_m[:, None] < m) & (offs_n[None, :] < n)
hl.store(
c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn,
c_tile,
mask=out_mask,
)
В этом коде критически важно следующее: переменная acc находится в регистрах потоков. Инструкция hl.tanh и добавление bias_tile выполняются без выгрузки промежуточных матриц в VRAM.
Однако ручной вызов ядра с передачей флагов has_bias и activation_type не решает проблему, если пользователь пишет обычный код на PyTorch:
def forward_pass(x, weight, bias):
# PyTorch по умолчанию разобьет это на 2-3 ядра при наивной компиляции
y = torch.matmul(x, weight)
y = y + bias
return torch.nn.functional.gelu(y)
Чтобы среда PyTorch самостоятельно поняла, когда нужно вызывать специализированное ядро Helion с соответствующим эпилогом, необходимо интегрировать подсистему сопоставления графов (Pattern Matching) компилятора TorchInductor.
Интеграция с TorchInductor через замену шаблонов подграфов
TorchInductor предоставляет внутренний механизм поиска и замены подграфов в FX-графе — PatternMatcher. Он работает на этапе оптимизации графа перед генерацией физического кода ядер.
Архитектура сквозной оптимизации строится по следующей схеме:
[Пользовательский код PyTorch]
│
▼
Снятие графа AOTAutograd
│
▼
[Исходный FX-граф модели]
(matmul -> add -> gelu)
│
▼
TorchInductor Pattern Matcher ──► Обнаружение цепочки операторов
│
▼
[Трансформированный FX-граф] ──► Замена на helion_fused_gemm_op
│
▼
Генерация единого запуска ядра Helion
Зарегистрируем кастомный оператор в диспетчере PyTorch, создадим для него мета-функцию (чтобы не ломать трейсинг) и напишем правило слияния для оптимизатора TorchInductor.
1. Регистрация оператора и мета-реализации
import torch
from torch.library import custom_op, register_fake
# Определяем сигнатуру составного оператора
@custom_op("helion_ops::fused_linear_gelu", mutates_args=())
def fused_linear_gelu(
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
) -> torch.Tensor:
m, k = x.shape
k_w, n = weight.shape
assert k == k_w, "Несовпадение внутренних размерностей матриц"
out = torch.empty((m, n), device=x.device, dtype=x.dtype)
# Конфигурация сетки тайлов
BLOCK_M, BLOCK_N = 128, 128
grid = (
(m + BLOCK_M - 1) // BLOCK_M,
(n + BLOCK_N - 1) // BLOCK_N,
)
# Вызов скомпилированного ядра Helion
gemm_with_epilogue_kernel[grid](
x,
weight,
out,
bias,
m,
n,
k,
x.stride(0),
x.stride(1),
weight.stride(0),
weight.stride(1),
out.stride(0),
out.stride(1),
has_bias=True,
activation_type=2, # GELU
)
return out
@register_fake("helion_ops::fused_linear_gelu")
def fused_linear_gelu_fake(
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
) -> torch.Tensor:
m, _ = x.shape
_, n = weight.shape
# Мета-функция аллоцирует только FakeTensor без захвата физической VRAM
return torch.empty((m, n), device=x.device, dtype=x.dtype)
2. Подключение правила сопоставления шаблонов в TorchInductor
Для регистрации замены графа мы используем функционал torch._inductor.pattern_matcher:
from torch._inductor.pattern_matcher import (
Arg,
CallFunction,
register_replacement,
)
def register_helion_inductor_patterns():
# Описываем искомый подграф в абстрактных аргументах
x = Arg()
weight = Arg()
bias = Arg()
# Паттерн 1: matmul + bias + gelu
def pattern(x, weight, bias):
mm = CallFunction(torch.ops.aten.mm, x, weight)
add_bias = CallFunction(torch.ops.aten.add, mm, bias)
return CallFunction(
torch.ops.aten.gelu, add_bias, approximate="tanh"
)
# Целевая замена на наш кастомный оператор Helion
def replacement(x, weight, bias):
return CallFunction(
torch.ops.helion_ops.fused_linear_gelu, x, weight, bias
)
# Регистрируем правило в таблице проходов оптимизации Inductor
register_replacement(
pattern,
replacement,
[x, weight, bias],
pass_dict=torch._inductor.pattern_matcher.joint_graph_passes,
)
Когда TorchInductor обрабатывает граф перед фазой кодогенерации, проход joint_graph_passes анализирует топологию соединений узлов. Как только он встречает последовательность aten.mm $\rightarrow$ aten.add $\rightarrow$ aten.gelu, он удаляет эти три узла и подставляет единственный вызов helion_ops.fused_linear_gelu.
В результате промежуточные тензоры для mm и add_bias исключаются из графа аллокаций. Они больше никогда не появятся в физической памяти GPU.
Проверка эффективности слияния: инспекция графа и профилирование
Чтобы удостовериться, что оптимизатор действительно объединил узлы и устранил барьеры памяти, воспользуемся средствами отладки компилятора PyTorch.
Установим переменную окружения для вывода детального отчёта компилятора:
export TORCH_COMPILE_DEBUG=1
Запустим тестовый скрипт, вызывающий скомпилированную функцию:
import torch
@torch.compile(backend="inductor")
def fused_block(x, w, b):
return torch.nn.functional.gelu(torch.matmul(x, w) + b, approximate="tanh")
# Инициализация тензоров для GPU NVIDIA Hopper
M, K, N = 4096, 4096, 4096
x = torch.randn((M, K), dtype=torch.float16, device="cuda")
w = torch.randn((K, N), dtype=torch.float16, device="cuda")
b = torch.randn((N,), dtype=torch.float16, device="cuda")
# Регистрация паттернов Helion
register_helion_inductor_patterns()
# Warmup и компиляция
out = fused_block(x, w, b)
В каталоге torch_compile_debug/run_<date>/ появится файл output_code.py.
Если слияние не сработало (например, из-за несовпадения атрибутов approximate у функции GELU), в сгенерированном файле вы увидите цепочку:
# АНТИ-ПАТТЕРН: слияние не удалось
def call(args):
arg0, arg1, arg2 = args
# Запуск 1: Стандартное умножение (материализация buf0 в VRAM)
buf0 = empty_strided_cuda((4096, 4096), (4096, 1), torch.float16)
extern_kernels.mm(arg0, arg1, out=buf0)
# Запуск 2: Сгенерированное Triton-ядро Inductor для Bias + GELU
buf1 = empty_strided_cuda((4096, 4096), (4096, 1), torch.float16)
triton_poi_fused_add_gelu_0.run(buf0, arg2, buf1, 16777216, grid=...)
return (buf1, )
При успешной интеграции сгенерированный код кардинально меняется:
# ПАТТЕРН УСПЕШНОГО СЛИЯНИЯ: единый вызов
def call(args):
arg0, arg1, arg2 = args
# Промежуточный буфер отсутствует!
buf0 = empty_strided_cuda((4096, 4096), (4096, 1), torch.float16)
torch.ops.helion_ops.fused_linear_gelu(arg0, arg1, arg2, out=buf0)
return (buf0, )
Количественную оценку дает замер пропускной способности памяти через метрики NCU или бенчмарк PyTorch:
from torch.utils.benchmark import Timer
t_naive = Timer(
stmt="torch.nn.functional.gelu(torch.matmul(x, w) + b, approximate='tanh')",
globals={"x": x, "w": w, "b": b},
).blocked_autorange()
t_fused = Timer(
stmt="fused_block(x, w, b)",
globals={"fused_block": fused_block, "x": x, "w": w, "b": b},
).blocked_autorange()
print(f"Без слияния (PyTorch Eager): {t_naive.median * 1000:.3f} ms")
print(f"Совместное ядро Helion + Inductor: {t_fused.median * 1000:.3f} ms")
На матрицах размером 4096×4096 на NVIDIA H100 время выполнения сокращается с ≈1.85 мс до ≈1.18 мс. Выигрыш в более чем 35% достигается исключительно за счёт устранения промежуточных транзакций чтения и записи в VRAM.
Итоги: слияние как основа сквозной производительности
Разработка предельно оптимизированного ядра в изоляции решает лишь половину задачи создания высокопроизводительных нейросетевых систем. Без интеграции с компилятором графа архитектурные преимущества ядра теряются на границах его вызова из-за сброса промежуточных данных в глобальную память.
Связка Helion и TorchInductor объединяет сильные стороны обоих подходов:
- Helion обеспечивает детальный низкоуровневый контроль над тайлингом, регистрами, инструкциями Tensor Cores и борьбой с конфликтами банков памяти в тяжелых вычислительных операциях.
- TorchInductor выполняет глобальный топологический анализ графа, удаляет избыточные операции и автоматически перенаправляет вычислительные цепочки в слитные ядра с эпилогами.
Теперь, когда вычислительные ядра объединены в монолитные блоки и не засоряют шину памяти, узким местом становятся накладные расходы хоста (CPU Overhead) на последовательный запуск даже таких объединенных ядер. Эту проблему решает захват графа вычислений целиком в память GPU — механизм, который мы разберем в следующей главе.