[JAX] Support vmapped cudnn attention

[qGentry]

Is your feature request related to a problem? Please describe.

Basically, cudnn attention doesn't work when vmapped because of this assertion here:
https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/cpp_extensions/attention.py#L563
This limits its usability when one wants to implement pipeline parallelism as it requires vmapping staged layers over stage buffer.

Describe the solution you'd like

Properly implemented vmap rules for the cudnn attention.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!