[JAX] WAR for XLA associative scan rewriter crash

[KshitijLakhani]

Description

An iterator-invalidation bug in XLA’s AssociativeScanRewriter causes intermittent compiler SIGSEGVs in CP fused-attention tests. The permanent fix is being handled in : openxla/xla#48879

This PR temporarily disables the affected XLA module pass for those CI tests. It will be reverted once the XLA fix is merged and available in the NVIDIA JAX image. This guarantees correctness in CI only (and not feature code) at the expense of some perf.

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 refactoring

Changes

The two command changes :

1. The first command runs all non CP tests normally.
2. The second runs only the CP class with the XLA pass disabled (thereby limiting the scope of the WAR)

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

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with the workaround limited to the affected context-parallel CI tests.

Summary

This PR narrows a temporary XLA workaround to the affected JAX context-parallel fused-attention tests:

• Runs non-context-parallel fused-attention tests normally.
• Runs the context-parallel test class separately with associative-scan-rewriter disabled.
• Produces separate JUnit reports for both test groups.
• Contains no changes since the previous review SHA.

_{Reviews (2) · Last reviewed commit: "Work around associative scan rewriter cr..."}

[jberchtold-nvidia]

LGTM pending CI, thanks!

[ptrendx]

@KshitijLakhani please sign your commit

[KshitijLakhani]

@KshitijLakhani please sign your commit

Just did @ptrendx . Thanks for catching it !

[KshitijLakhani]

Not merging this PR in as it only helps unblock CI and the XLA fix was just merged in.
JAX is expected to pick this up in their next nightly container. I'll run these tests out with the latest JAX container on 09/18 and the expectation is that it will pass in which case this PR will be closed without merging. However, if the fix isn't still available on 09/18, I will merge this PR to temporarily unblock CI (I confirmed that the pipeline launched on this PR completed successfully so the CI unblocking is verified)
@jberchtold-nvidia @aybchan @ptrendx