Skip to content

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

Closed
yeandy wants to merge 1 commit into
release/v26.3from
release/v26.3-te-sequence-descriptor-api-2.12
Closed

Fix SequenceDescriptor.from_segment_ids_and_pos() for TE >= 2.12#86
yeandy wants to merge 1 commit into
release/v26.3from
release/v26.3-te-sequence-descriptor-api-2.12

Conversation

@yeandy

@yeandy yeandy commented May 8, 2026

Copy link
Copy Markdown
Collaborator

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'

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'
@yeandy yeandy closed this May 18, 2026
@yeandy yeandy deleted the release/v26.3-te-sequence-descriptor-api-2.12 branch May 18, 2026 11:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant