Update

Author: Jokeren

bench/bench/bench_mlp.py

@@ -10,13 +10,14 @@
 from triton_bench.numerics import InFlexData
 from triton_bench.routing import routing_torch, simulate_expert_sharded_routing
 
+target = triton.runtime.driver.active.get_current_target()
+
 
 def is_hip_cdna4():
-    target = triton.runtime.driver.active.get_current_target()
     return target.backend == 'hip' and target.arch == 'gfx950'
 
 
-if torch.cuda.is_available() and not is_hip_cdna4():
+if torch.cuda.is_available() and not target.backend == "hip":
     from triton._C.libtriton import nvidia
     cublas_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
     cublas = nvidia.cublas.CublasLt(cublas_workspace)