[JAX] Add warning if using BSHD and max_segments_per_seq > 1

[jberchtold-nvidia]

Description

Adds a small warning if the user tries to use BSHD with max_segments_per_seq > 1

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Adds a warning if the user tries to use BSHD with max_segments_per_seq > 1
• Adds a new test to validate this warning is shown correctly

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[jberchtold-nvidia]

/te-ci jax

[greptile-apps]

Greptile Summary

This PR adds a UserWarning to fused_attn in transformer_engine/jax/attention.py when the caller sets max_segments_per_seq > 1 with a non-THD (BSHD) layout, since sequence packing is only meaningful for THD formats. The change is small, self-contained, and the warning logic — guarded by not qkv_layout.is_thd() — correctly covers all BSHD variants (BS3HD, BSHD_BS2HD, BSHD_BSHD_BSHD). The two issues flagged in earlier review rounds (stacklevel omission and the duplicate if causing IndentationError) are both resolved in the current state.

Key observations:

• Warning is correctly placed in the non-legacy code path (after the early-return for deprecated mask usage), so it won't fire spuriously for callers using the old mask-based API.
• stacklevel=2 is present, directing the warning to the caller's site rather than library internals.
• The PR description and checklist claim a test class (TestMaxSegmentsPerSeqWarning) was added to tests/jax/test_fused_attn.py, but it was removed in a subsequent commit (6d7ad99) with no stated reason, leaving the new warning path without automated test coverage.

Confidence Score: 5/5

Safe to merge — the one-line logic change is correct and non-breaking; all remaining feedback is P2 quality/documentation concerns.

All identified issues are P2 (missing test coverage, inaccurate PR checklist). No runtime defects, security issues, or logic errors remain. Previous P0/P1 issues (duplicate if / missing stacklevel) are resolved.

No files require special attention for merge safety, but tests/jax/test_fused_attn.py should ideally have the removed test reinstated.

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Adds a correctly placed UserWarning (with stacklevel=2) when max_segments_per_seq > 1 is used with a non-THD layout; prior review issues (duplicate if, missing stacklevel) are resolved.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["fused_attn(qkv, ..., max_segments_per_seq, qkv_layout)"]
    B{"sequence_descriptor is None\nor isinstance(jnp.ndarray)?"}
    C["warnings.warn(DeprecationWarning)\n+ raise ValueError if max_segments_per_seq != 1"]
    D["_legacy_fused_attn(...)  return"]
    E{"max_segments_per_seq > 1\nAND NOT qkv_layout.is_thd()?"}
    F["warnings.warn(UserWarning, stacklevel=2)\n'max_segments_per_seq only applies to THD'"]
    G["_fused_attn(...) → output  return"]

    A --> B
    B -- yes --> C --> D
    B -- no --> E
    E -- yes --> F --> G
    E -- no --> G

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into jberchtold/te-m..." | Re-trigger Greptile}

[jberchtold-nvidia]

/te-ci jax

[KshitijLakhani]

Thanks for testing this warning out, Jeremy !
I think negative testing is crucial in such cases..
However, I'm thinking that merging a test case for just checking a warning might be overkill.
We've got quite a few warnings in TE but we do not test them (from at least what I'm aware of) so I'd suggest we drop the test and just keep the warning change in fused attn.
Local negative testing to ensure that the warning gets triggered, followed by a passing CI for fused attn tests should be enough, I think . Thoughts ?

[jberchtold-nvidia]

That sounds good to me, I'll remove the new test

[jberchtold-nvidia]

Updated

[jberchtold-nvidia]

/te-ci jax

[KshitijLakhani]

LGTM !
Good to merge post successful CI runs for fused attn
Thanks !