[Feature Request, JAX] Support fused Softmax op when sequence length is larger than 2048

[MoFHeka]

Fail to Jax custom call when training a model with sequence length is larger than 2048. At nowadays, seqlen in most model was bigger than 2048, at least 4096. 2048 seqlen for TE is too small!

[zlsh80826]

Thanks for the feedback. We have plan to remove the restriction of 2048.

[ptrendx]

@zlsh80826 Is this already fixed?

[zlsh80826]

Yes, it is fixed in https://github.com/NVIDIA/TransformerEngine/pull/796. The softmax now supports up to seqlen = 16384.