Move bmm calls to stable cshim

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## Latest Changes
 
+- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead.
+
 ### v0.7.0 (2026-09-10)
 **Added**: 
 - Public XLA FFI registration provider

openequivariance/CMakeLists.txt

@@ -117,7 +117,6 @@ endfunction()
 
 find_package(CUDAToolkit QUIET)
 find_package(hip QUIET)
-find_package(rocblas QUIET)
 
 if(CUDAToolkit_FOUND)
     message(STATUS "Building stable extension with CUDA backend.")
@@ -138,14 +137,14 @@ if(CUDAToolkit_FOUND)
         CUDA::cudart
         CUDA::cuda_driver
         CUDA::nvrtc
-        CUDA::cublas
         cuda_stub_lib
     )
     add_stable_extension(oeq_stable_cuda CUDA_BACKEND "${CUDA_LINK_LIBS}")
 endif()
 
 if(hip_FOUND)
     message(STATUS "Building stable extension with HIP backend.")
+    find_package(hiprtc REQUIRED)
 
     add_library(hip_stub_lib SHARED ${EXT_DIR}/stubs/stream.cpp)
 
@@ -159,16 +158,10 @@ if(hip_FOUND)
         CXX_STANDARD 17
     )
 
-    if(TARGET roc::rocblas)
-        set(HIP_BLAS_LIB roc::rocblas)
-    else()
-        set(HIP_BLAS_LIB rocblas)
-    endif()
-
     set(HIP_LINK_LIBS
         hip_stub_lib
 	    hip::host
-        ${HIP_BLAS_LIB}
+        hiprtc::hiprtc
     )
     add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}")
 endif()

openequivariance/openequivariance/_torch/extlib/__init__.py

@@ -112,7 +112,7 @@ def load_jit_extension():
                 ],
             )
         if torch.version.cuda:
-            extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-lcublas"])
+            extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-ltorch_cuda"])
 
             try:
                 torch_libs, cuda_libs = library_paths("cuda")
@@ -127,6 +127,8 @@ def load_jit_extension():
         elif torch.version.hip:
             torch_libs = library_paths("cuda")[0]
             extra_link_args.append("-Wl,-rpath," + torch_libs)
+            extra_link_args.extend("-L" + path for path in library_paths("cuda"))
+            extra_link_args.extend(["-ltorch_hip", "-lhiprtc"])
             extra_cflags.append("-DHIP_BACKEND")
 
         torch_sources = [oeq_root + "/extension/" + src for src in torch_sources]

openequivariance/openequivariance/extension/group_mm.hpp

@@ -1,141 +1,73 @@
 #pragma once
 
+#include <array>
 #include <cstdint>
 #include <memory>
 #include <stdexcept>
+#include <type_traits>
 
-#ifdef CUDA_BACKEND
-    #include "cublas_v2.h"
-    #include <cuda_runtime.h>
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h>
 
-    struct BlasHandle {
-        cublasHandle_t handle;
-        BlasHandle() {
-            if (cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS)
-                throw std::logic_error("CUBLAS initialization failed");
-        }
-        ~BlasHandle() { cublasDestroy(handle); }
-    };
-#elif defined(HIP_BACKEND)
-    #include "rocblas/rocblas.h"
-    #include <hip/hip_runtime.h>
-
-    struct BlasHandle {
-        rocblas_handle handle;
-        BlasHandle() {
-            if (rocblas_create_handle(&handle) != rocblas_status_success)
-                throw std::logic_error("rocBLAS initialization failed");
-        }
-        ~BlasHandle() { rocblas_destroy_handle(handle); }
-    };
-#endif
+namespace oeq {
 
-inline BlasHandle& get_blas_handle() {
-    static BlasHandle handle;
-    return handle;
+inline void check_group_mm_shim(AOTITorchError status) {
+    if (status != AOTI_TORCH_SUCCESS)
+        throw std::runtime_error("group_gemm: PyTorch C shim failed");
 }
 
-template<typename T>
-void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
-        int64_t* ragged_counts, int num_W, int batch_size, int m, int k, int ragged_inner) {
+using GroupMMTensor = std::unique_ptr<
+    std::remove_pointer_t<AtenTensorHandle>,
+    decltype(&aoti_torch_delete_tensor_object)>;
+
+inline GroupMMTensor group_mm_view(
+        AtenTensorHandle tensor, std::array<int64_t, 3> sizes,
+        std::array<int64_t, 3> strides, int64_t offset) {
+    AtenTensorHandle view = nullptr;
+    check_group_mm_shim(aoti_torch__reinterpret_tensor(
+        tensor, 3, sizes.data(), strides.data(), offset, &view));
+    return GroupMMTensor(view, aoti_torch_delete_tensor_object);
+}
 
-    auto& blas = get_blas_handle();
-    T alpha = 1.0, beta = 0.0;
-    T* A_base = reinterpret_cast<T*>(A_raw);
-    T* B_base = reinterpret_cast<T*>(B_raw);
-    T* C_base = reinterpret_cast<T*>(C_raw);
+inline void group_gemm_torch(
+        AtenTensorHandle A, AtenTensorHandle B, AtenTensorHandle C,
+        const int64_t* ragged_counts, int64_t num_W, int64_t batch_size,
+        int64_t m, int64_t k, int64_t ragged_inner) {
+    if (batch_size == 0 || m == 0 || k == 0)
+        return;
 
-    int64_t ragged_offset = 0;
-    for (int i = 0; i < num_W; i++) {
-        int M, K, N, lda, ldb, ldc, strideA, strideB, strideC;
-        T *A, *B, *C;
-#ifdef CUDA_BACKEND
-        cublasOperation_t transa, transb;
-#elif defined(HIP_BACKEND)
-        rocblas_operation transa, transb;
-#endif
+    int64_t offset = 0;
+    for (int64_t i = 0; i < num_W; ++i) {
+        const int64_t n = ragged_counts[i];
+        if (n == 0)
+            continue;
 
         if (ragged_inner == 0) {
-            M = m; K = k; N = static_cast<int>(ragged_counts[i]);
-            A = A_base + (m * k * batch_size * i);
-            lda = k; strideA = M * K;
-            B = B_base + (k * batch_size * ragged_offset);
-            ldb = K * batch_size; strideB = K;
-            C = C_base + (m * batch_size * ragged_offset);
-            ldc = M * batch_size; strideC = M;
-#ifdef CUDA_BACKEND
-            transa = CUBLAS_OP_T; transb = CUBLAS_OP_N;
-#elif defined(HIP_BACKEND)
-            transa = rocblas_operation_transpose; transb = rocblas_operation_none;
-#endif
+            auto input = group_mm_view(B,
+                {batch_size, n, k}, {k, batch_size * k, 1},
+                offset * batch_size * k);
+            auto weight = group_mm_view(A,
+                {batch_size, k, m}, {m * k, 1, k},
+                i * batch_size * m * k);
+            auto output = group_mm_view(C,
+                {batch_size, n, m}, {m, batch_size * m, 1},
+                offset * batch_size * m);
+            check_group_mm_shim(aoti_torch_cuda_bmm_out(
+                output.get(), input.get(), weight.get()));
         } else {
-            M = k; K = static_cast<int>(ragged_counts[i]); N = m;
-            A = B_base + (k * batch_size * ragged_offset);
-            lda = k * batch_size; strideA = M;
-            B = A_base + (m * batch_size * ragged_offset);
-            ldb = m * batch_size; strideB = N;
-            C = C_base + (m * k * batch_size * i);
-            ldc = k; strideC = M * N;
-#ifdef CUDA_BACKEND
-            transa = CUBLAS_OP_N; transb = CUBLAS_OP_T;
-#elif defined(HIP_BACKEND)
-            transa = rocblas_operation_none; transb = rocblas_operation_transpose;
-#endif
-        }
-        ragged_offset += ragged_counts[i];
-
-        if (ragged_counts[i] > 0) {
-#ifdef CUDA_BACKEND
-            cublasStatus_t stat;
-            if (std::is_same<T, float>::value) {
-                stat = cublasSgemmStridedBatched(blas.handle,
-                    transa, transb, M, N, K,
-                    reinterpret_cast<float*>(&alpha),
-                    reinterpret_cast<float*>(A), lda, strideA,
-                    reinterpret_cast<float*>(B), ldb, strideB,
-                    reinterpret_cast<float*>(&beta),
-                    reinterpret_cast<float*>(C), ldc, strideC,
-                    batch_size);
-            } else if (std::is_same<T, double>::value) {
-                stat = cublasDgemmStridedBatched(blas.handle,
-                    transa, transb, M, N, K,
-                    reinterpret_cast<double*>(&alpha),
-                    reinterpret_cast<double*>(A), lda, strideA,
-                    reinterpret_cast<double*>(B), ldb, strideB,
-                    reinterpret_cast<double*>(&beta),
-                    reinterpret_cast<double*>(C), ldc, strideC,
-                    batch_size);
-            } else {
-                throw std::logic_error("Unsupported datatype for grouped GEMM!");
-            }
-            if (stat != CUBLAS_STATUS_SUCCESS)
-                throw std::logic_error("Grouped GEMM failed!");
-#elif defined(HIP_BACKEND)
-            rocblas_status stat;
-            if (std::is_same<T, float>::value) {
-                stat = rocblas_sgemm_strided_batched(blas.handle,
-                    transa, transb, M, N, K,
-                    reinterpret_cast<float*>(&alpha),
-                    reinterpret_cast<float*>(A), lda, strideA,
-                    reinterpret_cast<float*>(B), ldb, strideB,
-                    reinterpret_cast<float*>(&beta),
-                    reinterpret_cast<float*>(C), ldc, strideC,
-                    batch_size);
-            } else if (std::is_same<T, double>::value) {
-                stat = rocblas_dgemm_strided_batched(blas.handle,
-                    transa, transb, M, N, K,
-                    reinterpret_cast<double*>(&alpha),
-                    reinterpret_cast<double*>(A), lda, strideA,
-                    reinterpret_cast<double*>(B), ldb, strideB,
-                    reinterpret_cast<double*>(&beta),
-                    reinterpret_cast<double*>(C), ldc, strideC,
-                    batch_size);
-            } else {
-                throw std::logic_error("Unsupported datatype for grouped GEMM!");
-            }
-            if (stat != rocblas_status_success)
-                throw std::logic_error("Grouped GEMM failed!");
-#endif
+            auto left = group_mm_view(A,
+                {batch_size, m, n}, {m, 1, batch_size * m},
+                offset * batch_size * m);
+            auto right = group_mm_view(B,
+                {batch_size, n, k}, {k, batch_size * k, 1},
+                offset * batch_size * k);
+            auto output = group_mm_view(C,
+                {batch_size, m, k}, {m * k, k, 1},
+                i * batch_size * m * k);
+            check_group_mm_shim(aoti_torch_cuda_bmm_out(
+                output.get(), left.get(), right.get()));
         }
+        offset += n;
     }
 }
+
+}

openequivariance/openequivariance/extension/libtorch_tp_jit.cpp

@@ -2,17 +2,19 @@
 #include <pybind11/pybind11.h>
 
 #ifdef CUDA_BACKEND
-    #include <ATen/cuda/CUDAContext.h>
+    #include <c10/cuda/CUDAStream.h>
 #endif
 
 #ifdef HIP_BACKEND
     #include <c10/hip/HIPStream.h>
 #endif
 
 #include <ATen/Operators.h>
+#include <c10/core/DeviceGuard.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
 #include <torch/all.h>
+#include <torch/csrc/inductor/aoti_torch/utils.h>
 #include <torch/library.h>
 
 using Tensor = torch::Tensor;
@@ -29,6 +31,20 @@ constexpr Dtype kByte = torch::kByte;
 #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL
 #define REGISTER_LIBRARY TORCH_LIBRARY
 
+namespace {
+
+class TensorDeviceGuard {
+    c10::DeviceGuard guard;
+public:
+    explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.device()) {}
+};
+
+AtenTensorHandle tensor_handle(Tensor& tensor) {
+    return torch::aot_inductor::tensor_pointer_to_tensor_handle(&tensor);
+}
+
+}
+
 #include "torch_core.hpp"
 
 Tensor tensor_to_cpu_contiguous(const Tensor &tensor) {

openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp

@@ -27,6 +27,20 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte;
 #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL
 #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY
 
+namespace {
+
+class TensorDeviceGuard {
+    torch::stable::accelerator::DeviceGuard guard;
+public:
+    explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.get_device_index()) {}
+};
+
+AtenTensorHandle tensor_handle(Tensor& tensor) {
+    return tensor.get();
+}
+
+}
+
 #include "torch_core.hpp"
 
 Tensor tensor_to_cpu_contiguous(const Tensor &tensor) {

openequivariance/openequivariance/extension/stubs/stream.cpp

@@ -1,8 +1,15 @@
+#define USE_CUDA
+
 #include <cstdint>
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h>
 
 extern "C" {
     AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) {
-        return 0;
+        return AOTI_TORCH_FAILURE;
+    }
+
+    AOTITorchError aoti_torch_cuda_bmm_out(
+            AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2) {
+        return AOTI_TORCH_FAILURE;
     }
-}
+}