Kernel Selection and Autotuning

Relevant source files

• test/inductor/test_custom_op_autotune.py
• test/inductor/test_max_autotune.py
• test/inductor/test_max_autotune_blackwell.py
• test/inductor/test_template_heuristics_registry.py
• torch/_inductor/autotune_process.py
• torch/_inductor/choices.py
• torch/_inductor/codegen/subgraph.py
• torch/_inductor/heuristics/registry.py
• torch/_inductor/heuristics/template/__init__.py
• torch/_inductor/heuristics/template/aten.py
• torch/_inductor/heuristics/template/contiguous_mm.py
• torch/_inductor/heuristics/template/decompose_k.py
• torch/_inductor/heuristics/template/registry.py
• torch/_inductor/heuristics/template/triton.py
• torch/_inductor/heuristics/triton_codegen/__init__.py
• torch/_inductor/heuristics/triton_codegen/pointwise.py
• torch/_inductor/heuristics/triton_codegen/reduction.py
• torch/_inductor/kernel/bmm.py
• torch/_inductor/kernel/custom_op.py
• torch/_inductor/kernel/mm.py
• torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja
• torch/_inductor/kernel/templates/triton_bmm.py.jinja
• torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja
• torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja
• torch/_inductor/kernel/templates/triton_mm.py.jinja
• torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja
• torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja
• torch/_inductor/kernel/templates/triton_persistent_mm.py.jinja
• torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja
• torch/_inductor/select_algorithm.py
• torch/_inductor/virtualized.py

Purpose and Scope

This page documents the kernel selection and autotuning framework within PyTorch's TorchInductor backend. The process enables automatic selection and optimization of kernel implementations for fused operations during lowering and scheduling, choosing among various backends such as Triton, CUTLASS, ATen, and others.

Specifically, this page covers:

• The algorithm selection architecture based on ChoiceCaller abstractions.
• Configuration space generation and template heuristics guiding candidate selection.
• The autotuning flow including benchmarking in isolated subprocesses.
• Persistent caching of autotuning results to avoid repeated tuning.
• Subgraph and custom operator autotuning.
• How backend selection among Triton, CUTLASS, ATen, and CPU templates is integrated.

Overview: From Scheduled Nodes to Optimized Kernels

At a high level, the kernel selection and autotuning flow starts from scheduled nodes (produced by the scheduler), and picks the optimal kernel implementation. Multiple candidate kernel templates are generated and benchmarked to select the best-performing variant. This process is cached persistently for future reuse by shape.

The flow is:

This process optimizes kernel choices such as GEMM implementations, pointwise operators, and fused subgraphs.

Sources:
torch/_inductor/select_algorithm.py45-70
torch/_inductor/autotune_process.py31-180
torch/_inductor/template_heuristics/triton.py104-156
torch/_inductor/runtime/triton_heuristics.py40-60
torch/_inductor/codecache.py40-56

---

Template System and ChoiceCaller

Template Types and Backends

Inductor manages multiple kernel template abstractions, each representing a backend or kernel generation strategy. The key types are:

Template Type	Description	Primary Use	Source & Lines
TritonTemplate	Jinja2 templates for Triton GPU kernels	GPU matrix multiplications, reductions	torch/_inductor/select_algorithm.py:68-74 torch/_inductor/kernel/mm.py:87-110
ExternKernelChoice	Wrappers for ATen/cuBLAS/cuDNN external kernels	Stable fallback implementations	torch/_inductor/select_algorithm.py:72-78 torch/_inductor/kernel/mm.py:140-155
CUTLASS2xGemmTemplate	CUTLASS 2.x library based GEMM templates for CUDA	CUDA GEMM kernels	torch/_inductor/kernel/mm.py:24
CKGemmTemplate	AMD Composable Kernel (CK) templates for ROCm	ROCm GEMM kernels	torch/_inductor/kernel/mm.py:26
CppGemmTemplate	C++ based gemm templates for CPU	CPU GEMM kernels	torch/_inductor/kernel/mm.py:15
SubgraphTemplate	Autotuned arbitrary FX GraphModules	Complex fused subgraphs	torch/_inductor/codegen/subgraph.py:60-70

These template classes provide a uniform interface for kernel rendering, launching, and benchmarking.

ChoiceCaller Hierarchy

Kernel selection is built atop the abstract base class ChoiceCaller defined in the IR system. It defines the interface to:

• Benchmark candidate kernels.
• Obtain IR nodes to emit a kernel invocation.

Below is the ChoiceCaller subclass hierarchy used for various template types:

• TritonTemplateCaller benchmarks Triton-generated kernels.
• ExternKernelCaller benchmarks ATen and other external kernels as a baseline.
• SubgraphChoiceCaller benchmarks arbitrary FX subgraphs compiled into modules.

Sources:
torch/_inductor/select_algorithm.py56-80
torch/_inductor/codegen/subgraph.py60-75
torch/_inductor/ir.py40-42

---

Configuration Generation and Heuristics

Configuration Spaces for Kernel Templates

Kernel templates are parameterized by configuration classes that specify properties such as block sizes, number of warps, and pipeline stages. These configurations define the search space for autotuning.

Typical config dataclasses include:

Config Class	Description	Additional Fields	Source & Lines
BaseConfig / GemmConfig	Base GEMM tiling and scheduling parameters	block_m, block_n, block_k, num_stages, num_warps	torch/_inductor/heuristics/template/triton.py:114-130
ROCmGemmConfig	AMD ROCm-specific GEMM config extension	matrix_instr_nonkdim, waves_per_eu, kpack	torch/_inductor/heuristics/template/triton.py:236-247
BlackwellGPUGemmConfig	NVIDIA Blackwell GPU GEMM with persistent TMA specifics	epilogue_subtile, warp_specialize, flatten, and meta-WS knobs	torch/_inductor/heuristics/template/triton.py:155-175

These configs define optimization parameters that heuristics prune or prioritize.

Heuristic Classes

Heuristic classes generate candidate configurations for a given template, device type, and operator:

• CUDA: CUDAConfigHeuristic
• ROCm: ROCmConfigHeuristic
• XPU: XPUConfigHeuristic
• CPU: CPUConfigHeuristic

Heuristics encapsulate logic on allowed parameter combinations, device properties, and degraded fallback choices.

The registry system allows registering and retrieving heuristics by (template_name, device_type, op_name) keys.

Sources:
torch/_inductor/template_heuristics/triton.py104-265
torch/_inductor/choices.py131-147
torch/_inductor/heuristics/registry.py40-140

---

Autotuning Process

CachingAutotuner and autotune_select_algorithm

The CachingAutotuner class manages autotuning lifecycle:

• Checks PersistentCache for previously benchmarked best configs.
• If cache miss, launches precompilation and benchmarking of candidates.
• Selects the fastest choice and stores the result persistently.

The key entrypoint is autotune_select_algorithm which:

1. Generates choice candidates from template heuristics.
2. Normalizes inputs for benchmarking.
3. Checks persistent cache keyed by code hash and shapes.
4. If cached, returns best cached choice.
5. If not, asynchronously benchmarks candidates with AutotuneProcessPool.
6. Returns the winning choice for codegen.

The autotuning is essential to achieve high performance by exploring diverse kernel configurations and backend implementations.

Multi-Process Autotuning Infrastructure

To isolate benchmarking overhead and ensure stable measurements, tuning is performed in separate subprocesses managed by:

• TuningProcess: A single benchmark subprocess managing jobs via IPC pipes.
• AutotuneProcessPool: A pool of TuningProcess instances distributing benchmarking requests asynchronously.
• TritonBenchmarkRequest / ExternKernelBenchmarkRequest / SubgraphGPUBenchmarkRequest: Typed requests encapsulating kernel and input details to benchmark.

The pool supports environment variable isolation (e.g., CUDA_VISIBLE_DEVICES) and cache synchronization.

Sources:
torch/_inductor/select_algorithm.py48-70
torch/_inductor/autotune_process.py1-180
torch/_inductor/runtime/triton_heuristics.py40-57

---

Matrix Multiplication (MM) Selection

Inductor's MM kernel selection covers different backends and heuristics. It attempts candidate kernels from:

• Triton templates
• External ATen/cuBLAS kernels
• CUTLASS templates (NVIDIA CUDA)
• CK templates (AMD ROCm)
• CPU-specific templates (e.g. C++ GEMM)

The main dispatching and autotune calls occur in torch/_inductor/kernel/mm.py with utilities like autotune_select_algorithm.

This allows leveraging hardware-specific backends with autotuning to fine-tune block sizes and launch parameters.

Sources:
torch/_inductor/kernel/mm.py84-165
torch/_inductor/select_algorithm.py48-70

---

Subgraph and Custom Op Autotuning

Inductor supports autotuning for:

• Arbitrary fused subgraphs: Using SubgraphChoiceCaller which compiles FX GraphModules and compares runtime performance [codegen/subgraph.py].
• Custom operators: Using CustomOpConfig specifying multiple decompositions with input-size based dispatch via RangeBounds, allowing multiple implementations for the same logical op [kernel/custom_op.py].

This flexibility widens the autotuning capability beyond fixed kernels to arbitrary compositions and user-defined ops.

---

Pad MM Optimization

An auxiliary approach used in Inductor is padding matrix dimensions to multiples suitable for efficient kernels (e.g., aligning to 8 for float16). This enables optimized kernels from libraries or reduces the number of special cases.

Key points:

• Padding is statically decided when possible, considering symbolic shapes.
• Uses heuristics and checks device/dtype support.
• Padded versions are only used if performance gains are expected, validated by autotuning or heuristics.

Sources:
torch/_inductor/fx_passes/pad_mm.py100-200

---

Key Classes and Functions Reference

Symbol	Location	Description
autotune_select_algorithm	torch/_inductor/select_algorithm.py:49	Main kernel selection and autotuning routine
TritonTemplate	torch/_inductor/select_algorithm.py:68	Triton kernel template wrapper
ChoiceCaller	torch/_inductor/ir.py:40	Abstract base class for kernel choice benchmarking
CachingAutotuner	torch/_inductor/runtime/triton_heuristics.py:43	Caches best kernel configs, manages tuning
AutotuneProcessPool	torch/_inductor/autotune_process.py:31	Manages parallel autotuning subprocesses
BenchmarkTensors	torch/_inductor/select_algorithm.py:132	Holds input/output tensors used during benchmarking
SubgraphChoiceCaller	torch/_inductor/codegen/subgraph.py:60	Autotunes complex subgraph kernels

---

This completes the full technical overview of kernel selection and autotuning in TorchInductor.

Sources:
torch/_inductor/select_algorithm.py1-130
torch/_inductor/autotune_process.py1-180
torch/_inductor/template_heuristics/triton.py100-265
torch/_inductor/kernel/mm.py80-160
torch/_inductor/codegen/subgraph.py30-130
torch/_inductor/kernel/custom_op.py20-200