[TESTS] Enable parallel pytest in CI for CUDA (triton-lang#1905)

Run most of the pytest in parallel, this allows to speed up CI from
36min to 10min for A100 and 22min to 6min for H100. Some tests still
need to run serially like runtime tests.

.github/workflows/integration-tests.yml

@@ -66,6 +66,7 @@ jobs:
           python3 -m pip install --upgrade pip
           python3 -m pip install cmake==3.24
           python3 -m pip install --no-build-isolation -vvv '.[tests]'
+          python3 -m pip install pytest-xdist
 
       - name: Run lit tests
         if: ${{ env.BACKEND == 'CUDA'}}
@@ -82,7 +83,9 @@ jobs:
         if: ${{ env.BACKEND == 'CUDA'}}
         run: |
           cd python/test/unit
-          python3 -m pytest
+          python3 -m pytest -n 8 --ignore=runtime
+          # run runtime tests serially to avoid race condition with cache handling.
+          python3 -m pytest runtime/
 
       - name: Create artifacts archive
         if: ${{(matrix.runner[0] == 'self-hosted') && (matrix.runner[1] == 'V100' || matrix.runner[1] == 'A100' || matrix.runner[1] == 'H100')}}

python/test/unit/operators/test_cross_entropy.py

@@ -29,11 +29,13 @@ def test_op(M, N, dtype, mode):
     # backward pass
     elif mode == 'backward':
         dy = torch.randn_like(tt_y)
-        # triton backward
-        tt_y.backward(dy)
-        tt_dx = x.grad.clone()
+        # run torch first to make sure cuda context is initialized
         # torch backward
-        x.grad.zero_()
         th_y.backward(dy)
         th_dx = x.grad.clone()
+        # triton backward
+        x.grad.zero_()
+        tt_y.backward(dy)
+        tt_dx = x.grad.clone()
+
         torch.testing.assert_allclose(th_dx, tt_dx)