Настоящий материал представляет собой обзор главы 14 книги Chris Fregly, «AI Systems Performance Engineering: Optimizing Model Training and Inference Workloads with GPUs, CUDA, and PyTorch», O'Reilly Media, 2026 (Copyright © 2026 Flux Capacitor, LLC).
Прямые цитаты из книги выделены отдельными блоками с пометкой «Цитата из книги». Остальной текст — систематизированное изложение содержания главы, сопоставление с актуальными данными, проверка прогнозов и дополнительные замечания, выходящие за пределы текста главы. Цитаты приводятся в ограниченном объёме в целях обзора и критического анализа.
Глава 14 книги Fregly посвящена динамическому стеку компиляции PyTorch — от захвата графа на уровне байт-кода Python до генерации высокопроизводительных GPU-ядер через OpenAI Triton и альтернативный бэкенд XLA для TPU и других ускорителей. Материал охватывает как внутреннее устройство TorchDynamo, AOT Autograd, PrimTorch IR и TorchInductor, так и практические аспекты написания кастомных ядер, отладки компилятора и масштабирования на распределённые конфигурации.
«Компилятор PyTorch прошёл путь от экспериментальной возможности до production-ready стека, который обеспечивает производительность, близкую к рукописным библиотекам, при сохранении гибкости eager-режима разработки. Конвейер TorchDynamo → AOT Autograd → PrimTorch IR → TorchInductor → Triton/LLVM представляет собой вертикально интегрированный компилятор, где каждый этап решает свою задачу: захват графа, объединение проходов, канонизация, генерация кода и компиляция в PTX».
В главе подчёркивается, что производительность современных AI-моделей всё чаще определяется не только качеством отдельных ядер, но и качеством компиляции. Автоматическое слияние операций, подбор тайлов, генерация Tensor Core-кода — всё это берёт на себя компилятор. Задача инженера — понять, как работает этот стек, чтобы использовать его возможности и диагностировать проблемы.
Конвейер torch.compile состоит из четырёх последовательных этапов, каждый из которых решает свою задачу:
Далее код транслируется через Triton или LLVM в финальный PTX для NVIDIA GPU, либо через XLA для TPU и других ускорителей.
Пять уровней вертикально интегрированного конвейера
Компилятор PyTorch — вертикально интегрированный стек, где каждый уровень решает свою задачу. Понимание этой архитектуры необходимо для эффективной диагностики проблем производительности.
Производительность модели зависит не только от качества отдельных ядер, но и от качества компиляции. torch.compile автоматизирует большинство оптимизаций.
«TorchDynamo (или просто Dynamo) — первый этап компиляции. Он подключается к CPython через Frame Evaluation API, определённый в PEP 523. Этот механизм позволяет Dynamo перехватывать выполнение кадра Python до того, как интерпретатор начнёт его обрабатывать».
В главе TorchDynamo описывается как механизм захвата графа на уровне байт-кода. Когда интерпретатор Python выполняет тензорную операцию, Dynamo перехватывает её и добавляет как узел в граф FX, откладывая реальное выполнение. Этот процесс продолжается до тех пор, пока не встретится код, который Dynamo не может обработать. В этот момент происходит разрыв графа (graph break), и неподдерживаемые операции выполняются в обычном eager-режиме.
Для простой функции:
def f(x, y):
z = x.sin() + y
return z.sum()TorchDynamo создаст FX Graph, примерно эквивалентный следующему псевдокоду:
graph():
%x : Tensor = Placeholder[target=x]
%y : Tensor = Placeholder[target=y]
%sin : Tensor = torch.sin(%x)
%add : Tensor = torch.add(%sin, %y)
%sum : Tensor = torch.sum(%add)
return %sumКаждый узел соответствует вызову примитивного оператора ATen, такого как aten::sin, aten::add, aten::sum.
TorchDynamo расставляет guard'ы — проверки на значения Python, которые влияют на трассировку графа: формы тензоров, типы данных, значения глобальных переменных. Эти guard'ы гарантируют, что если что-то изменится, скомпилированный граф будет признан недействительным и при необходимости перекомпилирован.
Разрывы графа — главный враг производительности. Каждый разрыв означает, что оптимизированный граф обрезан, и вводится больше накладных расходов Python. Распространённые причины разрывов и способы их устранения:
print() не может быть захвачен в граф. Решение — обернуть его в проверку if not torch._dynamo.is_compiling().if tensor.sum() > 0: вызывает разрыв графа. Решения: torch.where(), torch.cond() или рефакторинг.x.add_(y) могут вызывать разрывы графа. Решение — переписать в out-of-place форме.torch.stack() вместо динамического построения списков.TorchDynamo захватывает граф на уровне байт-кода через Frame Evaluation API. Разрывы графа — главный враг производительности.
Устраняйте разрывы: заменяйте print() на проверку is_compiling(), используйте torch.where() вместо if, переписывайте in-place операции.
«AOT Autograd работает иначе. Он пропускает захваченный граф прямого прохода через движок autograd в функциональном режиме, записывая операции обратного прохода заранее. Результат — статический граф обратного прохода, объединённый с графом прямого прохода».
В главе AOT Autograd описывается как этап, на котором создаётся совместный граф прямого и обратного проходов. Это позволяет компилятору оптимизировать всю последовательность вычислений как единое целое, а не как набор отдельных операций.
AOT Autograd предоставляет несколько ключевых возможностей:
AOT Autograd создаёт совместный граф прямого и обратного проходов, что позволяет оптимизировать всю последовательность вычислений как единое целое.
Слияние ядер через границу проходов и повторное использование буферов снижают пиковое потребление памяти и ускоряют обучение.
«PyTorch имеет более 2000 операций в своём полном API. Для компилятора это огромная проблема — каждый бэкенд должен уметь обрабатывать все эти операции. PrimTorch IR решает эту задачу, сокращая разнообразие операций до примерно 250 примитивных».
В главе PrimTorch IR описывается как этап канонизации операций. Многие сложные операции PyTorch разлагаются на примитивы. Например, in-place операция x.add_(y) преобразуется в функциональный add с последующим явным копированием.
Стабильный интерфейс PrimTorch IR означает, что для поддержки нового ускорителя разработчикам нужно реализовать только 250 примитивов, а не тысячи операций ATen. Это радикально снижает порог входа для новых аппаратных платформ.
PrimTorch IR сокращает более 2000 операций PyTorch до ~250 примитивов, что упрощает поддержку новых аппаратных платформ.
In-place операции преобразуются в функциональные плюс отдельные узлы копирования, что упрощает оптимизацию и планирование памяти.
«TorchInductor (или просто Inductor) — стандартный бэкенд компилятора PyTorch. Он принимает оптимизированный FX Graph, состоящий из операций aten + PrimTorch IR, и генерирует высокопроизводительный код для NVIDIA GPU, AMD GPU, CPU и других платформ».
В главе TorchInductor описывается как бэкенд, который превращает граф в высокопроизводительный код. Для NVIDIA GPU он использует JIT-компилятор OpenAI Triton для генерации GPU-ядер. Triton понижается до NVIDIA PTX через LLVM NVPTX, не вызывая NVCC для компиляции ядер.
TorchInductor включает автотюнер, построенный на возможностях Triton. При первой компиляции кода с режимом max-autotune Inductor тратит дополнительное время на бенчмаркинг различных вариантов ядер — разные размеры блоков, тайлов, количества варпов. Самый быстрый вариант кэшируется и используется в дальнейшем.
TorchInductor использует CUDA Graphs для захвата последовательностей ядер во время выполнения. Это уменьшает накладные расходы на запуск при каждой итерации, что особенно полезно для инференса. Режимы reduce-overhead и max-autotune активируют CUDA Graphs.
Для стабильных форм и latency-критичного инференса предпочитайте torch.compile(mode="max-autotune") после того, как количество рекомпиляций под контролем. Держите границу захвата узкой (per-layer или attention block), чтобы уменьшить инвалидацию графа из-за ragged batching.
TorchInductor — стандартный бэкенд компилятора PyTorch. Для NVIDIA GPU он использует Triton, который понижается до PTX через LLVM NVPTX.
Автотюнинг и CUDA Graphs дают значительный выигрыш для больших и критичных нагрузок. Для небольших моделей выгода может быть незначительной.
«Triton использует модель SPMD (Single-Program, Multiple-Data), в отличие от SIMT-модели CUDA. Ядра Triton работают на уровне блоков потоков (CTA), а не отдельных потоков. Вам не нужно явно управлять потоками или варпами — компилятор Triton делает это за вас».
В главе Triton описывается как Python-ориентированный DSL для написания GPU-ядер. Он снижает барьер входа для разработчиков, не знакомых с CUDA C++. Под капотом Triton компилируется в PTX через LLVM.
Простое ядро сложения векторов:
import triton
import triton.language as tl
BLOCK_SIZE = 1024
@triton.jit
def vector_add_kernel(x_ptr, y_ptr, out_ptr, n_elements,
BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
result = x + y
tl.store(out_ptr + offsets, result, mask=mask)Здесь BLOCK_SIZE — константа времени компиляции. tl.arange(0, BLOCK_SIZE) возвращает вектор индексов, а tl.load и tl.store выполняют векторизованные загрузки и сохранения с маскированием на границах.
Чтобы ядро Triton было видимым для torch.compile и не вызывало разрыв графа, зарегистрируйте его как кастомную операцию:
from torch.library import triton_op, wrap_triton
@triton_op("my_triton_lib::vector_add", mutates_args=())
def vector_add(x: Tensor, y: Tensor) -> Tensor:
output = torch.empty_like(x)
n_elements = x.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
wrap_triton(vector_add_kernel)[grid](
x, y, output, n_elements, BLOCK_SIZE=1024
)
return output
c = torch.ops.my_triton_lib.vector_add(a, b)Производительность GPU-ядра высокочувствительна к параметрам времени компиляции. Встроенный автотюнер Triton автоматизирует поиск оптимальных настроек:
@triton.autotune(
configs=[
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64},
num_warps=4, num_stages=2),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64},
num_warps=8, num_stages=2),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128},
num_warps=8, num_stages=3,
num_consumer_groups=2,
num_buffers_warp_spec=3),
],
key=['M', 'N', 'K'],
)
@triton.jit
def gemm_kernel(...):
...Во время первого вызова Triton JIT-компилирует и тестирует каждую конфигурацию. Самая быстрая кэшируется для данного ключа, производного от характеристик входа.
Triton использует модель SPMD, абстрагируя управление потоками и варпами. Ядра пишутся на Python и компилируются в PTX через LLVM.
Регистрация кастомных ядер через torch.library.triton_op() делает их видимыми для torch.compile. Автотюнинг автоматизирует поиск оптимальных конфигураций.
«PyTorch XLA — отдельный бэкенд компиляции, ориентированный на Google Cloud TPU и другие ускорители, не поддерживаемые TorchInductor. XLA отображает операции PyTorch в IR графа XLA и выполняет их с использованием среды выполнения целевого оборудования».
В главе PyTorch XLA описывается как альтернативный бэкенд для не-CUDA ускорителей. Активация: torch.compile(..., backend="openxla"). XLA используется для Google Cloud TPU, внутреннего инференс-чипа Meta (MTIA), а также кастомных чипов AWS Inferentia и Trainium через AWS Neuron SDK.
PyTorch/XLA работает на модели ленивого выполнения (lazy execution): операции записываются в граф и выполняются только когда результат строго нужен. Это принципиально отличается от модели TorchDynamo, которая перехватывает байт-код.
Новый мост Dynamo для XLA (бета на 2025 год) использует guard-систему Dynamo для более надёжного захвата графов, чем ленивая трассировка.
PyTorch XLA — альтернативный бэкенд для TPU и не-CUDA ускорителей. Активируется через backend="openxla".
XLA оптимизирован для статических форм. Новые формы вызывают полную перекомпиляцию. Для TPU время компиляции может быть значительным.
«PyTorch DDP работает с TorchDynamo, явно разрывая графы в точках синхронизации all-reduce. Это ожидаемое поведение: компилятор компилирует редукцию каждого бакета градиента отдельно, чтобы коммуникация одного бакета перекрывалась с вычислением обратного прохода следующего».
В главе подчёркивается, что разрывы графа на границах коммуникации DDP необходимы для сохранения перекрытия вычислений и коммуникации. TorchDynamo + DDP работает с производительностью, аналогичной eager-режиму DDP, и может даже превосходить его в масштабе.
PyTorch FSDP поддерживается в скомпилированном режиме с use_original_params=True. Лучшая практика — оборачивать каждый подмодуль (например, каждый блок трансформера) в отдельный подмодуль FSDP. Dynamo создаёт явные разрывы на границах подмодулей FSDP, позволяя коммуникации каждого шарда перекрываться с вычислениями.
Разрывы графа на границах коммуникации DDP и FSDP необходимы для сохранения перекрытия вычислений и коммуникации.
Оборачивайте подмодули в FSDP индивидуально для пошардовой компиляции и экономии памяти. Тестируйте FSDP с torch.compile сначала на меньшей конфигурации.
Книга Fregly писалась в 2025 году. С тех пор появились новые данные, которые позволяют проверить некоторые из утверждений главы и уточнить цифры.
«Компилятор PyTorch прошёл путь от экспериментальной возможности до production-ready стека».
Подтверждено. TorchInductor и Triton широко используются в production. PyTorch 2.x с torch.compile стал стандартом для оптимизации производительности. NVIDIA продолжает обновлять Triton для поддержки новейших архитектур, включая Blackwell.
«На момент написания TorchInductor всё ещё использует Triton в качестве основного пути генерации GPU-кода».
Подтверждено. Triton остаётся основным путём генерации GPU-кода в TorchInductor. Однако появились конкуренты: NVIDIA анонсировала Python-ориентированные библиотеки CUDA (cuTile, CuTe DSL), а также появились альтернативные подходы, такие как Cypress и LLM-агенты, транслирующие ядра напрямую в PTX.
«PyTorch/XLA — отдельный бэкенд компиляции, ориентированный на Google Cloud TPU и другие ускорители, не поддерживаемые TorchInductor».
Подтверждено. PyTorch/XLA продолжает развиваться (версия 2.6 в 2025 году), но остаётся нишевым решением для TPU и не-CUDA ускорителей. Google Cloud продолжает поддерживать и улучшать его, но для NVIDIA GPU стандартом является TorchInductor + Triton.
«FlexAttention fused kernels достигают до 85–90% производительности современного FlashAttention».
Частично подтверждено. FlexAttention действительно достиг значительного прогресса. С новым Flash-бэкендом на GB200 он может достигать 1.6–3.2x ускорения над Triton-реализацией для прямого прохода. Однако для некоторых паттернов и конфигураций FlashAttention всё ещё остаётся быстрее.
«Современные компиляторные фреймворки автоматически выполняют часть оптимизаций».
Подтверждено. TorchInductor автоматически выполняет слияние операций, автотюнинг и использует CUDA Graphs. PyTorch/XLA 2.6 представил scan-оператор и host offloading. Тренд на автоматизацию продолжается.
Глава описывает компилятор PyTorch, но оставляет за скобками несколько аспектов, которые проявились в реальной практике.
Глава описывает поддержку динамических форм, но не рассматривает ограничения. CUDA Graphs требуют статических форм. При включении dynamic=True CUDA Graphs не используются. Для переменных длин последовательностей рекомендуется mode="default" или mode="max-autotune-no-cudagraphs". Это важный компромисс между гибкостью и производительностью.
Глава описывает инструменты отладки, но не подчёркивает сложность диагностики проблем компилятора. Graph breaks, рекомпиляции, проблемы с числовой корректностью — всё это требует глубокого понимания внутреннего устройства компилятора. Инструменты вроде TORCH_LOGS и torch._dynamo.explain() помогают, но кривая обучения остаётся высокой.
Глава фокусируется на TorchInductor и XLA, но не рассматривает взаимодействие с TensorRT. В production-средах TensorRT часто используется для инференса, и torch.compile может интегрироваться с ним. Это важный аспект для production-развёртываний, который остаётся за пределами главы.
Глава описывает Triton, но не упоминает экосистему вокруг него. Triton активно развивается, появляются новые возможности: TLX (Triton Language Extensions) для низкоуровневого контроля, поддержка Blackwell с инструкциями tcgen05, MXFP4/MXFP8. Сообщество GPU MODE активно способствует развитию Triton.
torch.where(), torch.cond() и рефакторинг — ключевая задача инженера.max-autotune даёт максимальную производительность. Для переменных форм требуется default или max-autotune-no-cudagraphs.Глава 14 книги Fregly — систематический разбор стека компиляции PyTorch. Ключевая идея: компилятор PyTorch прошёл путь от эксперимента до production-ready стека, который обеспечивает производительность, близкую к рукописным библиотекам. Для инженера это означает, что понимание внутреннего устройства компилятора — от TorchDynamo до TorchInductor — необходимо для эффективной оптимизации и диагностики проблем. По мере движения к моделям в сотни триллионов параметров и системам вроде NVL72 роль компилятора будет только расти.
Источники: Chris Fregly, «AI Systems Performance Engineering» (O'Reilly, 2026), Глава 14 — цитаты приводятся в ограниченном объёме в целях обзора и критического анализа. Документация PyTorch, TorchInductor, Triton, PyTorch/XLA; MLPerf Training/Inference v5.0; открытые публикации NVIDIA и Google Cloud.
This material is a review of Chapter 14 of Chris Fregly, «AI Systems Performance Engineering: Optimizing Model Training and Inference Workloads with GPUs, CUDA, and PyTorch», O'Reilly Media, 2026 (Copyright © 2026 Flux Capacitor, LLC).
Direct quotes from the book are marked with separate blocks labeled «Quote from the book». The rest of the text is a systematized exposition of the chapter's content, a comparison with current data, forecast verification, and additional observations that go beyond the chapter's text. Quotes are provided in limited volume for review and critical analysis purposes.
Chapter 14 of Fregly's book is dedicated to the dynamic PyTorch compilation stack — from graph capture at the Python bytecode level to the generation of high-performance GPU kernels via OpenAI Triton and the alternative XLA backend for TPUs and other accelerators. The material covers the internals of TorchDynamo, AOT Autograd, PrimTorch IR, and TorchInductor, as well as practical aspects of writing custom kernels, debugging the compiler, and scaling to distributed configurations.
«The PyTorch compiler has come a long way from an experimental feature to a production-ready stack that delivers performance close to hand-written libraries while retaining the flexibility of eager development. The pipeline TorchDynamo → AOT Autograd → PrimTorch IR → TorchInductor → Triton/LLVM is a vertically integrated compiler where each stage solves its own task: graph capture, pass merging, canonicalization, code generation and compilation to PTX».
The chapter emphasizes that the performance of modern AI models is increasingly determined not only by the quality of individual kernels but also by the quality of compilation. Automatic operation fusion, tile selection, Tensor Core code generation — all of this is handled by the compiler. The engineer's task is to understand how this stack works in order to leverage its capabilities and diagnose problems.
The torch.compile pipeline consists of four sequential stages, each solving its own task:
The code is then translated through Triton or LLVM into final PTX for NVIDIA GPUs, or through XLA for TPUs and other accelerators.
Five levels of a vertically integrated pipeline
The PyTorch compiler is a vertically integrated stack where each level solves its own task. Understanding this architecture is necessary for effective performance diagnostics.
Model performance depends not only on the quality of individual kernels but also on the quality of compilation. torch.compile automates most optimizations.
«TorchDynamo (or simply Dynamo) is the first stage of compilation. It hooks into CPython through the Frame Evaluation API defined in PEP 523. This mechanism lets Dynamo intercept execution of a Python frame before the interpreter begins processing it».
In the chapter, TorchDynamo is described as a bytecode-level graph capture mechanism. When the Python interpreter executes a tensor operation, Dynamo intercepts it and adds it as a node in an FX graph, deferring its actual execution. This continues until it encounters code that Dynamo cannot handle. At that point, a graph break occurs, and unsupported operations run in normal eager mode.
For a simple function:
def f(x, y):
z = x.sin() + y
return z.sum()TorchDynamo creates an FX Graph roughly equivalent to the following pseudocode:
graph():
%x : Tensor = Placeholder[target=x]
%y : Tensor = Placeholder[target=y]
%sin : Tensor = torch.sin(%x)
%add : Tensor = torch.add(%sin, %y)
%sum : Tensor = torch.sum(%add)
return %sumEach node corresponds to a call to a primitive ATen operator such as aten::sin, aten::add, aten::sum.
TorchDynamo installs guards — checks on Python values that affect graph tracing: tensor shapes, dtypes, global variable values. These guards ensure that if something changes, the compiled graph is invalidated and recompiled if needed.
Graph breaks are the main enemy of performance. Every break means the optimized graph is truncated and more Python overhead is introduced. Common causes of breaks and how to fix them:
print() cannot be captured in the graph. Fix: wrap in a check if not torch._dynamo.is_compiling().if tensor.sum() > 0: causes a graph break. Fixes: torch.where(), torch.cond(), or refactoring.x.add_(y) may cause breaks. Fix: rewrite out-of-place.torch.stack() instead.TorchDynamo captures graphs at the bytecode level via the Frame Evaluation API. Graph breaks are the main enemy of performance.
Eliminate breaks: replace print() with is_compiling() checks, use torch.where() instead of if, rewrite in-place operations.
«AOT Autograd works differently. It passes the captured forward graph through the autograd engine in functional mode, recording backward operations ahead of time. The result is a static backward graph merged with the forward graph».
In the chapter, AOT Autograd is described as the stage where a joint forward and backward graph is created. This allows the compiler to optimize the entire computation sequence as a single whole rather than as a collection of individual operations.
AOT Autograd provides several key capabilities:
AOT Autograd creates a joint forward and backward graph, allowing the entire computation sequence to be optimized as a single whole.
Cross-pass kernel fusion and buffer reuse reduce peak memory consumption and accelerate training.
«PyTorch has over 2000 operations in its full API. For a compiler, that is a huge problem — every backend must be able to handle all of them. PrimTorch IR solves this by reducing the diversity of operations to around 250 primitives».
In the chapter, PrimTorch IR is described as the operation canonicalization stage. Many complex PyTorch operations are decomposed into primitives. For example, the in-place operation x.add_(y) is transformed into a functional add followed by an explicit copy.
A stable PrimTorch IR interface means that to support a new accelerator, developers only need to implement 250 primitives instead of thousands of ATen operations. This radically lowers the barrier to entry for new hardware platforms.
PrimTorch IR reduces over 2000 PyTorch operations to ~250 primitives, simplifying support for new hardware platforms.
In-place operations are converted to functional plus separate copy nodes, simplifying optimization and memory planning.
«TorchInductor (or simply Inductor) is the default backend of the PyTorch compiler. It takes the optimized FX Graph consisting of aten + PrimTorch IR operations and generates high-performance code for NVIDIA GPUs, AMD GPUs, CPUs and other platforms».
In the chapter, TorchInductor is described as the backend that turns the graph into high-performance code. For NVIDIA GPUs, it uses the OpenAI Triton JIT compiler to generate GPU kernels. Triton lowers to NVIDIA PTX via LLVM NVPTX without invoking NVCC to compile the kernels.
TorchInductor includes an autotuner built on Triton's capabilities. On the first compile with the max-autotune mode, Inductor spends extra time benchmarking different kernel variants — different block sizes, tiles, warp counts. The fastest variant is cached and used subsequently.
TorchInductor uses CUDA Graphs to capture sequences of kernels at runtime. This reduces per-iteration launch overhead, which is especially useful for inference. The modes reduce-overhead and max-autotune enable CUDA Graphs.
For stable shapes and latency-critical inference, prefer torch.compile(mode="max-autotune") once recompilations are under control. Keep the capture boundary narrow (per-layer or attention block) to reduce graph invalidation from ragged batching.
TorchInductor is the default PyTorch compiler backend. For NVIDIA GPUs it uses Triton, which lowers to PTX via LLVM NVPTX.
Autotuning and CUDA Graphs provide significant gains for large, critical workloads. For small models, the benefit may be negligible.
«Triton uses the SPMD (Single-Program, Multiple-Data) model, unlike CUDA's SIMT model. Triton kernels operate at the level of thread blocks (CTAs), not individual threads. You don't have to explicitly manage threads or warps — the Triton compiler does it for you».
In the chapter, Triton is described as a Python-oriented DSL for writing GPU kernels. It lowers the barrier to entry for developers unfamiliar with CUDA C++. Under the hood, Triton compiles to PTX via LLVM.
A simple vector addition kernel:
import triton
import triton.language as tl
BLOCK_SIZE = 1024
@triton.jit
def vector_add_kernel(x_ptr, y_ptr, out_ptr, n_elements,
BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
result = x + y
tl.store(out_ptr + offsets, result, mask=mask)Here BLOCK_SIZE is a compile-time constant. tl.arange(0, BLOCK_SIZE) returns a vector of indices, and tl.load and tl.store perform vectorized loads and stores with boundary masking.
To make a Triton kernel visible to torch.compile and avoid graph breaks, register it as a custom operation:
from torch.library import triton_op, wrap_triton
@triton_op("my_triton_lib::vector_add", mutates_args=())
def vector_add(x: Tensor, y: Tensor) -> Tensor:
output = torch.empty_like(x)
n_elements = x.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
wrap_triton(vector_add_kernel)[grid](
x, y, output, n_elements, BLOCK_SIZE=1024
)
return output
c = torch.ops.my_triton_lib.vector_add(a, b)GPU kernel performance is highly sensitive to compile-time parameters. Triton's built-in autotuner automates the search for optimal settings:
@triton.autotune(
configs=[
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64},
num_warps=4, num_stages=2),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64},
num_warps=8, num_stages=2),
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128},
num_warps=8, num_stages=3,
num_consumer_groups=2,
num_buffers_warp_spec=3),
],
key=['M', 'N', 'K'],
)
@triton.jit
def gemm_kernel(...):
...On the first call, Triton JIT-compiles and benchmarks each configuration. The fastest one is cached for a given key derived from input characteristics.
Triton uses the SPMD model, abstracting thread and warp management. Kernels are written in Python and compiled to PTX via LLVM.
Registering custom kernels via torch.library.triton_op() makes them visible to torch.compile. Autotuning automates the search for optimal configurations.
«PyTorch XLA is a separate compilation backend targeting Google Cloud TPU and other accelerators not supported by TorchInductor. XLA maps PyTorch operations to the XLA graph IR and executes them using the target hardware runtime».
In the chapter, PyTorch XLA is described as an alternative backend for non-CUDA accelerators. Activation: torch.compile(..., backend="openxla"). XLA is used for Google Cloud TPU, Meta's internal inference chip (MTIA), and custom AWS Inferentia and Trainium chips via the AWS Neuron SDK.
PyTorch/XLA works on the lazy execution model: operations are recorded into a graph and only executed when the result is strictly needed. This is fundamentally different from the TorchDynamo model, which intercepts bytecode.
The new Dynamo bridge for XLA (beta in 2025) uses the Dynamo guard system for more robust graph capture than lazy tracing.
PyTorch XLA is an alternative backend for TPUs and non-CUDA accelerators. Activated via backend="openxla".
XLA is optimized for static shapes. New shapes cause full recompilation. For TPUs, compilation time can be significant.
«PyTorch DDP works with TorchDynamo, explicitly breaking graphs at all-reduce synchronization points. This is expected behavior: the compiler compiles the reduction for each gradient bucket separately so that communication of one bucket overlaps with the backward computation of the next».
In the chapter, it is emphasized that graph breaks at DDP communication boundaries are necessary to preserve compute-communication overlap. TorchDynamo + DDP works with performance similar to eager DDP and can even surpass it at scale.
PyTorch FSDP is supported in compiled mode with use_original_params=True. Best practice is to wrap each submodule (e.g., each transformer block) in a separate FSDP submodule. Dynamo creates explicit breaks at FSDP submodule boundaries, allowing communication of each shard to overlap with compute.
Graph breaks at DDP and FSDP communication boundaries are necessary to preserve compute-communication overlap.
Wrap submodules in FSDP individually for per-shard compilation and memory savings. Test FSDP with torch.compile on a smaller configuration first.
Fregly's book was written in 2025. Since then, new data has emerged that allows verification of some of the chapter's claims and clarification of figures.
«The PyTorch compiler has come a long way from an experimental feature to a production-ready stack».
Confirmed. TorchInductor and Triton are widely used in production. PyTorch 2.x with torch.compile has become the standard for performance optimization. NVIDIA continues to update Triton to support the latest architectures, including Blackwell.
«At the time of writing, TorchInductor still uses Triton as its primary GPU code generation path».
Confirmed. Triton remains the primary GPU code generation path in TorchInductor. However, competitors have emerged: NVIDIA announced Python-oriented CUDA libraries (cuTile, CuTe DSL), and alternative approaches such as Cypress and LLM agents translating kernels directly to PTX have appeared.
«PyTorch/XLA is a separate compilation backend targeting Google Cloud TPU and other accelerators not supported by TorchInductor».
Confirmed. PyTorch/XLA continues to evolve (version 2.6 in 2025) but remains a niche solution for TPUs and non-CUDA accelerators. Google Cloud continues to support and improve it, but for NVIDIA GPUs the standard is TorchInductor + Triton.
«FlexAttention fused kernels reach up to 85–90% of modern FlashAttention performance».
Partially confirmed. FlexAttention has indeed made significant progress. With the new Flash backend on GB200, it can achieve 1.6–3.2× speedup over the Triton implementation for the forward pass. However, for some patterns and configurations, FlashAttention remains faster.
«Modern compiler frameworks automatically perform some optimizations».
Confirmed. TorchInductor automatically performs operation fusion, autotuning, and uses CUDA Graphs. PyTorch/XLA 2.6 introduced the scan operator and host offloading. The trend toward automation continues.
The chapter describes the PyTorch compiler but omits several aspects that have emerged in real-world practice.
The chapter describes dynamic shape support but does not cover its limitations. CUDA Graphs require static shapes. When dynamic=True is enabled, CUDA Graphs are not used. For variable-length sequences, mode="default" or mode="max-autotune-no-cudagraphs" is recommended. This is an important trade-off between flexibility and performance.
The chapter describes debugging tools but does not emphasize the complexity of diagnosing compiler issues. Graph breaks, recompilations, numerical correctness issues — all require deep understanding of compiler internals. Tools like TORCH_LOGS and torch._dynamo.explain() help, but the learning curve remains high.
The chapter focuses on TorchInductor and XLA but does not consider interaction with TensorRT. In production environments, TensorRT is often used for inference, and torch.compile can integrate with it. This is an important aspect for production deployments that remains outside the chapter's scope.
The chapter describes Triton but does not mention the ecosystem around it. Triton is actively evolving, with new capabilities emerging: TLX (Triton Language Extensions) for low-level control, Blackwell support with tcgen05 instructions, MXFP4/MXFP8. The GPU MODE community actively contributes to Triton's development.
torch.where(), torch.cond(), and refactoring is a key engineering task.max-autotune delivers maximum performance. For variable shapes, default or max-autotune-no-cudagraphs is required.Chapter 14 of Fregly's book is a systematic review of the PyTorch compilation stack. The key idea: the PyTorch compiler has evolved from an experimental feature to a production-ready stack that delivers performance close to hand-written libraries. For an engineer, this means understanding the compiler's internals — from TorchDynamo to TorchInductor — is necessary for effective optimization and problem diagnosis. As we move toward models with hundreds of trillions of parameters and systems like NVL72, the role of the compiler will only grow.
Sources: Chris Fregly, «AI Systems Performance Engineering» (O'Reilly, 2026), Chapter 14 — quotes provided in limited volume for review and critical analysis purposes. PyTorch, TorchInductor, Triton, PyTorch/XLA documentation; MLPerf Training/Inference v5.0; open publications from NVIDIA and Google Cloud.
