[JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API by KshitijLakhani · Pull Request #2523 · NVIDIA/TransformerEngine

KshitijLakhani · 2025-12-16T23:12:00Z

Description

SequenceDescriptor's from_segment_ids_and_pos() accepts the segment_ids and an optional segment_pos as input. This class is supposed to serve as a convenience method to do two things:

Stuff the segment_ids and segment_pos in a SequenceDescriptor object for TE to use downstream
If the segment_pos is not passed, then calculate/extrapolate it

In it's current form, the second functionality gives incorrect results for THD + non-reordered and THD + reordered cases as it merely uses an arange to calculate the segment_pos naively. This could result in incorrect masking for these cases.

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 makes few changes:

Passed two new args to this function from_segment_ids_and_pos() : is_thd and is_segment_ids_reordered- the only cases that this function can currently guarantee to support is BSHD with and without load balancing and, THD without load balancing.
BSHD with load balancing is supported natively because the segment_ids and segment_pos are not reordered before passing onto from_segment_ids_and_pos(). However, if the segment_pos are reordered and passed to from_segment_ids_and_pos() it will assert
If THD + reordered use cases calls the function from_segment_ids_and_pos(), it will assert
The fused attn tests were modified to account for these two new args
- For THD fused attn non-CP tests, segment_pos=None is passed so as to exercise the newly added THD path in from_segment_ids_and_pos()
- For THD fused attn CP tests, segment_pos is explicitly passed (as before, not a new change)
- For BSHD fused attn CP tests, segment_pos=None is passed so as to exercise the default BSHD path to generate segment_pos (as before, not a new change)

Impact on user of the API:

These two new args, is_thd and is_segment_ids_reordered are not Optional and hence they will cause a TypeError for current users of this API - a breaking change. However, this is needed to ensure correct usage of this API
The user is now expected to let the API know whether this is a THD or BSHD layout and whether the segment_ids are reordered or not. It is expected that the segment_ids passed will be reordered only for THD load balancing. For all other cases the segment_ids should not be reordered

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 · 2025-12-16T23:14:13Z

/te-ci jax L0 L1

KshitijLakhani · 2025-12-17T02:16:02Z

/te-ci jax L0 L1

greptile-apps · 2025-12-17T02:18:32Z

Greptile Summary

Fixed incorrect segment position calculation in SequenceDescriptor.from_segment_ids_and_pos() by adding proper THD support. Previously used naive arange that resulted in incorrect masking for THD cases.

Key changes:

Added two new required parameters: is_thd and is_segment_ids_reordered to from_segment_ids_and_pos()
Implemented proper THD segment position generation that detects segment boundaries and calculates positions relative to each segment start
Added assertions to prevent unsupported combinations (THD + reordered when segment_pos=None)
Updated all test call sites to pass the new required parameters

Breaking change: Existing users must now provide is_thd and is_segment_ids_reordered arguments, causing TypeError if not updated.

Confidence Score: 1/5

This PR cannot be merged - contains critical logic bug that will cause runtime failures
Lines 833-844 in attention.py contain contradictory assertions that will always fail when is_segment_ids_reordered=True. The first assertion requires not is_thd (line 833) while the second requires is_thd (line 839), making it impossible to satisfy both. This bug wasn't caught by tests because they always pass explicit segment_pos when is_segment_ids_reordered=True, avoiding the buggy code path.
Pay immediate attention to transformer_engine/jax/attention.py lines 833-844 - the contradictory assertions must be fixed before merge

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Critical logic bug: contradictory assertions (lines 833-844) will always fail when `is_segment_ids_reordered=True`. Added THD segment position generation logic.
tests/jax/test_fused_attn.py	Updated tests to pass new required params `is_thd` and `is_segment_ids_reordered`. Minor comment inaccuracy on line 671.

Sequence Diagram

sequenceDiagram
    participant User
    participant SequenceDescriptor
    participant generate_default_pos
    
    User->>SequenceDescriptor: from_segment_ids_and_pos(segment_ids, segment_pos=None, is_thd, is_segment_ids_reordered)
    
    alt segment_pos is None
        SequenceDescriptor->>SequenceDescriptor: Check if is_segment_ids_reordered
        
        alt is_segment_ids_reordered = True
            SequenceDescriptor->>SequenceDescriptor: assert not is_thd (line 833)
            SequenceDescriptor->>SequenceDescriptor: assert is_thd (line 839)
            Note right of SequenceDescriptor: BUG: Contradictory assertions!<br/>Will always fail
        end
        
        SequenceDescriptor->>generate_default_pos: Call for q_seg_ids
        
        alt is_thd = True
            generate_default_pos->>generate_default_pos: Find segment boundaries
            generate_default_pos->>generate_default_pos: Calculate segment_start_offsets
            generate_default_pos->>generate_default_pos: Compute positions relative to segment start
            generate_default_pos->>generate_default_pos: Apply padding mask
            generate_default_pos-->>SequenceDescriptor: Return THD segment_pos
        else is_thd = False (BSHD)
            generate_default_pos->>generate_default_pos: Use simple arange
            generate_default_pos-->>SequenceDescriptor: Return BSHD segment_pos
        end
        
        SequenceDescriptor->>generate_default_pos: Call for kv_seg_ids
        generate_default_pos-->>SequenceDescriptor: Return kv segment_pos
        
    else segment_pos provided
        SequenceDescriptor->>SequenceDescriptor: Expand to pair
    end
    
    SequenceDescriptor-->>User: Return SequenceDescriptor with segment_ids and segment_pos

greptile-apps

Additional Comments (1)

transformer_engine/jax/attention.py, line 841-842 (link)

logic: q_seg_ids and kv_seg_ids are used here but not defined until line 847. This will cause a NameError at runtime when segment_pos is None.

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

KshitijLakhani · 2025-12-17T02:50:37Z

/te-ci jax L0 L1

greptile-apps

Additional Comments (1)

transformer_engine/jax/attention.py, line 829-832 (link)

style: f-strings in assert messages won't be evaluated until assertion fails

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

…es and load balanced cases in from_segment_ids_and_pos. Enforce passing of segment_pos for THD cases and lod balanced cases Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

greptile-apps

Additional Comments (1)

transformer_engine/jax/attention.py, line 833-837 (link)

logic: warning is shown even for the valid case (BSHD without load balancing)

this will cause users to see warnings every time they correctly use segment_pos=None with BSHD layout

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

2. Add support for THD in from_segment_ids_and_pos() 3. Assert if load balanced segment_ids is passed to generate a segment_pos Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

for more information, see https://pre-commit.ci

greptile-apps

Additional Comments (3)

transformer_engine/jax/attention.py, line 880-881 (link)

syntax: segment_ids is undefined, should be seg_ids (the parameter name on line 836)
tests/jax/test_fused_attn.py, line 680-681 (link)

syntax: is_thd and is_load_balanced are keyword-only parameters (note the * on line 800 in attention.py), so they must be passed as keyword arguments
tests/jax/test_fused_attn.py, line 709-710 (link)

syntax: is_thd and is_load_balanced are keyword-only parameters (note the * on line 800 in attention.py), so they must be passed as keyword arguments

_{2 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

greptile-apps

Additional Comments (3)

tests/jax/test_fused_attn.py, line 706-711 (link)

logic: passing segment_pos=None with is_load_balanced=True will trigger assertion error on attention.py:830

the BSHD + load balanced + SegmentIDs combination should either:
1. pass explicit segment_pos (like THD does on line 677-678), or
2. skip this format when load balanced
currently only works because default mesh_shape=(1,1,1) means cp_size=1
transformer_engine/jax/attention.py, line 794 (link)

style: misleading TODO - THD layout (non load balanced) IS now supported by the implementation below (lines 837-878)
transformer_engine/jax/attention.py, line 839-840 (link)

style: assumption not enforced - if sequence starts with padding (seg_ids[0]=0), first_is_segment=True will incorrectly mark it as segment start

consider asserting seg_ids[:, 0] != 0 or handle leading padding explicitly

_{2 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

…actually reordered segment_ids and segment_pos Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

for more information, see https://pre-commit.ci

…ncing Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

greptile-apps

Additional Comments (2)

transformer_engine/jax/attention.py, line 833-844 (link)

logic: contradictory assertions will always fail when is_segment_ids_reordered=True

line 833 asserts not is_thd (meaning is_thd must be False)
line 839 asserts is_thd (meaning is_thd must be True)

both assertions cannot be satisfied simultaneously, causing the function to always fail when segment_pos=None and is_segment_ids_reordered=True
transformer_engine/jax/attention.py, line 850-851 (link)

style: assumption that first token is always a valid segment may be incorrect if THD format allows leading padding

if seg_ids[0, 0] == 0 (first token is padding), this assumption causes incorrect segment_pos calculation. verify THD format guarantees no leading padding or add validation

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

for more information, see https://pre-commit.ci Fix lint errors Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

KshitijLakhani · 2025-12-24T01:47:31Z

/te-ci jax L0 L1

greptile-apps

Additional Comments (3)

transformer_engine/jax/attention.py, line 833-844 (link)

logic: contradictory assertions on lines 833 and 839 - both will execute when is_segment_ids_reordered=True, but line 833 requires is_thd=False while line 839 requires is_thd=True

this means when is_segment_ids_reordered=True:
- if is_thd=True: line 833 fails with "not supported for load balanced reordered (Striped) THD"
- if is_thd=False: line 839 fails with "not supported for load balanced reordered (Dual Chunk) BSHD"
based on line 822 comments and PR description, only THD with load balancing should set is_segment_ids_reordered=True, so line 839's assertion should be removed (it contradicts the intended behavior)
transformer_engine/jax/attention.py, line 856 (link)

style: potential edge case: (seg_ids[..., 1:] != 0) condition may not handle all segment transitions correctly

consider sequence: [1, 1, 0, 2, 2] where 0 is padding in the middle
- transition from 1→0 at index 2: (1 != 0) & (0 != 0) = False (correctly not marked as segment start)
- transition from 0→2 at index 3: (0 != 2) & (2 != 0) = True (marked as segment start)
however, the assumption on line 850 is that "the first token belongs to a segment and is not a padded token", which conflicts with the possibility of having padding in the middle. verify whether THD format allows mid-sequence padding, and if so, this logic needs adjustment. does THD format allow padding (segment_id=0) in the middle of a sequence, or is padding only at the end?
tests/jax/test_fused_attn.py, line 687-689 (link)

style: simplify boolean expression

_{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!}

_{2 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

KshitijLakhani · 2025-12-24T18:01:52Z

/te-ci jax L0 L1

jberchtold-nvidia

LGTM pending CI, thanks!

…eg-ids-only

ksivaman · 2025-12-31T07:51:14Z

/te-ci jax L0 L1

greptile-apps

Additional Comments (2)

transformer_engine/jax/attention.py, line 833-844 (link)

logic: contradictory assertions will always fail when is_segment_ids_reordered=True

Line 833 asserts not is_thd (THD must be False), but line 839 asserts is_thd (THD must be True). When is_segment_ids_reordered=True, the first assertion will pass only if is_thd=False, but then the second assertion will always fail because it requires is_thd=True.

Based on the comments and documentation (line 822: "Only THD with load balancing is expected to have this flag set to True"), the correct logic should be:
tests/jax/test_fused_attn.py, line 671-672 (link)

style: comment doesn't match implementation logic

Comment says "if no CP and load balancing", but the code on line 683 does the opposite - it explicitly passes segment_pos when cp_size > 1 and self.cp_load_balanced (i.e., WITH CP and load balancing), and passes None otherwise.

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

KshitijLakhani · 2026-01-05T22:37:35Z

@ksivaman this PR has a breaking change (FYI for TE release notes)

KshitijLakhani · 2026-01-05T22:39:11Z

@mgoldfarb-nvidia @huanghua1994 @mingxu1067 - even though this PR has been merged, it would be great to get your feedback on this, in case you spot something that needs to be reverted. Thanks !

TransformerEngine v2.12 (NVIDIA/TransformerEngine#2523) made `is_thd` and `is_segment_ids_reordered` required keyword arguments on `SequenceDescriptor.from_segment_ids_and_pos()` to fix incorrect segment position calculation for THD layouts. Since the packing branch in `cudnn_flash_attention` uses `qkv_layout="THD_THD_THD"` with standard (non-reordered) segment IDs, the correct values are `is_thd=True, is_segment_ids_reordered=False`. Without this fix, any configuration using `attention="cudnn_flash_te"` with `packing=True` and real data (`dataset_type != "synthetic"`) fails with: TypeError: SequenceDescriptor.from_segment_ids_and_pos() missing 2 required keyword-only arguments: 'is_thd' and 'is_segment_ids_reordered'

The existing test_gpu_packed_attention uses dataset_type=synthetic, which bypasses SequenceDescriptor.from_segment_ids_and_pos() entirely (takes the elif branch in cudnn_flash_attention). This meant the TE v2.12 API breakage went undetected. Add test_gpu_packed_attention_hf that uses HF parquet data with packing=True + attention=cudnn_flash_te, exercising the actual SequenceDescriptor codepath. This serves as a regression test for NVIDIA/TransformerEngine#2523.

KshitijLakhani self-assigned this Dec 16, 2025

KshitijLakhani added attention jax bug Something isn't working labels Dec 16, 2025

KshitijLakhani marked this pull request as ready for review December 17, 2025 02:16

greptile-apps Bot reviewed Dec 17, 2025

View reviewed changes

KshitijLakhani and others added 7 commits December 22, 2025 11:30

Fix incorrect calculation of segment pos from segment ids for thd cas…

2d2e33a

…es and load balanced cases in from_segment_ids_and_pos. Enforce passing of segment_pos for THD cases and lod balanced cases Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

65e6b4b

for more information, see https://pre-commit.ci

Correct the assert condition

9857577

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Modify fused attn tests to pass new args to from_segment_ids_and_pos()

b20ac22

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

03398a4

for more information, see https://pre-commit.ci

Calculate seg ids before pos

0a47eb6

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

217ea58

for more information, see https://pre-commit.ci

KshitijLakhani force-pushed the klakhani/fix/incorrect-sequence-descr-from-seg-ids-only branch from 499f4de to 217ea58 Compare December 22, 2025 19:31

greptile-apps Bot reviewed Dec 22, 2025

View reviewed changes

Kshitij Janardan Lakhani and others added 2 commits December 22, 2025 16:42

1. Change the signature for from_segment_ids_and_pos()

ca9d3bc

2. Add support for THD in from_segment_ids_and_pos() 3. Assert if load balanced segment_ids is passed to generate a segment_pos Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

0ee40a5

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Dec 23, 2025

View reviewed changes

Kshitij Janardan Lakhani added 2 commits December 22, 2025 16:55

Pass keyword-only args by name

ceec1ea

Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

nit: Fix typo to use seg_ids instead of segment_ids

ab00bb0

Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

KshitijLakhani changed the title ~~[JAX] Fix incorrect calculation of segment pos from segment ids~~ [JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API Dec 23, 2025

greptile-apps Bot reviewed Dec 23, 2025

View reviewed changes

nit: Fix comments

059c48d

Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

KshitijLakhani marked this pull request as draft December 23, 2025 01:15

Modify the function call to differentiate between load balancing and …

d524ad6

…actually reordered segment_ids and segment_pos Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>

pre-commit-ci Bot and others added 3 commits December 23, 2025 22:18

[pre-commit.ci] auto fixes from pre-commit.com hooks

d419f98

for more information, see https://pre-commit.ci

Fix the is_segment_ids_reordered to be set only when CP and load bala…

3efa504

…ncing Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Fix comments for from_segment_ids_and_pos()

e65062c

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani marked this pull request as ready for review December 24, 2025 00:23

greptile-apps Bot reviewed Dec 24, 2025

View reviewed changes

Code clean up

74a352e

for more information, see https://pre-commit.ci Fix lint errors Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/fix/incorrect-sequence-descr-from-seg-ids-only branch from 4101df5 to 74a352e Compare December 24, 2025 01:46

[pre-commit.ci] auto fixes from pre-commit.com hooks

e5381fb

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed Dec 24, 2025

View reviewed changes

KshitijLakhani requested review from cyanguwa, huanghua1994, jberchtold-nvidia, mgoldfarb-nvidia, mingxu1067 and zlsh80826 December 24, 2025 18:45

jberchtold-nvidia approved these changes Dec 24, 2025

View reviewed changes

Merge branch 'main' into klakhani/fix/incorrect-sequence-descr-from-s…

879afb4

…eg-ids-only

greptile-apps Bot reviewed Dec 31, 2025

View reviewed changes

ksivaman merged commit 26c82db into NVIDIA:main Dec 31, 2025
13 of 14 checks passed

KshitijLakhani mentioned this pull request Apr 1, 2026

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

Merged

13 tasks

yeandy mentioned this pull request Apr 7, 2026

Fix SequenceDescriptor API for TransformerEngine >= 2.12 AI-Hypercomputer/maxtext#3589

Closed

4 tasks

yeandy mentioned this pull request May 8, 2026

Fix SequenceDescriptor.from_segment_ids_and_pos() for TE >= 2.12 ROCm/maxtext#86

Closed

Uh oh!

Conversation

KshitijLakhani commented Dec 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

KshitijLakhani commented Dec 16, 2025

Uh oh!

KshitijLakhani commented Dec 17, 2025

Uh oh!

greptile-apps Bot commented Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 1/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

KshitijLakhani commented Dec 17, 2025

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (3)

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (3)

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (2)

Uh oh!

KshitijLakhani commented Dec 24, 2025

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (3)

Uh oh!

KshitijLakhani commented Dec 24, 2025

Uh oh!

jberchtold-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Dec 31, 2025

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (2)

Uh oh!

Uh oh!

KshitijLakhani commented Jan 5, 2026

Uh oh!

KshitijLakhani commented Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

KshitijLakhani commented Dec 16, 2025 •

edited

Loading

greptile-apps Bot commented Dec 17, 2025 •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading

KshitijLakhani commented Jan 5, 2026 •

edited

Loading