Skip to content

[JAX] Fallback to old triton ffi for autotuned kernels #3077

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jberchtold-nvidia merged 2 commits into from
Jun 3, 2026
Merged

[JAX] Fallback to old triton ffi for autotuned kernels #3077

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3dfd40e
Disable new triton ffi for autotuned kernels
jberchtold-nvidia Jun 2, 2026
f375ab1
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion transformer_engine/jax/triton_extensions/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -485,6 +485,7 @@ def _normalize_grid(grid_tuple):
# NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 to raise an error instead, prompting the
# user to upgrade JAX for improved performance.
is_autotuned = isinstance(kernel_fn, autotuner.Autotuner)
used_autotuned_launch = False
if is_autotuned and not is_triton_autotuned_alias_safe():
val = os.environ.get("NVTE_JAX_ENFORCE_TRITON_AUTOTUNING", "0")
try:
Expand Down Expand Up @@ -574,6 +575,7 @@ def _normalize_grid(grid_tuple):
kernel_calls,
input_output_aliases_with_sizes,
)
used_autotuned_launch = True

else:
# Regular kernel: compile single config.
Expand Down Expand Up @@ -630,7 +632,9 @@ def _normalize_grid(grid_tuple):
ffi_operand_output_aliases = None

compressed_call_proto = zlib.compress(call_proto)
if jax_version_meet_requirement(TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION):
if not used_autotuned_launch and jax_version_meet_requirement(
TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION
):
rule = jax.ffi.ffi_lowering(
"triton_kernel_call_ffi",
backend_config={"opaque": ir.StringAttr.get(compressed_call_proto)},
Expand Down
Loading