[JAX] Fix get_seqlens_and_offsets() to accept vmapped seg ids and non vmapped seg offsets

[KshitijLakhani]

Description

What is the bug ?

TE provides a convenience function from_segment_ids_and_pos() which allows users to pass only segment ids and the function returns a SequenceDescriptor with internally generated segment pos and passed segment ids.

As mentioned in Issue #2685 , if a user were to vmap a function forward() which i) accepts the q,k,v,segment ids and then ii) calls from_segment_ids_and_pos() followed by iii) a call to DPA(), what happens is that JAX sees the segment ids as vmapped hence an extra leading dimension is added (e.g. 1,2,128) whereas the segment offsets are not given a leading dimension (e.g. 2,128). This results in the FusedAttn primitive impl() assert being triggered due to a shape mismatch between seg ids and seg pos as mentioned in issue #2685

What is the root cause for the bug ?

On debugging, it can be seen that the shape starts differing when the batcher is being traced for the FusedAttn primitive.
segment_ids in the primitive: treated as vmapped inputs hence batched → (1, 2, 128).
segment_pos in the primitive: treated as derived within the function hence not batched → (2, 128).

Fixes #2685

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

There are two possible approaches to solve this:

1. Ensure that the issue is resolved at the source, i.e.. ensure that segment_pos has the same leading batching dims as segment_ids. Add any additional dims in the batcher for the same so that when impl() sees the shape they are the same. Pros: Issue resolved in a "JAX" way and at source. Cons: Increasing mem be expanding seg pos dims.
2. Resolve the issue when impl() is called, i.e. accomodate for mismatched seg id and seg pos dims when generating the seqlens and offsets. Pros: No extra mem needed as no expansion of dims. Cons: Not "truely" solved (at source)

Second approach is chosen here as it more optimized. After this PR merge the end user can vmap wrap the TE API calls without worrying about the batching in TE.
Accomodate for

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 jax L0 L1 L2

[greptile-apps]

Greptile Summary

Fixed shape mismatch bug when using from_segment_ids_and_pos() inside vmapped functions. The issue occurred because JAX added leading batch dimensions to user-provided segment_ids but not to internally-generated segment_pos, causing FusedAttn primitive assertions to fail.

Key changes:

• Modified get_seqlens_and_offsets() in SequenceDescriptor to detect and handle extra leading batch dims on segment_ids
• When segment_ids has more dimensions than segment_pos, the code now flattens extra batch dims, vmaps the seqlens/offsets computation with segment_pos broadcast, then reshapes outputs back
• Replaced strict shape equality assertions with more flexible validation that allows segment_ids to have additional leading dims
• Updated comments in FusedAttn primitive batchers to document that segment_ids/segment_pos may have different batch dimensions

Implementation notes:

• The vmap approach with in_axes=(0, 0, None, None) correctly broadcasts segment_pos across the batch dimension
• JAX will raise clear errors if q and kv segment_ids have incompatible batch sizes, which is appropriate for catching user errors
• Changes only affect THD layout; BSHD path remains unchanged

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation correctly solves the vmap shape mismatch issue with a clean approach: detecting extra batch dims on segment_ids, flattening them, vmapping the computation with broadcasted segment_pos, and reshaping outputs. The validation logic appropriately checks shape compatibility. Edge cases with mismatched batch sizes are caught by JAX's vmap error handling. Changes are scoped to THD layout only, minimizing risk to other code paths.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Modified get_seqlens_and_offsets() to handle vmapped segment_ids with broadcasted segment_pos by flattening extra batch dims, vmapping the computation, and reshaping outputs
transformer_engine/jax/cpp_extensions/attention.py	Updated batcher comments to clarify that segment_ids/segment_pos may have different batch dims and conversion is handled in attention.py

_{Last reviewed commit: 40e4d28}

[greptile-apps]

_{2 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[jberchtold-nvidia]

LGTM, thanks!

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

consider using seg_id_bdim instead of hardcoding 0 for consistency, even though check_valid_batch_dims ensures it's always 0

Suggested change

	updated_batch_dims[seg_pos_idx] = 0
	updated_batch_dims[seg_pos_idx] = seg_id_bdim

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

consider using seg_id_bdim instead of hardcoding 0 for consistency, even though check_valid_batch_dims ensures it's always 0

Suggested change

	updated_batch_dims[seg_pos_idx] = 0
	updated_batch_dims[seg_pos_idx] = seg_id_bdim

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

commented assertion could lead to unclear error if q and kv have mismatched batch sizes. vmap would fail but with a generic JAX error. consider uncommenting or adding a comment explaining why validation isn't needed

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[KshitijLakhani]

/te-ci jax L0 L1 L2

[KshitijLakhani]

CI passes. The only one failure is due to HF requests for the A100 L2 test.
Rerunning passes these