From b92647c580f93b6a1e598fb4fbf411216c6b25d5 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sun, 23 Nov 2025 05:22:21 +0000
Subject: [PATCH] Fix AMD Triton runtime error by reinstalling
 pytorch-triton-rocm

Uninstalling the conflicting `triton` package (upstream) on AMD systems often removes the `triton` directory shared with `pytorch-triton-rocm`, breaking the latter. This caused `ImportError: cannot import name 'Config' from 'triton'`.

This change adds a step to force reinstall `pytorch-triton-rocm` immediately after uninstalling `triton`, ensuring the correct package is present and intact for the runtime.
---
 speedrun.sh | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/speedrun.sh b/speedrun.sh
index 264710b..b6a1413 100644
--- a/speedrun.sh
+++ b/speedrun.sh
@@ -46,6 +46,9 @@ source .venv/bin/activate
 # if the NVIDIA version of triton (e.g. 3.4.0) is accidentally installed.
 if [ "$EXTRAS" == "amd" ]; then
     uv pip uninstall -q triton || true
+    # Uninstalling triton may have deleted the shared 'triton' directory, breaking pytorch-triton-rocm.
+    # Reinstall pytorch-triton-rocm to ensure it's intact.
+    uv pip install --force-reinstall --index-url https://repo.amd.com/rocm/whl/gfx1151 pytorch-triton-rocm
 fi
 
 # -----------------------------------------------------------------------------