[JAX] Support segment_ids/pos as FA inputs

[zlsh80826]

Description

This PR adds segment_ids/pos limited support and deprecated fused_attn_thd API.

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 refractor

Changes

• Add a new SequenceDescriptor class for different sequence descriptions scenario.
  • from_seqlens for non-THD
  • from_seqlens_and_offsets for THD
  • from_segment_ids_and_pos for THD + ring attn (haven't implemented)
• Change the old fused_attn mask parameter to SequenceDescriptor. Passing mask in the position argument will work for a while but generating deprecation warning.
• Deprecate fused_attn_thd API as the refactored fused_attn can also support THD format.
• Remove small inputs in test_fused_attn.py as the long sequence inputs should cover.
• Add different sequence inputs tests in test_fused_attn.py

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

[zlsh80826]

/te-ci jax L1

[zlsh80826]

/te-ci jax L1

[mgoldfarb-nvidia]

Are these removals to cut down on number of [redundant] test cases?

[zlsh80826]

Yes, I think those tests are unnecessary because we no longer use max_512 kernels for seqlen < 512. The kernels for seqlen=128 are the same as seqlen=2048, so the 2048 test cases already cover the 128 cases.

[mgoldfarb-nvidia]

Would it help to cut down on test cases by creating a standalone unit test for the SequenceDesc to cover and check all of the cases? Then in this unit test we can use either Seqlens or SegmentIDs depending on THD or BSHD?

[zlsh80826]

Good idea

[mgoldfarb-nvidia]

I couldn't see if we applied max_segments_per_seq anywhere when generating the seqlen and offset to cudnn. I found that this was a very useful to help limit the overhead of the jax code that sets up seqlen/offset rather than assuming max_seq_len (see _get_seqlens_and_offsets). We should do a quick benchmark to see how much overhead if any is incurred.

[zlsh80826]

I added fea6b0e and cfe00a8 to reduce the seqlen and offset shape to (batch, max_segments_per_seq)