Implementation:NVIDIA TransformerEngine PyTorch Extensions Header
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, PyTorch, Distributed |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Master C++ header declaring all extension function signatures exposed to Python through pybind11, organized by functional category.
Description
Organizes declarations into clearly separated sections: Router fusion (fused_topk, fused_moe_aux_loss), Permutation (moe_permute/unpermute), Attention (fused_attn_fwd/bwd, KV cache ops, format converters), GEMM (gemm, te_atomic_gemm, te_general_grouped_gemm), Activation functions, Normalization (layernorm, rmsnorm fwd/bwd), Softmax variants (scaled, masked, upper-triangular), Bias operations, Cast/Quantize/Dequantize, Transpose, Padding, Dropout, Communication-computation overlap (CommOverlapHelper/CommOverlap/CommOverlapP2P classes), Recipe management, RoPE, Multi-tensor operations, and miscellaneous utilities.
Usage
Serves as the single API surface for all PyTorch C++ extensions. Any new CUDA-accelerated operation must be declared here to be accessible from Python.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/pytorch/csrc/extensions.h- Lines
- 1--581
Signature
namespace transformer_engine::pytorch {
// Attention
py::object fused_attn_fwd(...);
std::vector<at::Tensor> fused_attn_bwd(...);
// GEMM
py::object gemm(at::Tensor A, at::Tensor B, ...);
py::object te_atomic_gemm(...);
py::object te_general_grouped_gemm(...);
// Normalization
py::object layernorm_fwd(at::Tensor input, ...);
at::Tensor layernorm_bwd(at::Tensor dz, ...);
py::object rmsnorm_fwd(at::Tensor input, ...);
// Communication overlap
class CommOverlapHelper { ... };
class CommOverlap { ... };
class CommOverlapP2P { ... };
// Activation, Cast, RoPE, Softmax, etc.
// ... (hundreds of function declarations)
} // namespace transformer_engine::pytorch
Import
#include "extensions.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| N/A | N/A | N/A | Header file -- declares function signatures used by all extension implementations |
Outputs
| Name | Type | Description |
|---|---|---|
| N/A | N/A | Provides function declarations for the entire C++ extension API surface |
Usage Examples
#include "extensions.h"
// Implement a new extension that uses existing TE functions
py::object my_custom_op(at::Tensor input) {
auto te_input = makeTransformerEngineTensor(input);
// ... call NVTE kernels ...
}