Skip to content

Add guard at lowest JAX version that still supports triton kernel calling #2741

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

Open
tdophung wants to merge 2 commits into
base: main
Choose a base branch
from
+15 −0
Open

Add guard at lowest JAX version that still supports triton kernel calling #2741

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ee87daf
add guard at bisected jax version where lower is segfault
tdophung Mar 6, 2026
14666c0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 6, 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
15 changes: 15 additions & 0 deletions transformer_engine/jax/triton_extensions/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,21 @@ def _check_triton_compatibility():
"If you don't need Triton, use transformer_engine.jax.cpp_extensions instead."
) from e

# Minimum jaxlib version required for Triton kernel dispatch to work correctly.
# jaxlib < 0.8.0.dev20250924 segfaults in pxla.py during Triton kernel execution
Copy link
Collaborator

@jberchtold-nvidia jberchtold-nvidia Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work! I think we should expose a function here, something like is_triton_extension_supported() so us or users can check as well. Then in our Triton tests we can add something like the following block at the top:

def test_some_triton_extension(...):
    if not is_triton_extension_supported():
        pytest.skip("... same info message about triton jax compatibility in older versions ...")
       return
    main test code

Similar to how we guard tests based on compute arch

Without this change, I think the tests on old containers would still fail, though the new error message is much more informative then before!

# (bisected: last known segfault = jax-2025-09-23, first known pass = jax-2025-09-24).
_JAXLIB_MIN_VERSION = "0.8.0.dev20250924"
import jaxlib # noqa: E402
Copy link
Contributor

@greptile-apps greptile-apps bot Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import placed mid-module when jax is already available

jax is already imported at line 42, and jax.__version__ always matches jaxlib.__version__ (they are co-released together). The mid-file import jaxlib requires a # noqa: E402 suppression and adds an extra package import for information that's already available via the existing jax import.

Alternatively, if you specifically want jaxlib.__version__ for precision, this import should be hoisted to the top-of-file import block alongside the other jax/jaxlib imports rather than sitting between executable module-level code.

Suggested change
import jaxlib # noqa: E402
import jaxlib # jaxlib version guard below

(move this line to the top of the file alongside the other imports at lines 41–43)

Or, to avoid the extra import entirely:

if version.parse(jax.__version__) < version.parse(_JAXLIB_MIN_VERSION):

since jax.__version__ == jaxlib.__version__ for any properly-installed JAX environment.

Copy link
Collaborator

@phu0ngng phu0ngng Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was also about to comment about jax vs jaxlib.

Besides, we have an existing helper here as well https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/quantize/helper.py#L72-L78, perhaps we should find a way to repurpose this helper, especially since we will need to have jax guard all over places in the future.

On the side note, should we check for a release version, i.e., >= 0.8.0 instead?

Copy link
Contributor

@greptile-apps greptile-apps bot Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great points! Yes, reusing _jax_version_meet_requirement would be much cleaner. You could move it to a shared utility module if multiple files need it.

Regarding dev vs release version:

  • >= 0.8.0.dev20250924 allows: 0.8.0.dev20250924, later dev versions, and all 0.8.x releases
  • >= 0.8.0 only allows official releases (excludes all dev versions)

Since dev versions sort before releases in PEP 440, using >= 0.8.0 would actually block 0.8.0.dev20250924 even though your bisection shows it works.

Recommendation: If most users will use official releases (not nightlies), >= 0.8.0 is cleaner and safer. Dev/nightly users typically expect to handle version quirks. But if you want to explicitly support those working dev versions, keep >= 0.8.0.dev20250924.

For the code:

# At top of file with other imports
from .quantize.helper import _jax_version_meet_requirement

# Later in the file
_JAXLIB_MIN_VERSION = "0.8.0"  # or keep dev version if needed
if not _jax_version_meet_requirement(_JAXLIB_MIN_VERSION):
    raise RuntimeError(f"JAX >= {_JAXLIB_MIN_VERSION} required for Triton kernel support")

This eliminates the mid-file import and reuses existing infrastructure.

Copy link
Collaborator Author

@tdophung tdophung Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@phu0ngng I think greptile has analyzed it pretty thoroughly on checking for > dev_version or >release_version. So the decision is between supporting dev versions after 0924 or not. I think we should support these dev versions, just in case other teams have not updated their pipeline to use a newer/release version. The ugly code to have the dev version instead of release is not that ugly and does not need any maintenance anyways

Copy link
Collaborator Author

@tdophung tdophung Mar 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

regarding the other poitns (jax vs jaxlib) and using the _jax_version_meets_requirement : I agree and will fix


if version.parse(jaxlib.__version__) < version.parse(_JAXLIB_MIN_VERSION):
raise RuntimeError(
f"jaxlib {jaxlib.__version__} is too old for transformer_engine.jax.triton_extensions.\n"
f"Triton kernel dispatch segfaults with jaxlib < {_JAXLIB_MIN_VERSION}.\n"
"Please upgrade: pip install --upgrade jax jaxlib\n"
"Or use a JAX nightly container dated 2025-09-24 or later.\n"
"If you don't need Triton, use transformer_engine.jax.cpp_extensions instead."
)


__all__ = ["triton_call_lowering", "get_triton_info"]

Expand Down