Updated XLA_FLAGS in ci/jax.sh

Author: VeeraRajasekhar

ci/jax.sh

@@ -58,6 +58,7 @@ run_test_config() {
     run_default_fa 1 test_custom_call_compute.py
     run_default_fa 1 test_functions.py
     run 1 test_fused_attn.py
+    XLA_FLAGS='--xla_gpu_graph_level=0' run 1 test_fused_attn.py -k 'test_ck_unfused_smallseq_backend' # CK smallseq with GPU graph disabled
     NVTE_CK_USES_FWD_V3=0 NVTE_CK_USES_BWD_V3=0 run_default_fa_lbl "v2" 3 test_fused_attn.py # Using FAv2 for forward and backward pass
     run_default_fa 1 test_helper.py
     run_default_fa 1 test_layer.py #it effectevly always uses unfused attention

tests/jax/test_fused_attn.py

@@ -1262,8 +1262,9 @@ def test_jax_new_rng():
 @pytest.fixture
 def ck_smallseq_env(monkeypatch):
     """Enable CK small-seq path and disable XLA GPU graphs for these tests."""
+    if "xla_gpu_graph_level=0" not in os.environ.get("XLA_FLAGS", ""):
+        pytest.skip("Run with XLA_FLAGS='--xla_gpu_graph_level=0' pytest ...")
     monkeypatch.setenv("NVTE_FUSED_ATTN_CK_SMALLSEQ", "1")
-    monkeypatch.setenv("XLA_FLAGS", "--xla_gpu_graph_level=0")
     yield
 
 @pytest.mark.parametrize("dtype", [jnp.bfloat16, jnp.float16], ids=["BF16", "FP16"])