[JAX] Fix: Use jitted kernels for generating THD (and BSHD) segment pos

[KshitijLakhani]

Description

This is a follow up to PR #2523.
PR #2523 made an API change to the user facing convenience method to create a SequenceDescriptor: from_segment_ids_and_pos.
However, what we've observed is that:

• The newer API introduced in PR [JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API #2523 is less intuitive and more difficult to use
• Most users prefer to generate their own segment_pos anyways

As there are multiple combinations: thd/bshd, cp/reordering, with/without segment_pos which can be difficult to support in an easy manner and because users usually tend to generate their own segment_pos, there isn't much value in continuing to support the ability to generate segment_pos from within from_segment_ids_and_posand the complexities associated with it.

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

This PR does two things:

• Simplifies the API to the way it was before PR [JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API #2523
• Makes the passing of segment_pos mandatory (in addition to the segment_id) to generate the SequenceDescriptor`

Comments point to test code where

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

[KshitijLakhani]

/te-ci L0

[greptile-apps]

Greptile Summary

This PR simplifies SequenceDescriptor.from_segment_ids_and_pos by removing the auto-generation of segment_pos and the is_thd/is_segment_ids_reordered keyword arguments, making segment_pos a required input. The test helpers are updated to always produce explicit segment_pos (now consistently int32) for both BSHD and THD paths, and the conditional CP/reorder logic that gated None vs explicit segment_pos is removed.

Confidence Score: 5/5

Safe to merge; the simplification is sound and the dtype concern from prior review has been resolved.

All previously raised concerns (dtype mismatch, breaking-change label) have been addressed. The only remaining finding is a minor typo in the error message (P2). No logic or correctness issues were found.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Removes complex auto-generation of segment_pos and the is_thd/is_segment_ids_reordered keyword args; now raises ValueError when segment_pos is None (with a TODO to make it a required arg by June 2026). Minor typo in the error message.
tests/jax/test_fused_attn.py	Renames gen_valid → generate_valid_segment_ids_and_pos and generate_random_segment_ids → generate_random_segment_ids_and_pos; explicitly passes segment_pos (now int32) for both BSHD and THD paths; removes the conditional CP/reorder logic that selected None vs explicit segment_pos.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["from_segment_ids_and_pos(segment_ids, segment_pos)"] --> B{segment_pos is None?}
    B -- Yes --> C["raise ValueError\n(segment_pos now required)"]
    B -- No --> D["_expand_to_pair(segment_ids)\n→ q_seg_ids, kv_seg_ids"]
    D --> E["_expand_to_pair(segment_pos)\n→ q_seg_pos, kv_seg_pos"]
    E --> F["return SequenceDescriptor(\n  segment_ids=(q, kv),\n  segment_pos=(q, kv)\n)"]

Loading

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Breaking API change not flagged in checklist

from_segment_ids_and_pos previously accepted segment_pos=None (optional) plus two keyword-only parameters is_thd: bool and is_segment_ids_reordered: bool. All three have been removed, making segment_pos a required positional argument. Any caller passing is_thd=True or is_segment_ids_reordered=False will receive a TypeError at runtime, and any caller relying on the segment_pos=None default path for BSHD or non-load-balanced THD will break silently (missing required argument). The PR checklist marks "Breaking change" as unchecked, which is incorrect for this function, even with its "Experimental" label.

[jberchtold-nvidia]

If you have time, it may be nice to have the following:

def from_segment_ids_and_pos(
    cls,
    ...,
    segment_pos: Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]] = None,
):
    assert segment_pos is not None, "segment_pos is now required as the previous automatic generation of segment_pos did not have the proper context to generate correct a correct segment_pos to account for all load-balancing and context-parallelism strategies the user may be using when calling TransformerEngine's DotProductAttention."

[KshitijLakhani]

Added in 86038f6

[jberchtold-nvidia]

LGTM, thanks!

[KshitijLakhani]

CI passed for all jobs (L0, L1, L2) except a few which were related to CI node issues
Safe to merge

[KshitijLakhani]

NOTE: The title of the PR is not rightly frames (which also reflects in the PR merge commit)
"[JAX] Fix: Use jitted kernels for generating THD (and BSHD) segment pos" was the initial scope, but should have been revised to "[JAX] Fix: Force user to pass segment_pos in from_segment_ids_and_pos"
Latter is more appropriate as that was the final scope of the PR