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 issue where pre-compiled Flash Attention wheels with local version labels (e.g. 2.8.3+cu130torch2.11) were incorrectly rejected by the upper-bound check version <= max_version, because PEP 440 defines 2.8.3+local > 2.8.3. The fix applies Version.public to strip the local segment at all three affected call sites (flash-attn, flash-attn-3, jax).

Confidence Score: 5/5

Safe to merge — the fix is correct, minimal in scope, and all remaining call sites not covered have only >= checks that are unaffected by local version labels.

All three changed call sites correctly apply Version.public before range comparisons. The one uncovered call site in _common.py uses only a >= check, which is not broken by local version labels (local > base, so the minimum check still passes). No regressions introduced.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Strips local version labels for flash-attn and flash-attn-3 before range comparison; fix is correct and covers both version slots.
transformer_engine/jax/version_utils.py	Strips local label from JAX version before the >= check; defensive improvement even though >= alone is not broken by local labels.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["importlib.metadata.version(package)"] --> B["Raw version string\ne.g. '2.8.3+cu130torch2.11'"]
    B --> C["PkgVersion(raw_string)"]
    C --> D[".public property\ne.g. '2.8.3'"]
    D --> E["PkgVersion(public_string)"]
    E --> F{"Range check\nmin_version <= version <= max_version"}
    F -->|Pass| G["Package accepted as installed"]
    F -->|Fail| H["Warning / not installed"]

    style B fill:#f9a,stroke:#c33
    style D fill:#afa,stroke:#393
    style E fill:#afa,stroke:#393

Loading

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

[ksivaman]

LGTM

[ksivaman]

/te-ci

[pstjohn]

/te-ci