[DOC] Fix documentation pipeline (#6702)

1. Change the default build directory from `python/` to the root of
triton.
2. Update on the CI runners.
3. Increase the relative tolerance of 08-group-gemm tutorial due to its
failure on some GPUs using mmav2.
Tested https://github.com/triton-lang/triton/actions/runs/14816446953

Author: Jokeren

.github/workflows/documentation.yml

@@ -8,7 +8,7 @@ permissions: read-all
 
 jobs:
   Build-Documentation:
-    runs-on: [a100-runner-set]
+    runs-on: [nvidia-a100]
     timeout-minutes: 30
 
     steps:

docs/conf.py

@@ -43,7 +43,7 @@ def get_cmake_dir():
     plat_name = sysconfig.get_platform()
     python_version = sysconfig.get_python_version()
     dir_name = f"cmake.{plat_name}-{sys.implementation.name}-{python_version}"
-    cmake_dir = Path("../python") / "build" / dir_name
+    cmake_dir = Path("../build") / dir_name
     return cmake_dir
 
 
@@ -100,7 +100,7 @@ def setup(app):
     app.connect("autodoc-process-signature", process_sig)
     max_jobs = os.getenv("MAX_JOBS", str(2 * os.cpu_count()))
     print(f"Installing Triton Python package using {max_jobs} threads")
-    subprocess.run("pip install -e ../python", shell=True, env=os.environ.copy())
+    subprocess.run("pip install -e ../", shell=True, env=os.environ.copy())
 
     setup_generated_mlir_docs()
 

python/tutorials/08-grouped-gemm.py

@@ -392,12 +392,12 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]):
 tri_out = group_gemm_fn(group_A, group_B)
 ref_out = [torch.matmul(a, b) for a, b in zip(group_A, group_B)]
 for i in range(group_size):
-    assert torch.allclose(ref_out[i], tri_out[i], atol=1e-2, rtol=0)
+    assert torch.allclose(ref_out[i], tri_out[i], atol=1e-2, rtol=1e-2)
 
 if supports_tma():
     tri_tma_out = group_gemm_tma_fn(group_A, group_B_T)
     for i in range(group_size):
-        assert torch.allclose(ref_out[i], tri_tma_out[i], atol=1e-2, rtol=0)
+        assert torch.allclose(ref_out[i], tri_tma_out[i], atol=1e-2, rtol=1e-2)
 
 
 # only launch the kernel, no tensor preparation here to remove all overhead