[rocm-jaxlib-v0.8.2] Backporting CI Benchmark related changes and fixes

[mmakevic-amd]

Motivation

Currently, CI benchmarks are failing on v0.8.2. This PR fixes it by backporting changes from #691 and #730.

Note: #622 will be closed in favour of this PR

Test Plan

I will manually trigger CI check before merging

Test Result

Workflow run as expected but gemma2 failed due to device time being above threshold: https://github.com/ROCm/xla/actions/runs/23635202015/job/68844227492
After unsetting HLO arg mode (from uninitialized ) job is running as expected: https://github.com/ROCm/xla/actions/runs/24037623275/job/70190205343

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[mmakevic-amd]

@hsharsha gemma2 is exhibiting some run-to-run variations (from ~12ms to ~600ms GPU device time), so the job sometimes fails. Workflow itself works as expected.

[i-chaochen]

once this PR is merged, are we going to have a benchmark CI on presubmit check? because seems I don't see this benchmark CI on presubmit check on 0.9.1 branch? for example this PR #756 maybe need to rebase or shall we create a new PR to check on 0.9.1?

[mmakevic-amd]

once this PR is merged, are we going to have a benchmark CI on presubmit check? because seems I don't see this benchmark CI on presubmit check on 0.9.1 branch? for example this PR #756 maybe need to rebase or shall we create a new PR to check on 0.9.1?

No, only as a postsubmit. I can enable it as a presubmit no problem, but the whole workflow lasts ~40min, as you can see https://github.com/ROCm/xla/actions/workflows/postsubmit_benchmark.yml, so I'm not sure if we want that

[mmakevic-amd]

One can trigger it manually before merging if that's necessary

[i-chaochen]

once this PR is merged, are we going to have a benchmark CI on presubmit check? because seems I don't see this benchmark CI on presubmit check on 0.9.1 branch? for example this PR #756 maybe need to rebase or shall we create a new PR to check on 0.9.1?

No, only as a postsubmit. I can enable it as a presubmit no problem, but the whole workflow lasts ~40min, as you can see https://github.com/ROCm/xla/actions/workflows/postsubmit_benchmark.yml, so I'm not sure if we want that

I think we can choose label to activate benchmark CI, just like what we have in claude code review and TSAN/ASAN. And it's best to let it run on pre-submit.

cc @nurmukhametov @alekstheod ?

[nurmukhametov]

once this PR is merged, are we going to have a benchmark CI on presubmit check? because seems I don't see this benchmark CI on presubmit check on 0.9.1 branch? for example this PR #756 maybe need to rebase or shall we create a new PR to check on 0.9.1?

No, only as a postsubmit. I can enable it as a presubmit no problem, but the whole workflow lasts ~40min, as you can see https://github.com/ROCm/xla/actions/workflows/postsubmit_benchmark.yml, so I'm not sure if we want that

IMO, 40 mins is fine. However, I don't understand two things:

• Why does the checkout take up to 15 min sometimes?
• Why don't we hit caches for builds?