[JAX] Support new JAX triton_kernel_call_ffi for cuda-graph support (NVIDIA#3055)

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>

Author: jberchtold-nvidia

transformer_engine/jax/triton_extensions/utils.py

@@ -44,14 +44,17 @@
 from packaging import version
 
 from jax import core
+from jaxlib.mlir import ir
 import jax
 import jax.numpy as jnp
 
 from ..version_utils import (
     TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION,
+    TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION,
     TRITON_EXTENSION_MIN_JAX_VERSION,
     is_triton_autotuned_alias_safe,
     is_triton_extension_supported,
+    jax_version_meet_requirement,
 )
 
 
@@ -626,12 +629,19 @@ def _normalize_grid(grid_tuple):
     else:
         ffi_operand_output_aliases = None
 
-    # Use JAX FFI lowering with compressed protobuf
-    rule = jax.ffi.ffi_lowering(
-        "triton_kernel_call",  # Custom call target registered in gpu_triton.py
-        api_version=2,
-        backend_config=zlib.compress(call_proto),
-        operand_output_aliases=ffi_operand_output_aliases,
-    )
+    compressed_call_proto = zlib.compress(call_proto)
+    if 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)},
+            operand_output_aliases=ffi_operand_output_aliases,
+        )
+    else:
+        rule = jax.ffi.ffi_lowering(
+            "triton_kernel_call",  # Custom call target registered in gpu_triton.py
+            api_version=2,
+            backend_config=compressed_call_proto,
+            operand_output_aliases=ffi_operand_output_aliases,
+        )
 
     return rule(ctx, *array_args)

transformer_engine/jax/version_utils.py

@@ -25,6 +25,9 @@ def jax_version_meet_requirement(version: str):
 # Minimum JAX version required for Triton kernel dispatch (jaxlib < 0.8.0 segfaults).
 TRITON_EXTENSION_MIN_JAX_VERSION = "0.8.0"
 
+# Minimum JAX version for non-legacy Triton kernel FFI (supports CUDA graph capture).
+TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION = "0.10.1"
+
 # Nightly and stable floors for safe input_output_aliases in TritonAutotunedKernelCall.
 # jaxlib/gpu/triton_kernels.cc had a bug in the autotuning save/restore loop:
 # it iterated over all declared aliases unconditionally, but input_copies only
@@ -76,5 +79,6 @@ def is_triton_extension_supported() -> bool:
     "is_triton_autotuned_alias_safe",
     "is_triton_extension_supported",
     "TRITON_EXTENSION_MIN_JAX_VERSION",
+    "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION",
     "TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION",
 ]