Debug CI tests on Ada

[timmoon10]

This applies the changes in #393 to the PyTorch and Paddle tests. In particular, it only runs tests involving cuDNN fused attention on compute capabilities 8.0 and 9.0.

[timmoon10]

Pipeline 9489089

[timmoon10]

Running on an L40, I found that the JAX FP8 GEMM tests on integer matrices were failing. It seems cuBLAS chooses a split-k kernel that prevents us from getting bit-wise correct results, although it is still within the expected FP8 error. I've changed the matrix dims to help cuBLAS pick a nicer kernel.

Pipeline 9504876.

[timmoon10]

Pipeline 9617488 is green.

[timmoon10]

I've tweaked the PyTorch and JAX fused attention tests so we check if there's a supported backed (namely F16_arbitrary_seqlen on Ada). These pass when I run manually on an L40 and I've launched pipeline 9938409.

#403 adds some PyTorch attention tests and #411 adds backend detection logic to Paddle. We should hold off on merging until those are in.

[timmoon10]

This PR is now good to go, pending pipeline 10094388 pipeline 70748350.

[cyanguwa]

Just that one comment, otherwise looks good!

[timmoon10]

Tests passed in pipeline 10211932.

[ksivaman]

Looks good