How to Convert PyTorch Dispatch Macros to V2 Format with at-dispatch-v2
When maintaining PyTorch ATen kernels, converting legacy AT_DISPATCH macros to the AT_DISPATCH_V2 format involves multiple transformation rules: argument reordering, lambda wrapping, and type group expansion. This manual process is error-prone.
This article explains how to safely convert AT_DISPATCH macros to V2 format using the at-dispatch-v2 skill.
What This Skill Does
at-dispatch-v2 converts legacy PyTorch AT_DISPATCH macros to AT_DISPATCH_V2 format:
- Targets AT_DISPATCH_ALL_TYPES_AND*, AT_DISPATCH_FLOATING_TYPES_AND*, and similar macros
- Reorders arguments (scalar_type, name moved to front)
- Wraps lambdas with AT_WRAP()
- Expands type groups with AT_EXPAND()
- Separates individual types (kHalf, kBFloat16, etc.)
Designed for PyTorch core developers maintaining ATen and CUDA kernels.
Installation
Prerequisites
- Claude Code installed
- Access to PyTorch repository source code
Install Command
claude mcp add github.com/pytorch/pytorch/tree/main/.claude/skills/at-dispatch-v2
Usage
Basic Usage
Specify the file to convert:
Convert AT_DISPATCH macros in this file to V2
Transformation Rules
Old format:
AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool,
dtype, "kernel_name", [&]() {
kernel<scalar_t>();
});
New format:
AT_DISPATCH_V2(dtype, "kernel_name", AT_WRAP([&]() {
kernel<scalar_t>();
}), AT_EXPAND(AT_ALL_TYPES), kBFloat16, kHalf, kBool);
Type Group Mapping
| Old Macro Base | V2 Type Group |
|---|---|
| ALL_TYPES | AT_EXPAND(AT_ALL_TYPES) |
| FLOATING_TYPES | AT_EXPAND(AT_FLOATING_TYPES) |
| INTEGRAL_TYPES | AT_EXPAND(AT_INTEGRAL_TYPES) |
| COMPLEX_TYPES | AT_EXPAND(AT_COMPLEX_TYPES) |
Important Considerations
AT_WRAP() Is Mandatory
Lambdas containing commas will break macro parsing without AT_WRAP(). It cannot be omitted.
No AT_EXPAND() for Individual Types
Individual types like kHalf and kBFloat16 are listed directly. Do not write AT_EXPAND(kHalf).
Keep Dispatch.h
Do not remove #include <ATen/Dispatch.h> when adding Dispatch_v2.h — other code may depend on it.
Summary
at-dispatch-v2 enables safe conversion of legacy AT_DISPATCH macros to V2 format following argument reordering, lambda wrapping, and type group expansion rules. V2 format is more extensible and eliminates the need for AND2/AND3 macro variants.
For full specifications, visit the skill detail page.