[TEST] Fix triton_kernels IGC after bea27e3

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

python/triton_kernels/tests/conftest.py

@@ -37,4 +37,5 @@ def pytest_configure(config):
     if worker_id is not None and worker_id.startswith("gw"):
         import torch
         gpu_id = int(worker_id[2:])  # map gw0 → 0, gw1 → 1, ...
-        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id % torch.cuda.device_count())
+        if torch.cuda.is_available():
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id % torch.cuda.device_count())

python/triton_kernels/triton_kernels/matmul_ogs.py

@@ -700,7 +700,7 @@ def matmul_ogs_torch(x, w, bias,
             if k > 0:
                 out[expt] = matmul_ogs_torch(
                     x[:, start_x:start_x+k], w[start_w:start_w+k, :], None,
-                    None, None, None, None, betas, gammas, None, round_x, round_y
+                    None, None, None, None, betas, gammas, None, round_x, round_y, device
                 )
             padded_k = triton.cdiv(k, block_k) * block_k
             start_x += padded_k if inner_routing_data.x_is_padded else k