Implementation:NVIDIA TransformerEngine JAX Triton Utils

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, JAX
Last Updated	2026-02-07 14:00 GMT

Overview

Provides utility functions for integrating Triton kernels into JAX's XLA compilation pipeline, including Triton package detection, kernel compilation, and MLIR lowering.

Description

_detect_triton_package identifies which Triton variant is installed (standard or pytorch-triton) and validates compatibility. _check_triton_compatibility emits appropriate warnings for mixed JAX+PyTorch environments. compile_triton compiles a Triton kernel to a CUDA module with deterministic naming based on content hashing. triton_call_lowering is the main integration point -- it compiles a Triton kernel, creates a custom call operation in MLIR with the compiled PTX/CUBIN, and returns the results. get_triton_dtype maps JAX dtypes to Triton dtypes.

This module is the bridge between Triton's kernel compilation pipeline and JAX's XLA compiler, enabling Triton-authored GPU kernels (primarily for MoE permutation) to be used as JAX custom primitives with proper compilation, caching, and MLIR lowering.

Usage

Use triton_call_lowering when creating JAX custom primitives that are backed by Triton kernels. Use get_triton_info to check Triton availability and version. Use compile_triton when you need to compile Triton kernels outside of the JAX lowering pipeline.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/jax/triton_extensions/utils.py

Lines

1--517

Signature

def get_triton_info() -> Dict[str, Any]:
    """Returns Triton package info (available, version, package name)."""
    ...

def get_triton_dtype(aval) -> "triton.language.dtype":
    """Maps JAX abstract values to Triton dtypes."""
    ...

def compile_triton(
    triton_kernel,
    signature: Dict,
    constants: Dict,
    grid: Tuple,
    num_warps: int,
    num_stages: int,
    **kwargs,
) -> Tuple[str, bytes, int]:
    """Compiles a Triton kernel and returns (name, binary, shared_mem)."""
    ...

def triton_call_lowering(
    ctx,
    *args,
    triton_kernel,
    grid: Tuple,
    num_warps: int,
    num_stages: int,
    signature: Dict,
    constants: Dict,
    out_types: Sequence,
    **kwargs,
):
    """MLIR lowering for Triton kernel calls in JAX."""
    ...

Import

from transformer_engine.jax.triton_extensions.utils import triton_call_lowering, get_triton_info, compile_triton

I/O Contract

Inputs

Name	Type	Required	Description
triton_kernel	Callable	Yes	Triton kernel function to compile
signature	Dict	Yes	Kernel argument type signature
constants	Dict	Yes	Compile-time constant values
grid	Tuple	Yes	Kernel launch grid dimensions
num_warps	int	Yes	Number of warps per block
num_stages	int	Yes	Number of pipeline stages
out_types	Sequence	Yes	Output abstract types for MLIR

Outputs

Name	Type	Description
results	Sequence	MLIR operation results from the custom call

Usage Examples

from transformer_engine.jax.triton_extensions.utils import get_triton_info

# Check if Triton is available
info = get_triton_info()
if info["available"]:
    print(f"Triton {info['version']} available via {info['package']}")

# Triton call lowering is used internally in primitive definitions:
# class MyPrimitive(BasePrimitive):
#     @staticmethod
#     def lowering(ctx, *args, **kwargs):
#         return triton_call_lowering(
#             ctx, *args,
#             triton_kernel=my_triton_kernel,
#             grid=(grid_x,), num_warps=4, num_stages=2,
#             ...
#         )

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements