[ROCm] Deselect crashing FusedAttention bwd8 + add pytest-timeout (ROCM-25626)#798
Open
srinivamd wants to merge 1 commit into
Open
[ROCm] Deselect crashing FusedAttention bwd8 + add pytest-timeout (ROCM-25626)#798srinivamd wants to merge 1 commit into
srinivamd wants to merge 1 commit into
Conversation
test_fused_attention_bwd8 causes a SIGSEGV in the XLA Triton compiler
(KernelReuseCache::GetWithStatus via ThunkEmitter::EmitTritonCustomCall)
on gfx942. pytest-xdist replaces the crashed worker, but the replacement
deadlocks during GPU runtime re-initialization, blocking the entire test
suite for 9+ hours until the job-level timeout fires.
Changes:
1. Deselect test_fused_attention_bwd8 in ci/run_pytest_rocm.sh (sgpu only).
This variant is NOT covered by the existing in-test skipTest guards,
which skip bwd{1,2,7,9} but not bwd8.
2. Add pytest-timeout to build/test-requirements.txt and pass
--timeout=1800 (30 min) to both sgpu and mgpu pytest invocations.
This is a safety net: if any future test hangs, it will be killed
after 30 minutes instead of wasting the full 10-hour job budget.
Tracked in: ROCM-25626
Upstream: openxla/xla#38972, jax-ml#34722
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the 9+ hour JAX UT hang on gfx942 (MI300X/MI308X) by:
Root Cause
test_fused_attention_bwd8causes a SIGSEGV in the XLA Triton compiler (KernelReuseCache::GetWithStatusviaThunkEmitter::EmitTritonCustomCall) on gfx942. pytest-xdist replaces the crashed worker, but the replacement deadlocks during GPU runtime re-initialization (corrupted GPU state from the SEGFAULT). With no--timeoutflag, the stuck replacement worker blocks the entire suite for 9h44m until the job-level 600-minute timeout fires.Why bwd8 is not caught by existing in-test skips
v0.9.2 already has
skipTestguards ingpu_ops_test.pylines 208-228 forbwd{1,2,7,9}, butbwd8is not in the skip list. This variant maps to a different parameter tuple that was not included in the upstream skip PR (jax-ml/jax#34722).Why v0.9.1 doesn't hang
On v0.9.1,
bwd8either doesn't hit the SEGFAULT codepath (different XLA/Triton version) or has been deselected (ROCm/jax#797).Changes
ci/run_pytest_rocm.sh--deselect=tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_bwd8to the single-accelerator pytest invocation (sgpu only — FusedAttentionTest is notmultiaccelerator)--timeout=1800(30 min) to both sgpu and mgpu pytest invocations. This is a safety net: if any future test hangs, it will be killed after 30 minutes instead of wasting the full 10-hour job budgetbuild/test-requirements.txtpytest-timeoutdependency (required for the--timeoutflag)Evidence
What this does NOT fix
Test plan
pytest-timeoutinstalls correctly from test-requirements.txtTracked in: ROCM-25626
Related: ROCm/jax#797 (v0.9.1 LDS deselect)