Strip local version labels from package version checks

[pstjohn]

Pre-compiled Flash Attention wheels (e.g. from
mjun0812/flash-attention-prebuild-wheels) embed build metadata in their package version string (e.g. "2.8.3+cu130torch2.11"). While flash_attn.version returns the clean "2.8.3", TE reads the version via importlib.metadata which returns the full string including the local segment. Under PEP 440, "2.8.3+local" > "2.8.3", causing version range checks like min_version <= version <= max_version to incorrectly reject a compatible installation.

Use Version.public to strip the local label before comparison at all get_pkg_version call sites (flash-attn, flash-attn-3, jax)

[greptile-apps]

Greptile Summary

This PR fixes a version comparison bug where pre-compiled wheels embedding local build metadata (e.g. "2.8.3+cu130torch2.11") caused min_version <= version <= max_version range checks to incorrectly reject compatible flash-attn and JAX installations. The fix applies Version.public to strip the local segment before comparison at all three affected importlib.metadata call sites.

Confidence Score: 5/5

Safe to merge — the fix is correct, minimal, and covers all affected call sites.

All three version-comparison call sites are correctly patched using PEP 440's .public property. No P0/P1 issues found. The only feedback is a P2 style suggestion to centralise the double-wrapping pattern into a helper function.

No files require special attention.

Vulnerabilities

No security concerns identified.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Strips local version labels from flash-attn v2 and v3 metadata before comparison by double-wrapping with PkgVersion().public; fix is correct and complete for both call sites.
transformer_engine/jax/version_utils.py	Applies the same local-label-stripping fix to the JAX version check; semantically correct.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["importlib.metadata.version(pkg)"] --> B["Returns full string\ne.g. '2.8.3+cu130torch2.11'"]
    B --> C["PkgVersion(full_string)"]
    C --> D[".public property\nstrips local label"]
    D --> E["'2.8.3'"]
    E --> F["PkgVersion('2.8.3')"]
    F --> G{"min_version <= version\n<= max_version"}
    G -->|"Pass ✓"| H["Install accepted"]
    G -->|"Fail ✗"| I["Install rejected\n(was wrong without fix)"]

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into pstjohn/strip-p..." | Re-trigger Greptile}

[ksivaman]

LGTM

[ksivaman]

/te-ci