This adds support for accessing peer memory from Trition.
TODO: b/347711619 - Upstream after we've upstreamed the rest of the launcher.

diff --git b/python/test/unit/runtime/test_peer_access.py b/python/test/unit/runtime/test_peer_access.py
new file mode 100644
--- /dev/null
+++ b/python/test/unit/runtime/test_peer_access.py
@@ -0,0 +1,24 @@
+import pytest
+import torch
+
+import triton
+import triton.language as tl
+
+
+def test_peer_access(device):
+    if not hasattr(torch, device):
+       pytest.skip(f"{device} does not support peer access")
+    if getattr(torch, device).device_count() < 2:
+       pytest.skip("need at least 2 devices to test peer access")
+
+    @triton.jit
+    def device_accumulate(my_ptr, peer_ptr):
+        tl.store(my_ptr, tl.load(my_ptr) + tl.load(peer_ptr))
+
+    my_tensor = torch.randn(1, device=f"{device}:0")
+    peer_tensor = torch.randn(1, device=f"{device}:1")
+    expected = my_tensor + peer_tensor.to(device=f"{device}:0")
+
+    device_accumulate[(1,1,1)](my_tensor, peer_tensor)
+
+    torch.testing.assert_close(my_tensor, expected)

--- a/third_party/nvidia/backend/driver.c
+++ b/third_party/nvidia/backend/driver.c
@@ -965,6 +965,37 @@ cleanup:
   return NULL;
 }
 
+// Enable peer access if dev_ptr is allocated on a different device than the
+// device on which we will execute the kernel.
+PyObject* enablePeerAccessIfNecessary(CUdeviceptr dev_ptr) {
+  CUmemorytype mem_type = CU_MEMORYTYPE_HOST;
+  CUresult status = cuPointerGetAttribute(
+      &mem_type, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, dev_ptr);
+  if (status != CUDA_SUCCESS || mem_type != CU_MEMORYTYPE_DEVICE) {
+    // Not peer memory
+    Py_RETURN_NONE;
+  }
+  int mem_device_ordinal = 0;
+  CUDA_CHECK_AND_RETURN_NULL(cuPointerGetAttribute(
+      &mem_device_ordinal, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, dev_ptr));
+  CUdevice mem_device = 0;
+  CUDA_CHECK_AND_RETURN_NULL(cuDeviceGet(&mem_device, mem_device_ordinal));
+  CUdevice compute_device = 0;
+  CUDA_CHECK_AND_RETURN_NULL(cuCtxGetDevice(&compute_device));
+  if (mem_device != compute_device) {
+    CUcontext mem_ctx = NULL;
+    CUDA_CHECK_AND_RETURN_NULL(cuDevicePrimaryCtxRetain(&mem_ctx, mem_device));
+    CUresult status = cuCtxEnablePeerAccess(mem_ctx, /*flags=*/0);
+    if (status == CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED) {
+      status = CUDA_SUCCESS;
+    }
+    CUDA_CHECK_AND_RETURN_NULL(status);
+  }
+  Py_RETURN_NONE;
+cleanup:
+  return NULL;
+}
+
 static void _launch(int gridX, int gridY, int gridZ, int num_warps,
                     int num_ctas, int launch_cooperative_grid, int launch_pdl,
                     int shared_memory, CUstream stream, CUfunction function,
@@ -1076,6 +1107,9 @@ bool extractPointer(void *ptr, PyObject *obj) {
   if (*dev_ptr == 0) {
     return true; // valid nullptr
   }
+  if (enablePeerAccessIfNecessary(*dev_ptr) == NULL) {
+    return false;
+  }
   CUresult status = cuPointerGetAttribute(
       dev_ptr, CU_POINTER_ATTRIBUTE_DEVICE_POINTER, *dev_ptr);
   if (status == CUDA_ERROR_INVALID_VALUE) {