IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/extra/thunder/cuda/include/pyutils/broker.cuh

@@ -0,0 +1,551 @@
+/**
+ * @file broker.cuh
+ * @brief Utility for multiprocess data exchange and synchronization.
+ * 
+ * This file provides the KittensBroker class, which enables efficient inter-process
+ * communication and synchronization using POSIX shared memory, semaphores, and sockets.
+ * The broker is designed to work in multi-GPU environments where processes need to
+ * exchange data and synchronize execution across different local ranks.
+ * 
+ * @note This implementation relies on POSIX IPC mechanisms and is intended for
+ *       Unix-like systems. All processes must be running on the same node.
+ */
+
+#pragma once
+
+#include <cerrno>
+#include <cstdint>
+#include <cstring>
+#include <fcntl.h>
+#include <semaphore.h>
+#include <stdexcept>
+#include <sys/mman.h>
+#include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/un.h>
+#include <sys/uio.h>
+#include <unistd.h>
+#include <vector>
+
+#if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
+    #error "KittensBroker is not supported on Windows"
+#endif
+
+namespace kittens {
+
+namespace detail {
+namespace broker {
+
+static constexpr int MAX_LOCAL_WORLD_SIZE = 72;
+static constexpr int VAULT_SIZE_PER_RANK = 64; // sizeof(cudaIpcMemHandle_t)
+
+struct KittensVault {
+    static constexpr int INIT_CODE = 0x43617473; // "Cats"
+    int init;
+    int barrier;
+    int sense;
+    uint8_t data[MAX_LOCAL_WORLD_SIZE * VAULT_SIZE_PER_RANK];
+};
+
+static constexpr int SHM_SIZE = (sizeof(KittensVault) + 4095) / 4096 * 4096;
+
+__host__ inline static void init_sync(
+    int local_rank,
+    volatile KittensVault *vault
+) {
+    if (local_rank == 0) {
+        // initialize barrier resources
+        vault->barrier = 0;
+        vault->sense = 0;
+        __sync_synchronize(); // make previous writes visible
+        vault->init = KittensVault::INIT_CODE;
+    } else {
+        while (vault->init != KittensVault::INIT_CODE) usleep(1);
+        __sync_synchronize(); // see leader's previous writes
+    }
+}
+
+__host__ inline static void sync(
+    int local_world_size,
+    volatile KittensVault *vault
+) {
+    if (vault->init != KittensVault::INIT_CODE)
+        throw std::runtime_error("KittensBroker: KittensVault not initialized");
+
+    // Phase 1
+    int arrived = __sync_add_and_fetch(&vault->barrier, 1);
+    if (arrived == local_world_size) vault->sense = 1;
+    while (!vault->sense) usleep(1);
+
+    // Make previous writes visible
+    __sync_synchronize();
+
+    // Phase 2
+    arrived = __sync_add_and_fetch(&vault->barrier, -1);
+    if (arrived == 0) vault->sense = 0;
+    while (vault->sense) usleep(1);
+}
+
+__host__ inline void *create_shm(const char *key, size_t size) {
+    int shm_fd;
+    shm_fd = shm_open(key, O_RDWR | O_CREAT | O_EXCL | O_CLOEXEC, 0600);
+
+    if (shm_fd < 0) {
+        if (errno == EEXIST)
+            throw std::runtime_error("KittensBroker: Named shared memory already exists");
+        throw std::runtime_error("KittensBroker: Failed to create shared memory");
+    }
+
+    if (ftruncate(shm_fd, size) != 0) {
+        shm_unlink(key);
+        close(shm_fd);
+        throw std::runtime_error("KittensBroker: Failed to truncate shared memory");
+    }
+
+    void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
+    close(shm_fd);
+    if (addr == MAP_FAILED) {
+        shm_unlink(key);
+        throw std::runtime_error("KittensBroker: Failed to map to shared memory");
+    }
+
+    return addr;
+}
+
+__host__ inline void *open_shm(const char *key, size_t size) {
+    int shm_fd;
+    while (true) {
+        shm_fd = shm_open(key, O_RDWR | O_CLOEXEC, 0);
+        if (shm_fd >= 0)
+            break;
+        if (errno != ENOENT) 
+            throw std::runtime_error("KittensBroker: Failed to open shared memory");
+        usleep(1);
+    }
+
+    struct stat shm_st;
+    do {
+        if (fstat(shm_fd, &shm_st) != 0) {
+            shm_unlink(key);
+            close(shm_fd);
+            throw std::runtime_error("KittensBroker: Failed to open shared memory stats");
+        }
+        usleep(1);
+    } while ((size_t)shm_st.st_size < size);
+
+    void *addr = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
+    close(shm_fd);
+    if (addr == MAP_FAILED) {
+        shm_unlink(key);
+        throw std::runtime_error("KittensBroker: Failed to map to shared memory");
+    }
+
+    return addr;
+}
+
+__host__ inline void unlink_shm(const char *key) {
+    shm_unlink(key);
+}
+
+__host__ inline void unmap_shm(void *addr, size_t size) {
+    munmap(addr, size);
+}
+
+__host__ inline int create_socket(const char *key, int local_rank) {
+    int sock_fd;
+    if ((sock_fd = socket(AF_UNIX, SOCK_DGRAM | SOCK_CLOEXEC, 0)) < 0)
+        throw std::runtime_error("KittensBroker: Socket creation error");
+
+    struct sockaddr_un addr;
+    memset(&addr, 0, sizeof(addr));
+    addr.sun_family = AF_UNIX;
+
+    char unique_key[64];
+    int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank);
+    if (n < 0 || n >= (int)sizeof(unique_key)) {
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Socket name too long"); 
+    }
+
+    size_t len = strnlen(unique_key, sizeof(addr.sun_path));
+    if (len > (sizeof(addr.sun_path) - 1)) {
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Socket name too long");
+    }
+    strcpy(addr.sun_path, unique_key);
+    unlink(unique_key);
+
+    if (bind(sock_fd, (struct sockaddr *)&addr, SUN_LEN(&addr)) < 0) {
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Failed to bind socket");
+    }
+
+    return sock_fd;
+}
+
+__host__ inline void send_fd(
+    int sock_fd,
+    int data_fd,
+    const char *dst_key,
+    int dst_local_rank,
+    int src_local_rank
+) {
+    union {
+      struct cmsghdr cm;
+      char* control;
+    } control_un;
+
+    size_t sizeof_control = CMSG_SPACE(sizeof(int));
+    control_un.control = reinterpret_cast<char *>(malloc(sizeof_control));
+    if (!control_un.control) {
+        close(sock_fd);
+        close(data_fd);
+        throw std::runtime_error("KittensBroker: Failed to allocate a control buffer");
+    }
+  
+    struct msghdr msg {};
+    msg.msg_control = control_un.control;
+    msg.msg_controllen = sizeof_control;
+  
+    struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg);
+    cmptr->cmsg_len = CMSG_LEN(sizeof(int));
+    cmptr->cmsg_level = SOL_SOCKET;
+    cmptr->cmsg_type = SCM_RIGHTS;
+    memmove(CMSG_DATA(cmptr), &data_fd, sizeof(data_fd));
+
+    struct sockaddr_un addr {};
+    addr.sun_family = AF_UNIX;
+    char dst_unique_key[64];
+    int n = snprintf(dst_unique_key, sizeof(dst_unique_key), "%s%d", dst_key, dst_local_rank);
+    if (n < 0 || n >= (int)sizeof(dst_unique_key)) { 
+        free(control_un.control);
+        close(sock_fd);
+        close(data_fd);
+        throw std::runtime_error("KittensBroker: dst path too long"); 
+    }
+    strcpy(addr.sun_path, dst_unique_key);
+    msg.msg_name = (void *)&addr;
+    msg.msg_namelen = sizeof(struct sockaddr_un);
+  
+    int payload = src_local_rank;
+    struct iovec iov[1];
+    iov[0].iov_base = &payload;
+    iov[0].iov_len  = sizeof(payload);
+    msg.msg_iov = iov;
+    msg.msg_iovlen = 1;
+  
+    while (true) {
+        ssize_t sent = sendmsg(sock_fd, &msg, 0);
+        if (sent <= 0) {
+            if (errno == EINTR) continue;
+            close(sock_fd);
+            close(data_fd);
+            free(control_un.control);
+            throw std::runtime_error("KittensBroker: Failed to send FD over socket");
+        }
+        break;
+    }
+
+    free(control_un.control);
+}
+
+__host__ inline void recv_fd(int sock_fd, int *data_fd, int *src_local_rank) {
+    union {
+      struct cmsghdr cm;
+      char* control;
+    } control_un;
+
+    size_t sizeof_control = CMSG_SPACE(sizeof(int));
+    control_un.control = reinterpret_cast<char *>(malloc(sizeof_control));
+    if (!control_un.control) {
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Failed to allocate a control buffer");
+    }
+
+    struct msghdr msg {};
+    msg.msg_control = control_un.control;
+    msg.msg_controllen = sizeof_control;
+
+    int payload = -1;
+    struct iovec iov[1];
+    iov[0].iov_base = &payload;
+    iov[0].iov_len  = sizeof(payload);
+    msg.msg_iov = iov;
+    msg.msg_iovlen = 1;
+
+    while (true) {
+        ssize_t received = recvmsg(sock_fd, &msg, 0);
+        if (received < 0 && errno == EINTR) {
+            msg.msg_controllen = sizeof_control;
+            msg.msg_iovlen = 1;
+            continue;
+        }
+        if (received < static_cast<ssize_t>(sizeof(*data_fd))) {
+            free(control_un.control);
+            close(sock_fd);
+            throw std::runtime_error("KittensBroker: Failed to receive data over socket");
+        }
+        break;
+    }
+
+    if (msg.msg_flags & MSG_CTRUNC) {
+        free(control_un.control);
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Control data truncated");
+    }
+
+    struct cmsghdr *cmptr = CMSG_FIRSTHDR(&msg);
+    if (!cmptr ||
+        cmptr->cmsg_len != CMSG_LEN(sizeof(int)) ||
+        cmptr->cmsg_level != SOL_SOCKET ||
+        cmptr->cmsg_type != SCM_RIGHTS) {
+        free(control_un.control);
+        close(sock_fd);
+        throw std::runtime_error("KittensBroker: Failed to receive data over socket");
+    }
+
+    memmove(data_fd, CMSG_DATA(cmptr), sizeof(*data_fd));
+    free(control_un.control);
+    *src_local_rank = payload;
+}
+
+__host__ inline void unlink_socket(const char *key, int local_rank) {
+    char unique_key[64];
+    int n = snprintf(unique_key, sizeof(unique_key), "%s%d", key, local_rank);
+    if (n < 0 || n >= (int)sizeof(unique_key))
+        throw std::runtime_error("KittensBroker: Socket name too long");
+    unlink(unique_key);
+}
+
+__host__ inline void close_socket(int sock_fd) {
+    close(sock_fd);
+}
+
+} // namespace broker
+} // namespace detail
+
+/**
+    @brief KittensBroker utility for multiprocess data exchange.
+
+    Note that the code relies on POSIX sockets/shared memory/semaphores for 
+    inter-process communication and synchronization.
+
+    The main functions meant to be used by the user are:
+
+        KittensBroker broker(local_rank, local_world_size);
+        broker.exchange_data(dst, src, size); // exchange data between all processes
+        broker.exchange_fds(dst, src_fd); // exchange file descriptors between all processes
+        broker.broadcast_fd(dst, src_fd, src_rank); // broadcast file descriptor from src_rank to all processes
+        broker.sync(); // wait until all processes reach here
+ */
+struct KittensBroker {
+    // TODO: make unique per process group
+    static inline constexpr const char *SHM_KEY_ = "/kittens_broker_shm";
+    static inline constexpr const char *SOCK_KEY_ = "/tmp/kittens_broker.sock";
+
+    int local_rank_;
+    int local_world_size_;
+
+    void *shm_raw_;
+    volatile detail::broker::KittensVault *shm_;
+    int sock_;
+
+    __host__ inline KittensBroker(int local_rank, int local_world_size)
+        : local_rank_(local_rank), 
+          local_world_size_(local_world_size),
+          shm_raw_(nullptr),
+          shm_(nullptr),
+          sock_(-1) {
+        if (local_rank_ < 0)
+            throw std::runtime_error("KittensBroker: Local rank must be non-negative");
+        if (local_rank_ >= local_world_size_)
+            throw std::runtime_error("KittensBroker: Local rank is greater than local world size");
+        if (local_world_size_ > detail::broker::MAX_LOCAL_WORLD_SIZE)
+            throw std::runtime_error("KittensBroker: Local world size is greater than MAX_LOCAL_WORLD_SIZE");
+
+        if (local_rank_ == 0) {
+            shm_raw_ = detail::broker::create_shm(SHM_KEY_, sizeof(detail::broker::KittensVault));
+            shm_ = reinterpret_cast<volatile detail::broker::KittensVault *>(shm_raw_);
+            memset(shm_raw_, 0, sizeof(detail::broker::KittensVault));
+        } else {
+            shm_raw_ = detail::broker::open_shm(SHM_KEY_, sizeof(detail::broker::KittensVault));
+            shm_ = reinterpret_cast<volatile detail::broker::KittensVault *>(shm_raw_);
+        }
+        detail::broker::init_sync(local_rank_, shm_);
+        detail::broker::sync(local_world_size_, shm_);
+
+        if (local_rank_ ==0)
+            detail::broker::unlink_shm(SHM_KEY_);
+        detail::broker::sync(local_world_size_, shm_);
+
+        sock_ = detail::broker::create_socket(SOCK_KEY_, local_rank_);
+        detail::broker::sync(local_world_size_, shm_);
+    }
+
+    KittensBroker(const KittensBroker&) = delete;
+    KittensBroker& operator=(const KittensBroker&) = delete;
+
+    __host__ inline KittensBroker(KittensBroker&& other) noexcept
+        : local_rank_(other.local_rank_),
+          local_world_size_(other.local_world_size_),
+          shm_raw_(other.shm_raw_),
+          shm_(other.shm_),
+          sock_(other.sock_) {
+        other.local_rank_ = -1;
+        other.local_world_size_ = -1;
+        other.shm_raw_ = nullptr;
+        other.shm_ = nullptr;
+        other.sock_ = -1;
+    }
+
+    __host__ inline void destroy() {
+        if (shm_raw_) {
+            detail::broker::unmap_shm(shm_raw_, sizeof(detail::broker::KittensVault));
+            shm_raw_ = nullptr;
+            shm_ = nullptr;
+        }
+        if (sock_ >= 0) {
+            detail::broker::unlink_socket(SOCK_KEY_, local_rank_);
+            detail::broker::close_socket(sock_);
+            sock_ = -1;
+        }
+        local_rank_ = -1;
+        local_world_size_ = -1;
+    }
+
+    __host__ inline KittensBroker& operator=(KittensBroker&& other) noexcept {
+        if (this != &other) {
+            destroy();
+            local_rank_ = other.local_rank_;
+            local_world_size_ = other.local_world_size_;
+            shm_raw_ = other.shm_raw_;
+            shm_ = other.shm_;
+            sock_ = other.sock_;
+            other.local_rank_ = -1;
+            other.local_world_size_ = -1;
+            other.shm_raw_ = nullptr;
+            other.shm_ = nullptr;
+            other.sock_ = -1;
+        }
+        return *this;
+    }
+
+    __host__ inline ~KittensBroker() {
+        destroy();
+    }
+
+    __host__ inline void sync(int num_ranks = -1) {
+        if (num_ranks == -1)
+            num_ranks = local_world_size_;
+        else if (num_ranks < 0 || num_ranks > local_world_size_)
+            throw std::runtime_error("KittensBroker: Invalid number of ranks");
+
+        detail::broker::sync(num_ranks, shm_);
+    }
+
+    __host__ inline void exchange_data(void *dst_, const void *src_, size_t size) {
+        if (size > detail::broker::VAULT_SIZE_PER_RANK)
+            throw std::runtime_error("KittensBroker: Size is greater than VAULT_SIZE_PER_RANK");
+
+        uint8_t *dst = reinterpret_cast<uint8_t *>(dst_);
+        const uint8_t *src = reinterpret_cast<const uint8_t *>(src_);
+
+        // Exchange data
+        sync(); // ensure all processes enter together
+        memcpy(const_cast<uint8_t *>(shm_->data) + local_rank_ * detail::broker::VAULT_SIZE_PER_RANK, src, size);
+        sync(); // ensure all processes exit together
+    
+        // Pack and copy back to destination
+        for (int i = 0; i < local_world_size_; i++)
+            memcpy(dst + i * size, const_cast<uint8_t *>(shm_->data) + i * detail::broker::VAULT_SIZE_PER_RANK, size);
+    }
+
+    __host__ inline void exchange_fds(int *dst, const int data_fd) {
+        if (dst == nullptr)
+            throw std::runtime_error("KittensBroker: dst is null");
+        if (data_fd < 0)
+            throw std::runtime_error("KittensBroker: source fd is negative");
+
+        // Initialize dst buffer
+        for (int i = 0; i < local_world_size_; ++i)
+            dst[i] = -1;
+
+        // Ensure all processes enter together
+        sync();
+
+        if (local_rank_ == 0) {
+            // Rank 0 receives all FDs from and distributes them to other ranks
+            dst[0] = data_fd;
+            for (int i = 0; i < local_world_size_ - 1; i++) {
+                int received_fd;
+                int src_local_rank;
+                detail::broker::recv_fd(sock_, &received_fd, &src_local_rank);
+                if (received_fd < 0)
+                    throw std::runtime_error("KittensBroker: Failed to receive FD over socket");
+                if (src_local_rank == local_rank_)
+                    throw std::runtime_error("KittensBroker: Invalid source rank");
+                dst[src_local_rank] = received_fd;
+            }
+            for (int dst_local_rank = 1; dst_local_rank < local_world_size_; dst_local_rank++) {
+                for (int src_local_rank = 0; src_local_rank < local_world_size_; src_local_rank++) {
+                    if (dst_local_rank == src_local_rank)
+                        continue;
+                    detail::broker::send_fd(sock_, dst[src_local_rank], SOCK_KEY_, dst_local_rank, src_local_rank);
+                }
+            }
+            close(dst[0]); // no longer needed
+            dst[0] = -1;
+        } else {
+            // The rest sends its FD to and receives the other FDs from rank 0
+            detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, 0, local_rank_);
+            close(data_fd); // no longer needed
+            for (int i = 0; i < local_world_size_ - 1; i++) {
+                int received_fd;
+                int src_local_rank;
+                detail::broker::recv_fd(sock_, &received_fd, &src_local_rank);
+                if (received_fd < 0)
+                    throw std::runtime_error("KittensBroker: Failed to receive FD over socket");
+                if (src_local_rank == local_rank_)
+                    throw std::runtime_error("KittensBroker: Invalid source rank");
+                dst[src_local_rank] = received_fd;
+            }
+        }
+
+        // Ensure all processes exit together
+        sync();
+    }
+
+    __host__ inline void broadcast_fd(int *dst, const int data_fd, const int src_local_rank) {
+        if (src_local_rank < 0 || src_local_rank >= local_world_size_)
+            throw std::runtime_error("KittensBroker: Invalid source rank");
+
+        // Ensure all processes enter together
+        sync();
+
+        if (local_rank_ == src_local_rank) {
+            if (data_fd < 0)
+                throw std::runtime_error("KittensBroker: Source rank has invalid FD");
+            for (int dst_local_rank = 0; dst_local_rank < local_world_size_; dst_local_rank++) {
+                if (dst_local_rank == src_local_rank)
+                    continue;
+                detail::broker::send_fd(sock_, data_fd, SOCK_KEY_, dst_local_rank, src_local_rank);
+            }
+            close(data_fd); // no longer needed
+        } else {
+            if (!dst)
+                throw std::runtime_error("KittensBroker: Destination rank has invalid buffer");
+            int _src_local_rank;
+            detail::broker::recv_fd(sock_, dst, &_src_local_rank);
+            if (*dst < 0)
+                throw std::runtime_error("KittensBroker: Failed to receive valid FD over socket");
+            if (_src_local_rank != src_local_rank)
+                throw std::runtime_error("KittensBroker: Invalid source rank");
+        }
+
+        // Ensure all processes exit together
+        sync();
+    }
+};
+
+} // namespace kittens

tinygrad_repo/extra/thunder/cuda/include/pyutils/club.cuh

@@ -0,0 +1,122 @@
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+/*
+    CUDA-specific ThreadPool
+
+    Example usage
+
+    // Construction
+    KittensClub club(device_ids, NUM_DEVICES);
+
+    // Dispatch work to all threads (no need to set device)
+    club.execute([&](int dev_idx) {
+        int dev;
+        CUDACHECK(cudaGetDevice(&dev));
+        if (dev != dev_idx) {
+            fprintf(stderr, "Device mismatch: expected %d, got %d\n", dev_idx, dev);
+            exit(1);
+        }
+    });
+*/
+class KittensClub {
+public:
+    __host__ inline KittensClub(const int *device_ids, const int num_devices);
+    __host__ inline KittensClub(const int *device_ids, const cudaStream_t *streams, const int num_devices);
+    __host__ inline ~KittensClub();
+
+    // Dispatches `task` to all threads, and waits for all threads to finish (using cv)
+    __host__ inline void execute(std::function<void(int, cudaStream_t)> task);
+
+private:
+    // Condition indicators
+    bool stop;
+    std::vector<bool> task_available;
+    int n_task_done;
+
+    // Threadpool
+    std::vector<std::thread> workers;
+    
+    // Streams for each device
+    std::vector<cudaStream_t> streams;
+    
+    // Main entry point for each thread
+    __host__ inline void worker(int worker_id, int device_id);
+
+    // Used to dispatch work to all threads
+    std::function<void(int, cudaStream_t)> current_task;
+
+    // Synchronization
+    std::mutex mutex;
+    std::condition_variable cond_task_available;
+    std::condition_variable cond_task_done;
+};
+    
+__host__ inline KittensClub::KittensClub(const int *device_ids, const int num_devices) : stop(false), n_task_done(0) {
+    for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
+        task_available.push_back(false);
+        streams.push_back(0); // Use default stream (null stream)
+        workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); });
+    }
+}
+
+__host__ inline KittensClub::KittensClub(const int *device_ids, const cudaStream_t *streams_in, const int num_devices) : stop(false), n_task_done(0) {
+    for (size_t dev_idx = 0; dev_idx < num_devices; ++dev_idx) {
+        task_available.push_back(false);
+        streams.push_back(streams_in[dev_idx]);
+        workers.emplace_back([this, dev_idx, device_ids] { worker(dev_idx, device_ids[dev_idx]); });
+    }
+}
+    
+__host__ inline KittensClub::~KittensClub() {
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        stop = true;
+    }
+    cond_task_available.notify_all();
+    for (std::thread &worker : workers) {
+        worker.join();
+    }
+}
+    
+__host__ inline void KittensClub::execute(std::function<void(int, cudaStream_t)> task) {
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        current_task = task;
+        for (size_t i = 0; i < task_available.size(); ++i)
+            task_available[i] = true;
+    }
+    cond_task_available.notify_all();
+    {
+        std::unique_lock<std::mutex> lock(mutex);
+        cond_task_done.wait(lock, [this] { return n_task_done == workers.size(); });
+        n_task_done = 0;
+    }
+}
+
+__host__ inline void KittensClub::worker(int worker_id, int device_id) {
+    cudaSetDevice(device_id); // done once and never again! This saves a LOT of time
+    while (true) {
+        std::function<void(int, cudaStream_t)> task;
+        {
+            std::unique_lock<std::mutex> lock(mutex);
+            cond_task_available.wait(lock, [this, worker_id] { return stop || task_available[worker_id]; });
+
+            if (stop)
+                return;
+
+            task = current_task;
+            task_available[worker_id] = false;
+        }
+        task(worker_id, streams[worker_id]);
+        {
+            std::lock_guard<std::mutex> lock(mutex); // adds about 10 microseconds overhead
+            ++n_task_done;
+            if (n_task_done == workers.size())
+                cond_task_done.notify_one();
+        }
+    }
+}

tinygrad_repo/extra/thunder/cuda/include/pyutils/parallel_tensor.cuh

@@ -0,0 +1,336 @@
+#pragma once
+
+#include <iostream>
+#include <map>
+#include <vector>
+
+#include <ATen/ops/from_blob.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/utils/pybind.h>
+
+#include "../types/device/vmm.cuh"
+#include "../types/device/ipc.cuh"
+#include "broker.cuh"
+
+namespace kittens {
+namespace py {
+
+/**
+ * @brief Distributed tensor wrapper for multi-GPU IPC sharing and multicast.
+ *        Can be later used for easy PGL creation right before a kernel call.
+ *        Meant to be used as a single object per thread/process.
+ */
+struct TKParallelTensor {
+    inline static std::map<std::pair<int, int>, KittensBroker> brokers_; // lazily initialized
+
+    at::Tensor data_; // for direct access from PyTorch
+    std::vector<int64_t> shape_;
+    at::ScalarType dtype_;
+
+    std::vector<void *> raw_ptrs_;
+    size_t allocated_size_;
+
+    int local_rank_; // identical to device index
+    int local_world_size_;
+
+    bool multicast_;
+    void *multicast_ptr_;
+    size_t multicast_allocated_size_;
+
+    detail::ipc::flavor ipc_flavor_;
+
+    __host__ inline TKParallelTensor(
+        const at::Tensor &tensor,
+        int local_rank,
+        int local_world_size,
+        bool multicast
+    ) : data_(tensor),
+        shape_(tensor.sizes().vec()),
+        dtype_(tensor.scalar_type()),
+        raw_ptrs_(local_world_size, nullptr),
+        allocated_size_(tensor.nbytes()),
+        local_rank_(local_rank),
+        local_world_size_(local_world_size),
+        multicast_(multicast),
+        multicast_ptr_(nullptr),
+        multicast_allocated_size_(0),
+        ipc_flavor_(detail::ipc::flavor::LEGACY) {
+
+        TORCH_CHECK(tensor.is_cuda(), "Tensor must be on CUDA device");
+        TORCH_CHECK(tensor.is_contiguous(), "Tensor must be contiguous");
+        TORCH_CHECK(tensor.dim() <= 4, "Only tensors with dim <= 4 are supported for TKParallelTensor");
+        TORCH_CHECK(tensor.device().index() == local_rank_, "Tensor device index must match local_rank");
+        TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative");
+        TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size");
+        TORCH_CHECK(!multicast, "Multicast is not supported for pre-allocated tensors");
+
+        brokers_.try_emplace(
+            {local_rank_, local_world_size_},
+            local_rank_, local_world_size_
+        );
+
+        if (brokers_.size() > 1)
+            std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl;
+
+        c10::cuda::CUDAGuard device_guard(local_rank_);
+        exchange_ipc_handles<detail::ipc::flavor::LEGACY>();
+    }
+
+    __host__ inline TKParallelTensor(
+        const std::vector<int64_t> &shape,
+        const at::ScalarType dtype,
+        int local_rank,
+        int local_world_size,
+        bool multicast
+    ) : shape_(shape),
+        dtype_(dtype),
+        raw_ptrs_(local_world_size, nullptr),
+        allocated_size_(0),
+        local_rank_(local_rank),
+        local_world_size_(local_world_size),
+        multicast_(multicast),
+        multicast_ptr_(nullptr),
+        multicast_allocated_size_(0),
+        ipc_flavor_(detail::ipc::flavor::VMM) {
+
+        TORCH_CHECK(local_rank_ >= 0, "local_rank must be non-negative");
+        TORCH_CHECK(local_rank_ < local_world_size_, "local_rank must be less than local_world_size");
+
+        brokers_.try_emplace(
+            {local_rank_, local_world_size_},
+            local_rank_, local_world_size_
+        );
+
+        if (brokers_.size() > 1)
+            std::cerr << "WARNING: 2 KittensBroker instances created in the same process. This is not safe." << std::endl;
+
+        c10::cuda::CUDAGuard device_guard(local_rank_);
+        create_shareable_cuda_tensor();
+        exchange_ipc_handles<detail::ipc::flavor::VMM>();
+
+        if (multicast_)
+            initialize_multicast();
+    }
+
+    TKParallelTensor(const TKParallelTensor&) = delete;
+    TKParallelTensor& operator=(const TKParallelTensor&) = delete;
+    TKParallelTensor& operator=(TKParallelTensor&& other) = delete;
+
+    __host__ inline TKParallelTensor(TKParallelTensor&& other) :
+        data_(std::move(other.data_)),
+        shape_(std::move(other.shape_)),
+        dtype_(std::move(other.dtype_)),
+        raw_ptrs_(std::move(other.raw_ptrs_)),
+        allocated_size_(other.allocated_size_),
+        local_rank_(other.local_rank_),
+        local_world_size_(other.local_world_size_),
+        multicast_(other.multicast_),
+        multicast_ptr_(other.multicast_ptr_),
+        multicast_allocated_size_(other.multicast_allocated_size_),
+        ipc_flavor_(other.ipc_flavor_) {
+        other.data_ = at::Tensor();
+        other.shape_.clear();
+        other.dtype_ = at::ScalarType::Undefined;
+        other.raw_ptrs_.clear();
+        other.allocated_size_ = 0;
+        other.local_rank_ = -1;
+        other.local_world_size_ = -1;
+        other.multicast_ = false;
+        other.multicast_ptr_ = nullptr;
+        other.multicast_allocated_size_ = 0;
+    }
+
+    __host__ inline ~TKParallelTensor() {
+        destroy();
+    }
+
+    __host__ inline at::Tensor data() const {
+        return data_;
+    }
+
+    __host__ inline void create_shareable_cuda_tensor() {
+        c10::cuda::CUDAGuard device_guard(local_rank_);
+
+        TORCH_CHECK(!shape_.empty(), "Shape must be non-empty");
+        TORCH_CHECK(shape_.size() <= 4, "Shape must have at most 4 dimensions for TKParallelTensor");
+        size_t size = c10::elementSize(dtype_);
+        for (auto dim : shape_) {
+            TORCH_CHECK(dim > 0, "Size dimensions must be positive");
+            size *= static_cast<size_t>(dim);
+        }
+
+        void *raw_ptr;
+        detail::vmm::vm_alloc_map_set_access(
+            &raw_ptr, &allocated_size_, size, local_rank_, local_world_size_);
+
+        // Create local copies for capture
+        int local_rank = local_rank_;
+        size_t allocated_size = allocated_size_;
+
+        auto deleter = [local_rank, raw_ptr, allocated_size](void* p) mutable {
+            if (!p) return;
+            c10::cuda::CUDAGuard device_guard(local_rank);
+            auto stream = c10::cuda::getCurrentCUDAStream().stream();
+            CUDACHECK(cudaStreamSynchronize(stream));
+            detail::vmm::vm_unmap(raw_ptr, allocated_size);
+        };
+
+        at::TensorOptions options = at::TensorOptions()
+            .dtype(dtype_)
+            .device(at::kCUDA, local_rank_);
+
+        data_ = at::from_blob(raw_ptr, shape_, std::move(deleter), options);
+    }
+
+    template <detail::ipc::flavor IPC_FLAVOR>
+    __host__ inline void exchange_ipc_handles() {
+        using handle_t = detail::ipc::handle<IPC_FLAVOR>;
+
+        // Get IPC handle
+        detail::ipc::check_support(local_rank_);
+        void *raw_ptr = reinterpret_cast<void *>(data_.data_ptr());
+        handle_t ipc_handle;
+        detail::ipc::export_handle(&ipc_handle, raw_ptr);
+
+        // Exchange IPC handles
+        std::vector<handle_t> all_ipc_handles(local_world_size_);
+        if constexpr (IPC_FLAVOR == detail::ipc::flavor::LEGACY) {
+            brokers_.at({local_rank_, local_world_size_}).exchange_data(
+                reinterpret_cast<void *>(all_ipc_handles.data()),
+                reinterpret_cast<void *>(&ipc_handle),
+                sizeof(handle_t)
+            );
+        } else if constexpr (IPC_FLAVOR == detail::ipc::flavor::VMM) {
+            brokers_.at({local_rank_, local_world_size_}).exchange_fds(
+                reinterpret_cast<int *>(all_ipc_handles.data()),
+                ipc_handle.handle_
+            );
+        } else {
+            throw std::runtime_error("Invalid IPC flavor");
+        }
+
+        // Import IPC handles
+        for (int i = 0; i < local_world_size_; i++) {
+            if (i == local_rank_)
+                raw_ptrs_[i] = raw_ptr;
+            else
+                detail::ipc::import_handle(&raw_ptrs_[i], all_ipc_handles[i], allocated_size_, local_world_size_);
+        }
+    }
+
+    __host__ inline void initialize_multicast() {
+        using handle_t = detail::ipc::handle<detail::ipc::flavor::VMM>;
+
+        detail::vmm::multicast_check(local_rank_);
+        detail::ipc::check_support(local_rank_);
+        detail::vmm::handle multicast_handle;
+
+        if (local_rank_ == 0) {
+            // Create multicast handle; only a single rank should create MC handle
+            detail::vmm::multicast_create_handle(
+                &multicast_handle,
+                &multicast_allocated_size_,
+                allocated_size_,
+                local_world_size_
+            );
+
+            // Currently, non-rank-0 path assumes allocated_size_ == multicast_allocated_size_
+            if (allocated_size_ != multicast_allocated_size_)
+                throw std::runtime_error("Multicast allocated size does not match memory allocated size");
+
+            // Get IPC handle
+            handle_t ipc_handle;
+            detail::ipc::export_handle(&ipc_handle, multicast_handle);
+
+            // Broadcast the IPC multicast handle
+            brokers_.at({local_rank_, local_world_size_}).broadcast_fd(nullptr, ipc_handle.handle_, 0);
+        } else {
+            // Receive the IPC multicast handle from rank 0
+            handle_t ipc_handle;
+            brokers_.at({local_rank_, local_world_size_}).broadcast_fd(&ipc_handle.handle_, -1, 0);
+            multicast_allocated_size_ = allocated_size_;
+            detail::ipc::import_handle(&multicast_handle, ipc_handle, multicast_allocated_size_, local_world_size_);
+        }
+
+        // Add all devices to the MC handle. Must sync
+        detail::vmm::multicast_bind_device(multicast_handle, local_rank_);
+        brokers_.at({local_rank_, local_world_size_}).sync(); // must ensure all devices are added
+
+        // Bind all memory to the MC handle and map to a virtual address; must be done after adding all devices
+        detail::vmm::handle memory_handle;
+        detail::vmm::vm_retrieve_handle(&memory_handle, raw_ptrs_[local_rank_]);
+        detail::vmm::multicast_bind_memory(multicast_handle, memory_handle, allocated_size_);
+        brokers_.at({local_rank_, local_world_size_}).sync();
+
+        // Map virtual address to multicast handle and set access; must be done after adding all devices
+        detail::vmm::vm_map(&multicast_ptr_, multicast_handle, multicast_allocated_size_);
+        detail::vmm::vm_set_access(multicast_ptr_, multicast_allocated_size_, local_world_size_);
+
+        // Free the handles immediately
+        detail::vmm::vm_free(multicast_handle);
+        detail::vmm::vm_free(memory_handle);
+    }
+
+    __host__ inline void destroy() {
+        // 1. Multicast cleanup
+        if (multicast_ && multicast_ptr_) {
+            brokers_.at({local_rank_, local_world_size_}).sync();
+            detail::vmm::handle multicast_handle;
+            detail::vmm::vm_retrieve_handle(&multicast_handle, multicast_ptr_);
+            detail::vmm::vm_unmap(multicast_ptr_, multicast_allocated_size_);
+            detail::vmm::multicast_unbind_device(multicast_handle, multicast_allocated_size_, local_rank_);
+            brokers_.at({local_rank_, local_world_size_}).sync();
+            detail::vmm::vm_free(multicast_handle);
+        }
+
+        // 2. Imported handle cleanup
+        for (int i = 0; i < local_world_size_; i++) {
+            if (i != local_rank_ && i < raw_ptrs_.size()) {
+                if (ipc_flavor_ == detail::ipc::flavor::LEGACY) {
+                    detail::ipc::free_handle<detail::ipc::flavor::LEGACY>(raw_ptrs_[i], allocated_size_);
+                } else if (ipc_flavor_ == detail::ipc::flavor::VMM) {
+                    detail::ipc::free_handle<detail::ipc::flavor::VMM>(raw_ptrs_[i], allocated_size_);
+                } else {
+                    throw std::runtime_error("Invalid IPC flavor");
+                }
+            }
+        }
+        brokers_.at({local_rank_, local_world_size_}).sync(); // must sync before destroying the tensor
+
+        // 3. Tensor cleanup
+        if (data_.defined())
+            data_.reset(); // properly decreases the ref count
+
+        // 4. Member variables cleanup
+        shape_.clear();
+        dtype_ = at::ScalarType::Undefined;
+        raw_ptrs_.clear();
+        allocated_size_ = 0;
+        local_rank_ = -1;
+        local_world_size_ = -1;
+        multicast_ = false;
+        multicast_ptr_ = nullptr;
+        multicast_allocated_size_ = 0;
+    }
+};
+
+} // namespace py
+} // namespace kittens
+
+#define BIND_TK_PARALLEL_TENSOR(m) \
+    pybind11::class_<kittens::py::TKParallelTensor>(m, "TKParallelTensor") \
+        .def(pybind11::init<const at::Tensor&, int, int, bool>(), \
+             pybind11::arg("tensor"), \
+             pybind11::arg("local_rank"), \
+             pybind11::arg("local_world_size"), \
+             pybind11::arg("multicast") = false) \
+        .def(pybind11::init<const std::vector<int64_t>&, const at::ScalarType&, int, int, bool>(), \
+             pybind11::arg("shape"), \
+             pybind11::arg("dtype"), \
+             pybind11::arg("local_rank"), \
+             pybind11::arg("local_world_size"), \
+             pybind11::arg("multicast") = false) \
+        .def("data", &kittens::py::TKParallelTensor::data) \
+        .def_readonly("data_", &kittens::py::TKParallelTensor::data_) \
+        .def_readonly("local_rank_", &kittens::py::TKParallelTensor::local_rank_) \
+        .def_readonly("local_world_size_", &kittens::py::TKParallelTensor::local_world_size_)

tinygrad_repo/extra/thunder/cuda/include/pyutils/pyutils.cuh

@@ -0,0 +1,235 @@
+#pragma once
+
+#include "util.cuh"
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h> // for automatic Python list -> std::vector conversion
+
+namespace kittens {
+namespace py {
+
+template<typename T> struct from_object {
+    static T make(pybind11::object obj) {
+        return obj.cast<T>();
+    }
+    static T unwrap(pybind11::object obj, int dev_idx) {
+        return make(obj); // Scalars should be passed in as a scalar
+    }
+};
+template<ducks::gl::all GL> struct from_object<GL> {
+    static GL make(pybind11::object obj) {
+        // Check if argument is a torch.Tensor
+        if (pybind11::hasattr(obj, "__class__") && 
+            obj.attr("__class__").attr("__name__").cast<std::string>() == "Tensor") {
+        
+            // Check if tensor is contiguous
+            if (!obj.attr("is_contiguous")().cast<bool>()) {
+                throw std::runtime_error("Tensor must be contiguous");
+            }
+            if (obj.attr("device").attr("type").cast<std::string>() == "cpu") {
+                throw std::runtime_error("Tensor must be on CUDA device");
+            }
+            
+            // Get shape, pad with 1s if needed
+            std::array<int, 4> shape = {1, 1, 1, 1};
+            auto py_shape = obj.attr("shape").cast<pybind11::tuple>();
+            size_t dims = py_shape.size();
+            if (dims > 4) {
+                throw std::runtime_error("Expected Tensor.ndim <= 4");
+            }
+            for (size_t i = 0; i < dims; ++i) {
+                shape[4 - dims + i] = pybind11::cast<int>(py_shape[i]);
+            }
+            
+            // Get data pointer using data_ptr()
+            uint64_t data_ptr = obj.attr("data_ptr")().cast<uint64_t>();
+            
+            // Create GL object using make_gl
+            return make_gl<GL>(data_ptr, shape[0], shape[1], shape[2], shape[3]);
+        }
+        throw std::runtime_error("Expected a torch.Tensor");
+    }
+    static GL unwrap(pybind11::object obj, int dev_idx) {
+        if (!pybind11::isinstance<pybind11::list>(obj))
+            throw std::runtime_error("GL unwrap expected a Python list.");
+        pybind11::list lst = pybind11::cast<pybind11::list>(obj);
+        if (dev_idx >= lst.size())
+            throw std::runtime_error("Device index out of bounds.");
+        return *lst[dev_idx].cast<std::shared_ptr<GL>>();
+    }
+};
+template<ducks::pgl::all PGL> struct from_object<PGL> {
+    static PGL make(pybind11::object obj) {
+        static_assert(!PGL::MULTICAST, "Multicast not yet supported on pyutils. Please initialize the multicast pointer manually.");
+        if (!pybind11::isinstance<pybind11::list>(obj))
+            throw std::runtime_error("PGL from_object expected a Python list.");
+        pybind11::list tensors = pybind11::cast<pybind11::list>(obj);
+        if (tensors.size() != PGL::num_devices)
+            throw std::runtime_error("Expected a list of " + std::to_string(PGL::num_devices) + " tensors");
+        std::array<int, 4> shape = {1, 1, 1, 1};
+        uint64_t data_ptrs[PGL::num_devices];
+        for (int i = 0; i < PGL::num_devices; i++) {
+            auto tensor = tensors[i];
+            if (!pybind11::hasattr(tensor, "__class__") || 
+                tensor.attr("__class__").attr("__name__").cast<std::string>() != "Tensor")
+                throw std::runtime_error("Expected a list of torch.Tensor");
+            if (!tensor.attr("is_contiguous")().cast<bool>())
+                throw std::runtime_error("Tensor must be contiguous");
+            if (tensor.attr("device").attr("type").cast<std::string>() == "cpu")
+                throw std::runtime_error("Tensor must be on CUDA device");
+            auto py_shape = tensor.attr("shape").cast<pybind11::tuple>();
+            size_t dims = py_shape.size();
+            if (dims > 4)
+                throw std::runtime_error("Expected Tensor.ndim <= 4");
+            for (size_t j = 0; j < dims; ++j) {
+                if (i == 0)
+                    shape[4 - dims + j] = pybind11::cast<int>(py_shape[j]);
+                else if (shape[4 - dims + j] != pybind11::cast<int>(py_shape[j]))
+                    throw std::runtime_error("All tensors must have the same shape");
+            }
+            data_ptrs[i] = tensor.attr("data_ptr")().cast<uint64_t>();
+        }
+        return make_pgl<PGL>(data_ptrs, shape[0], shape[1], shape[2], shape[3]);
+    }
+    static PGL unwrap(pybind11::object obj, int dev_idx) {
+        return *obj.cast<std::shared_ptr<PGL>>();
+    }
+};
+
+static std::unordered_set<std::string> registered;
+template<typename T> static void register_pyclass(pybind11::module &m) {
+    if constexpr (ducks::gl::all<T> || ducks::pgl::all<T>) {
+        std::string _typename = typeid(T).name();
+        if (registered.find(_typename) == registered.end()) {
+            pybind11::class_<T, std::shared_ptr<T>>(m, _typename.c_str());
+            registered.insert(_typename);
+        }
+    }
+}
+template<typename T> static pybind11::object multigpu_make(pybind11::object obj) {
+    if constexpr (ducks::gl::all<T>) {
+        if (!pybind11::isinstance<pybind11::list>(obj))
+            throw std::runtime_error("multigpu_make [GL] expected a Python list.");
+        pybind11::list lst = pybind11::cast<pybind11::list>(obj);
+        std::vector<std::shared_ptr<T>> gls;
+        for (int i = 0; i < lst.size(); i++)
+            gls.push_back(std::make_shared<T>(from_object<T>::make(lst[i])));
+        return pybind11::cast(gls);
+    } else if constexpr (ducks::pgl::all<T>) {
+        return pybind11::cast(std::make_shared<T>(from_object<T>::make(obj)));
+    } else {
+        return pybind11::cast(from_object<T>::make(obj));
+    }
+}
+
+template<typename T> concept has_dynamic_shared_memory = requires(T t) { { t.dynamic_shared_memory() } -> std::convertible_to<int>; };
+template<typename T> concept is_multigpu_globals = requires { 
+    { T::num_devices } -> std::convertible_to<std::size_t>;
+    { T::dev_idx } -> std::convertible_to<std::size_t>;
+} && T::num_devices >= 1;
+
+template<typename> struct trait;
+template<typename MT, typename T> struct trait<MT T::*> { using member_type = MT; using type = T; };
+template<typename> using object = pybind11::object;
+template<auto kernel, typename TGlobal> static void bind_kernel(auto m, auto name, auto TGlobal::*... member_ptrs) {
+    m.def(name, [](object<decltype(member_ptrs)>... args, pybind11::kwargs kwargs) {
+        TGlobal __g__ {from_object<typename trait<decltype(member_ptrs)>::member_type>::make(args)...};
+        cudaStream_t raw_stream = nullptr;
+        if (kwargs.contains("stream")) {
+            // Extract stream pointer
+            uintptr_t stream_ptr = kwargs["stream"].attr("cuda_stream").cast<uintptr_t>();
+            raw_stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+        }
+        if constexpr (has_dynamic_shared_memory<TGlobal>) {
+            int __dynamic_shared_memory__ = (int)__g__.dynamic_shared_memory();
+            cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, __dynamic_shared_memory__);
+            kernel<<<__g__.grid(), __g__.block(), __dynamic_shared_memory__, raw_stream>>>(__g__);
+        } else {
+            kernel<<<__g__.grid(), __g__.block(), 0, raw_stream>>>(__g__);
+        }
+    });
+}
+template<auto function, typename TGlobal> static void bind_function(auto m, auto name, auto TGlobal::*... member_ptrs) {
+    m.def(name, [](object<decltype(member_ptrs)>... args) {
+        TGlobal __g__ {from_object<typename trait<decltype(member_ptrs)>::member_type>::make(args)...};
+        function(__g__);
+    });
+}
+static void bind_multigpu_boilerplate(auto m) {
+    m.def("enable_all_p2p_access", [](const std::vector<int>& device_ids) {
+        int device_count;
+        CUDACHECK(cudaGetDeviceCount(&device_count));
+        if (device_count < device_ids.size())
+            throw std::runtime_error("Not enough CUDA devices available");
+        for (int i = 0; i < device_ids.size(); i++) {
+            CUDACHECK(cudaSetDevice(device_ids[i]));
+            for (int j = 0; j < device_ids.size(); j++) {
+                if (i == j) continue;
+                int can_access = 0;
+                CUDACHECK(cudaDeviceCanAccessPeer(&can_access, device_ids[i], device_ids[j]));
+                if (!can_access)
+                    throw std::runtime_error("Device " + std::to_string(device_ids[i]) + " cannot access device " + std::to_string(device_ids[j]));
+                cudaError_t res = cudaDeviceEnablePeerAccess(device_ids[j], 0);
+                if (res != cudaSuccess && res != cudaErrorPeerAccessAlreadyEnabled) {
+                    CUDACHECK(res);
+                }
+            }
+        }
+    });
+    pybind11::class_<KittensClub, std::shared_ptr<KittensClub>>(m, "KittensClub")
+        .def(pybind11::init([](const std::vector<int>& device_ids) {
+            int device_count;
+            CUDACHECK(cudaGetDeviceCount(&device_count));
+            if (device_count < device_ids.size())
+                throw std::runtime_error("Not enough CUDA devices available");
+            auto club = std::make_shared<KittensClub>(device_ids.data(), device_ids.size());
+            club->execute([&](int dev_idx, cudaStream_t stream) {}); // warmup
+            return club;
+        }), pybind11::arg("device_ids"))
+        .def(pybind11::init([](const std::vector<int>& device_ids, const std::vector<pybind11::object>& streams) {
+            int device_count;
+            CUDACHECK(cudaGetDeviceCount(&device_count));
+            if (device_count < device_ids.size())
+                throw std::runtime_error("Not enough CUDA devices available");
+            if (streams.size() != device_ids.size())
+                throw std::runtime_error("Number of streams must match number of devices");
+            
+            std::vector<cudaStream_t> raw_streams(streams.size());
+            for (size_t i = 0; i < streams.size(); ++i) {
+                uintptr_t stream_ptr = streams[i].attr("cuda_stream").cast<uintptr_t>();
+                raw_streams[i] = reinterpret_cast<cudaStream_t>(stream_ptr);
+            }
+            
+            auto club = std::make_shared<KittensClub>(device_ids.data(), raw_streams.data(), device_ids.size());
+            club->execute([&](int dev_idx, cudaStream_t stream) {}); // warmup
+            return club;
+        }), pybind11::arg("device_ids"), pybind11::arg("streams"));
+}
+template<auto kernel, typename TGlobal> static void bind_multigpu_kernel(auto m, auto name, auto TGlobal::*... member_ptrs) {
+    static_assert(is_multigpu_globals<TGlobal>, "Multigpu globals must have a member num_devices >= 1 and dev_idx");
+    (register_pyclass<typename trait<decltype(member_ptrs)>::member_type>(m), ...);
+    m.def((std::string("make_globals_")+name).c_str(), [](object<decltype(member_ptrs)>... args) -> std::vector<pybind11::object> {
+        return {multigpu_make<typename trait<decltype(member_ptrs)>::member_type>(args)...};
+    });
+    m.def(name, [](std::shared_ptr<KittensClub> club, object<decltype(member_ptrs)>... args) {
+        std::vector<TGlobal> __g__;
+        for (int i = 0; i < TGlobal::num_devices; i++) {
+            __g__.emplace_back(from_object<typename trait<decltype(member_ptrs)>::member_type>::unwrap(args, i)...);
+            __g__.back().dev_idx = i;
+        }
+        if constexpr (has_dynamic_shared_memory<TGlobal>) {
+            club->execute([&](int dev_idx, cudaStream_t stream) {
+                int __dynamic_shared_memory__ = (int)__g__[dev_idx].dynamic_shared_memory();
+                cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, __dynamic_shared_memory__);
+                kernel<<<__g__[dev_idx].grid(), __g__[dev_idx].block(), __dynamic_shared_memory__, stream>>>(__g__[dev_idx]);
+            });
+        } else {
+            club->execute([&](int dev_idx, cudaStream_t stream) {
+                kernel<<<__g__[dev_idx].grid(), __g__[dev_idx].block(), 0, stream>>>(__g__[dev_idx]);
+            });
+        }
+    });
+    // TODO: PGL destructor binding
+}
+
+} // namespace py
+} // namespace kittens

tinygrad_repo/extra/thunder/cuda/include/pyutils/torch_helpers.cuh

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <torch/extension.h>
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

tinygrad_repo/extra/thunder/cuda/include/pyutils/torchutils.cuh

@@ -0,0 +1,180 @@
+#pragma once
+
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/core/Tensor.h>
+
+#include "kittens.cuh"
+#include "parallel_tensor.cuh"
+
+namespace kittens {
+namespace py {
+
+template <typename Config>
+concept has_min_blocks_per_sm = requires { std::integral_constant<int, int(Config::MIN_BLOCKS_PER_SM)>{}; };
+
+template <typename Config>
+consteval int min_blocks_per_sm() {
+    if constexpr(has_min_blocks_per_sm<Config>)
+        return Config::MIN_BLOCKS_PER_SM;
+    else
+        return 1;
+}
+
+template <typename Config, typename Globals, auto Kernel>
+__global__
+__launch_bounds__(Config::NUM_THREADS, min_blocks_per_sm<Config>())
+void global_kernel_unclustered(const __grid_constant__ Globals G) {
+    Kernel(G);
+}
+
+template <typename Config, typename Globals, auto Kernel>
+__global__
+__launch_bounds__(Config::NUM_THREADS, min_blocks_per_sm<Config>())
+__cluster_dims__(Config::CLUSTER_SIZE)
+void global_kernel_clustered(const __grid_constant__ Globals G) {
+    Kernel(G);
+}
+
+template <typename Layout>
+static inline void tensor_check(const at::Tensor &t) {
+    TORCH_CHECK(t.is_cuda(), "Tensor must be on CUDA device")
+    TORCH_CHECK(t.is_contiguous(), "Tensor must be contiguous")
+    TORCH_CHECK(t.dim() <= 4, "Expected Tensor.dim() <= 4");
+
+    if constexpr (std::is_same_v<typename Layout::dtype, char>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Char, "Tensor has invalid dtype (expected int8)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, short>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Short, "Tensor has invalid dtype (expected int16)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, int>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Int, "Tensor has invalid dtype (expected int32)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, long>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Long, "Tensor has invalid dtype (expected int64)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, ::kittens::fp8e4m3>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e4m3fn, "Tensor has invalid dtype (expected fp8e4m3)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, ::kittens::fp8e5m2>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Float8_e5m2, "Tensor has invalid dtype (expected fp8e5m2)");
+#ifdef KITTENS_BLACKWELL
+    } else if constexpr (std::is_same_v<typename Layout::dtype, ::kittens::fp8e8m0>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Byte, "Tensor has invalid dtype (expected fp8e8m0 represented as uint8)");
+#endif
+    } else if constexpr (std::is_same_v<typename Layout::dtype, ::kittens::bf16>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::BFloat16, "Tensor has invalid dtype (expected bfloat16)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, ::kittens::half>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Half, "Tensor has invalid dtype (expected float16)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, float>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Float, "Tensor has invalid dtype (expected float32)");
+    } else if constexpr (std::is_same_v<typename Layout::dtype, double>) {
+        TORCH_CHECK(t.dtype() == at::ScalarType::Double, "Tensor has invalid dtype (expected float64)");
+    } else {
+        TORCH_CHECK(false, "Unsupported dtype");
+    }
+}
+
+template <kittens::ducks::pgl::all PGL>
+static inline void parallel_tensor_check(const TKParallelTensor& t) {
+    tensor_check<PGL>(t.data_);
+    TORCH_CHECK(t.data_.sizes().vec() == t.shape_, "Shape mismatch between TKParallelTensor and the underlying tensor");
+    TORCH_CHECK(t.data_.dtype() == t.dtype_, "Dtype mismatch between TKParallelTensor and the underlying tensor");
+    TORCH_CHECK(t.raw_ptrs_.size() == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor");
+    TORCH_CHECK(t.local_rank_ == t.data_.device().index(), "Current tensor device index mismatch within TKParallelTensor");
+    TORCH_CHECK(t.local_world_size_ == PGL::num_devices, "Number of devices mismatch between PGL and TKParallelTensor");
+    TORCH_CHECK(t.multicast_ == PGL::multicast, "Multicast mismatch between PGL and TKParallelTensor");
+    TORCH_CHECK(t.raw_ptrs_[t.local_rank_] == reinterpret_cast<void *>(t.data_.data_ptr()), "Current tensor data pointer not found in TKParallelTensor's raw_ptrs_");
+}
+
+template <kittens::ducks::gl::all GL>
+static inline GL tensor_to_gl(const at::Tensor &t) {
+    tensor_check<GL>(t);
+
+    std::array<int, 4> shape = {1, 1, 1, 1};
+    for (int i = 0; i < static_cast<int>(t.dim()); ++i)
+        shape[4 - t.dim() + i] = static_cast<int>(t.size(i));
+
+    uint64_t data_ptr = reinterpret_cast<uint64_t>(t.data_ptr());
+
+    return ::kittens::make_gl<GL>(data_ptr, shape[0], shape[1], shape[2], shape[3]);
+}
+
+template <kittens::ducks::pgl::all PGL>
+static inline PGL parallel_tensor_to_pgl(TKParallelTensor &t) {
+    parallel_tensor_check<PGL>(t);
+
+    std::array<int, 4> shape = {1, 1, 1, 1};
+    for (int i = 0; i < static_cast<int>(t.data_.dim()); ++i) {
+        shape[4 - t.data_.dim() + i] = static_cast<int>(t.data_.size(i));
+    }
+
+    if constexpr (PGL::multicast)
+        return ::kittens::make_pgl<PGL>(
+            reinterpret_cast<uint64_t>(t.multicast_ptr_), reinterpret_cast<uint64_t *>(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]);
+    else
+        return ::kittens::make_pgl<PGL>(
+            reinterpret_cast<uint64_t *>(t.raw_ptrs_.data()), shape[0], shape[1], shape[2], shape[3]);
+}
+
+template <kittens::ducks::gl::all GL>
+static inline GL make_fake_gl(const int batch, const int depth, const int rows, const int cols) {
+    return ::kittens::make_gl<GL>(reinterpret_cast<uint64_t>(nullptr), batch, depth, rows, cols);
+}
+
+static inline void _device_check(const at::Tensor& first, const at::Tensor& second) {
+    TORCH_CHECK(first.device() == second.device(), "All tensors must be on the same device");
+}
+
+template <typename T1, typename... Ts>
+static inline void device_check(const T1& first, const Ts&... rest) {
+    (_device_check(first, rest), ...);
+}
+
+static inline void _parallel_tensor_check(const TKParallelTensor& first, const TKParallelTensor& second) {
+    TORCH_CHECK(first.local_rank_ == second.local_rank_, "All parallel tensors must have the same local_rank");
+    TORCH_CHECK(first.local_world_size_ == second.local_world_size_, "All parallel tensors must have the same local_world_size");
+}
+
+template <typename T1, typename... Ts>
+static inline void parallel_tensor_check(const T1& first, const Ts&... rest) {
+    (_parallel_tensor_check(first, rest), ...);
+}
+
+template <typename Config>
+concept static_grid = requires { Config::NUM_BLOCKS; };
+
+template <typename Config>
+concept static_block = requires { Config::NUM_THREADS; };
+
+template <typename Config>
+concept static_dynamic_shared_memory = requires { Config::DYNAMIC_SHARED_MEMORY; };
+
+template <typename Config, typename Globals, auto Kernel>
+static inline void launch_kernel(const Globals &G) {
+    dim3 grid;
+    if constexpr (static_grid<Config>)
+        grid = dim3{Config::NUM_BLOCKS, 1, 1};
+    else
+        grid = G.grid();
+
+    dim3 block;
+    if constexpr (static_block<Config>)
+        block = dim3{Config::NUM_THREADS, 1, 1};
+    else
+        block = G.block();
+
+    int dynamic_shared_memory;
+    if constexpr (static_dynamic_shared_memory<Config>)
+        dynamic_shared_memory = static_cast<int>(Config::DYNAMIC_SHARED_MEMORY);
+    else
+        dynamic_shared_memory = G.dynamic_shared_memory();
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    if constexpr (Config::CLUSTER_SIZE <= 1) {
+        CUDACHECK(cudaFuncSetAttribute(global_kernel_unclustered<Config, Globals, Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_shared_memory));
+        global_kernel_unclustered<Config, Globals, Kernel><<<grid, block, dynamic_shared_memory, stream>>>(G);
+    } else {
+        CUDACHECK(cudaFuncSetAttribute(global_kernel_clustered<Config, Globals, Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_shared_memory));
+        global_kernel_clustered<Config, Globals, Kernel><<<grid, block, dynamic_shared_memory, stream>>>(G);
+    }
+}
+
+} // namespace py
+} // namespace kittens

tinygrad_repo/extra/thunder/cuda/include/pyutils/util.cuh

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../ops/ops.cuh"
+#include "club.cuh"
+#include <iostream>
+
+#define CHECK_CUDA_ERROR(val) check((val), #val, __FILE__, __LINE__)
+template <typename T>
+void check(T err, char const* const func, char const* const file,
+           int const line)
+{
+    if (err != cudaSuccess)
+    {
+        std::cerr << "CUDA Runtime Error at: " << file << ":" << line
+                  << std::endl;
+        std::cerr << cudaGetErrorString(err) << " " << func << std::endl;
+        //std::exit(EXIT_FAILURE);
+    }
+}