[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]

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

[zlsh80826]

I think we still need to keep those different SeqDescFormat in the test_fused_attn for a while. I don't have a better way to verify the converted seqlens/offsets are correct, the best way to verify that is checking the FA results. The current time for CI is still under 1hr so I think it is ok for now. But yeap, if we found some of SeqDescFormat are rarely used in the future, we can remove them or just testing on the specific input shape instead of all input shapes.

[mgoldfarb-nvidia]

LGTM.

[zlsh80826]

/te-ci jax L1

[zlsh80826]

/te-ci jax L1

[zlsh80826]

/te-ci jax L1