[ROCm][CI] Retrying in case of batch variance effects and reducing flakiness

[AndreasKaratzas]

On ROCm there is non-determinism in the case of batch sizes greater than 1, and in the case of preemption test, this is exactly what is being tested. Therefore, we add a retry so that any chance of this test failing is met with a second attempt. Furthermore, we disable skinny GEMMs which is particularly useful for MI355 test deployment.

cc @kenroche

[gemini-code-assist]

Code Review

This pull request introduces several changes to improve CI stability on ROCm platforms. The main changes include adding a retry mechanism for a flaky test (test_with_ngram_gpu_spec_decoding), disabling skinny GEMM via an environment variable to mitigate non-determinism, and disabling prefix caching for the asynchronous scheduling tests. These changes are well-contained and only affect ROCm environments. The implementation appears correct and should contribute to more reliable CI runs.

_{Note: Security Review is unavailable for this PR.}

[AndreasKaratzas]

Adding ready label for tests to start.

[gshtras]

Disabling skinny GEMMs for tests is not a good idea. This is a production feature that is meant to always be on. Without it the tests are testing a synthetic use case that doesn't match production reality.
If you believe the GEMMs are bugged, we need to disable them until fixed. If not, the tests need to acomodate for it
If there is a specific batch invariance test that is meant to test a specific deterministic run, it may fit there, but it needs to be clearly stated in the docs.

[AndreasKaratzas]

Disabling skinny GEMMs for tests is not a good idea. This is a production feature that is meant to always be on. Without it the tests are testing a synthetic use case that doesn't match production reality. If you believe the GEMMs are bugged, we need to disable them until fixed. If not, the tests need to acomodate for it If there is a specific batch invariance test that is meant to test a specific deterministic run, it may fit there, but it needs to be clearly stated in the docs.

@gshtras Skinny GEMMs are bugged and I am working with @amd-hhashemi to resolve this. However, I think that having skinny GEMM failure affect all tests is not a good principle. I am working on introducing a skinny GEMM test amongst others here: #35183

That way, the skinny GEMM test will fail but others will remain intact.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tjtanaa]

LGTM

[noooop]

I couldn't help but laugh when I saw this headline.

Based on my understanding, batch variance eliminates all randomness.

batch variance? flakiness? retrying?

This sounds like a bug.

[AndreasKaratzas]

I couldn't help but laugh when I saw this headline.

Based on my understanding, batch variance eliminates all randomness.

batch variance? flakiness? retrying?

This sounds like a bug.

@noooop So, the flakiness is likely due to batch variance (not batch invariance). ROCm does not have batch invariance yet like CUDA. At the same time, I am still exploring why this test fails sometimes (rarely but still). I will confess however, that retry might have been an overkill 😅

[noooop]

Sorry, I misread that. I thought it was batch invariance.