From 69463cb4c4c606b3a929f3d7b372ef098f1f2f84 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:43:55 -0700 Subject: [PATCH 01/18] Add stable C shim BMM prototype and migration design --- sandbox/cshim_group_gemm/DESIGN.md | 292 ++++++++++++++++++++++++ sandbox/cshim_group_gemm/group_gemm.cpp | 99 ++++++++ sandbox/cshim_group_gemm/smoke_test.py | 71 ++++++ 3 files changed, 462 insertions(+) create mode 100644 sandbox/cshim_group_gemm/DESIGN.md create mode 100644 sandbox/cshim_group_gemm/group_gemm.cpp create mode 100644 sandbox/cshim_group_gemm/smoke_test.py diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md new file mode 100644 index 00000000..3b919e07 --- /dev/null +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -0,0 +1,292 @@ +# Move Torch cuBLAS calls to the stable C shim + +Replace OEQ's direct CUDA and ROCm BLAS calls for `libtorch_tp_jit::group_gemm` with PyTorch's `aoti_torch_cuda_bmm_out` C API. Keep the existing grouped operation, tensor layouts, output allocation, and custom autograd rules. PyTorch will own the underlying BLAS handles, streams, library calls, and backend selection. + +Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). + +## Problem and scope + +OEQ's grouped GEMM is a host loop over nonempty ragged groups, issuing one strided batched GEMM per group. It does not require a vendor's heterogeneous grouped-GEMM API. The only direct BLAS computation calls found are float32/float64 strided batched GEMMs in `extension/group_mm.hpp`, called from the Torch backend in `extension/torch_core.hpp`. + +Using PyTorch's stable ABI to obtain a cuBLAS handle does not make that handle interchangeable with a separately loaded cuBLAS implementation. Earlier H100 experiments passed with matching handle/function owners, while some foreign-owner combinations returned incorrect streams, failed, or crashed. Successful version queries did not establish handle compatibility. Moving the GEMM itself behind PyTorch's C ABI removes this borrowed-handle boundary. + +The current checkout at `dbe854415da7771eba33195534c171adbca5677b` creates its own static BLAS handle; the inspected PR snapshot at `b551d7469db6d7ac688859cec46a3bb2b1c71a5b` borrows Torch's handle. The replacement removes either form of OEQ-side handle management. The final patch must be based on the branch being merged. + +The stable wheels currently target PyTorch 2.10. Use official stable C declarations, with no dependency on ATen's C++ ABI in that build. The C entry point was also exercised successfully on the available PyTorch 2.7 installation; that observation does not establish support for every older Torch release. [PyTorch stable ABI documentation](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). + +## Public interface + +Include `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`, which declares: + +```cpp +AOTITorchError aoti_torch_cuda_bmm_out( + AtenTensorHandle out, + AtenTensorHandle self, + AtenTensorHandle mat2); +``` + +Use this generated API rather than the deprecated `aoti_torch_bmm_out` spelling. The API takes PyTorch tensor handles, not cuBLAS/hipBLAS handles. OEQ does not call a BLAS version getter or select a vendor library at runtime. [PyTorch 2.10 declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74). + +Add a Torch-specific helper, such as `group_mm_torch.hpp`, taking the existing input/output pointers, dimensions, CPU ragged counts, dtype, and explicit device index. Both Torch extension variants should use it where their supported Torch versions expose the required C symbols. Keep this helper outside the framework-independent CUDA/HIP kernel backend. + +Create temporary, non-owning tensor views with `aoti_torch_create_tensor_from_blob_v2`. Specify sizes and strides in elements, storage offset zero, and the actual input device. Release each temporary tensor handle with `aoti_torch_delete_tensor_object` through RAII, including error paths. Releasing metadata must not release the caller's input or output storage. [PyTorch's blob implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L525). + +## Layout mapping + +Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Pointer offsets and strides below are measured in elements. + +| Mode | View | Base pointer | Shape | Strides | +| --- | --- | --- | --- | --- | +| `ragged_inner == 0` | Left / input | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| `ragged_inner == 0` | Right / weights | `A + i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | +| `ragged_inner == 0` | Output | `C + o*b*m` | `[b,n,m]` | `[m,b*m,1]` | +| `ragged_inner == 1` | Left | `A + o*b*m` | `[b,m,n]` | `[m,1,b*m]` | +| `ragged_inner == 1` | Right | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| `ragged_inner == 1` | Output | `C + i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | + +Call BMM-out once for each nonempty group. Preserve the existing zero-initialized output and skip empty groups, including the all-empty case. This matters for the weight-gradient output blocks in mode 1. Use 64-bit dimensions and offsets throughout instead of the current narrowing casts to `int`. + +These views preserve OEQ's interleaved batch layout without explicitly transposing or copying the underlying buffers. PyTorch can handle BLAS-compatible strides directly, but some layouts can trigger its internal copy path; avoiding all copies is not an unconditional API guarantee. [BMM layout handling](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). + +## Device, stream, and autograd behavior + +Guard the inputs' actual GPU device before making contiguous copies, allocating the output, constructing views, or invoking BMM. Restore the caller's device on return. The 2.10 stable implementation can use `aoti_torch_create_device_guard` / `aoti_torch_delete_device_guard`, whose implementation selects Torch's active accelerator; this avoids requiring a CUDA-specific guard implementation in the shared helper. The existing standalone prototype uses the older CUDA-named guard because it was tested on Torch 2.7. [Generic guard implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L1468). + +Use Torch's current stream for that device. Do not set a BLAS stream, create a stream, or synchronize inside the operator. Original tensors must remain valid through launch; asynchronous storage lifetime follows the usual PyTorch current-stream contract. No global or cached raw pointers or tensor views are needed. + +Validate the raw-pointer helper's preconditions at the tensor boundary: same GPU device and dtype, supported float32/float64 types, contiguous CPU int64 counts, valid counts length and nonnegative values, compatible input shapes, and a valid mode. Check that counts describe the available rows before constructing views. Keep the existing custom operator schema, fake implementation, and backward formulas. BMM-out is an internal implementation detail; the custom autograd registration remains responsible for derivatives. + +PyTorch's BMM implementation uses alpha 1 and beta 0, matching the existing GEMMs. It may choose a different BLAS backend or kernel and apply Torch's precision/determinism settings, so require numerical agreement rather than bitwise equivalence. [BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L680). + +## Why this should work on ROCm + +Source inspection of PyTorch **v2.10.0** supports using the same generated function on both GPU platforms: + +1. PyTorch's build defines `GENERATED_CXX_TORCH_CUDA` to contain `c_shim_cuda.cpp`, then explicitly adds that source to `torch_hip` under `USE_ROCM`. The generated shim is therefore part of the ROCm library too. [Generated source definition](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L331), [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941). +2. The shim generator names the API using the `CUDA` dispatch key and generates a call into the corresponding ATen backend. The exported name remains `aoti_torch_cuda_bmm_out`; do not invent an `aoti_torch_hip_bmm_out` symbol. [Shim generator](https://github.com/pytorch/pytorch/blob/v2.10.0/torchgen/gen_aoti_c_shim.py#L492). +3. BMM reaches Torch's GEMM/batched GEMM implementation. Its HIP conversion maps `cublasSgemmStridedBatched` and `cublasDgemmStridedBatched` to their `hipblas` equivalents. Torch also contains ROCm-specific backend selection, including a double-precision fallback when hipBLASLt cannot handle the operation. [HIP mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [float/double backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). + +Torch's ROCm tensors use the `cuda` device interface. The helper should use Torch's CUDA device-type value for both these Torch builds, rather than confusing that convention with a distinct external HIP tensor type. [ROCm semantics](https://docs.pytorch.org/docs/main/notes/hip.html). + +This is **source-level evidence**, not an AMD hardware test or a binary export audit of a ROCm wheel. Verify the exports in the supported wheel and run one AMD smoke case before marking ROCm complete. + +## Build and packaging changes + +- Remove direct cuBLAS/rocBLAS includes, calls, and handle management from the Torch grouped-GEMM path. On the PR branch, also remove the now-unused borrowed-handle adapters. +- Remove `CUDA::cublas`, the grouped-GEMM-only `find_package(rocblas)` / `${HIP_BLAS_LIB}`, and the JIT loader's explicit `-lcublas` link flag after confirming there are no remaining consumers. Torch still brings its own BLAS dependencies. +- Resolve the generated BMM symbol through the installed Torch GPU library: `libtorch_cuda` for CUDA and `libtorch_hip` for ROCm. Common tensor/guard C APIs are supplied by Torch's common library. +- OEQ's wheel build currently downloads CPU LibTorch headers/libraries and creates GPU link stubs from `extension/stubs/stream.cpp`. Extend those build-only stubs with the exact generated BMM declaration/definition, sharing the header so signature drift causes a build error. If retaining the prototype's GPU-specific guards, add their GPU exports too; the proposed 2.10 generic guard avoids that need. Never package or execute the fake implementations. Check that installed extensions resolve against the real Torch libraries. +- Apply changes to both the Python extension and the AOTI shared library targets, and to source/JIT builds. The existing HIP CMake target is named `torch_stable_hip` while import/install expectations use `oeq_stable_hip`; resolve that naming mismatch as part of making the HIP artifact load correctly. +- Keep the existing CUDA/HIP runtime and runtime-compiler dependencies needed by OEQ-generated kernels. Removing the BLAS dependency does not make the entire extension independent of the GPU platform. + +The stable-wheel baseline remains 2.10. Confirm the supported source/JIT baseline before sharing every C helper with that build: the prototype's helpers were exercised on 2.7, whereas the proposed generic guard needs its own baseline check. If older JIT support must be preserved, a JIT-only adapter using that installed Torch's ordinary BMM-out/device guard is acceptable; it still eliminates direct vendor BLAS calls and does not enter the stable wheel. + +## JAX scope + +Neither the inspected PR snapshot nor the current JAX extension calls `group_gemm_blas`; its CMake target links the runtime, driver, and NVRTC, without BLAS. This replacement therefore remains entirely within Torch. + +JAX's public FFI has buffers and a GPU stream, but no equivalent BMM shim or public BLAS-handle getter. If a future JAX operation needs this GEMM, express it in the compiled JAX graph where shapes permit, or provide a separate FFI implementation whose BLAS handles and calls come from the same library. [JAX GPU FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu), [public XLA FFI API](https://github.com/openxla/xla/blob/main/xla/ffi/api/c_api.h#L757). + +## Existing prototype and validation + +A standalone C++ prototype already implements both layout branches using only Torch C APIs. Its single float32 forward smoke case passed on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against an independent CPU float64 reference. + +Earlier ctypes calls to the same C shim passed 40 cases spanning both dtypes and modes. Four selected profiles showed one GEMM kernel with no observed copy or GPU allocation inside the BMM call. These measurements do not cover the CPU cost of temporary tensor metadata or establish end-to-end performance. + +The integrated extension, its autograd registration, the final wheel link setup, and ROCm execution remain to be validated. No additional GPU experiments were run for this design. + +## Acceptance criteria + +- Both grouped-GEMM modes use Torch BMM, with no direct vendor BLAS symbols or handles in OEQ's Torch extension. +- Existing layouts, empty groups, float32/float64 behavior, operator schema, fake implementation, and gradients are preserved. +- Device guarding and a nondefault current stream work; the caller's current device is restored. +- Stable CUDA and HIP artifacts resolve the official C symbols in the real installed Torch GPU library; link stubs do not ship. Audit direct dynamic dependencies to confirm BLAS is now Torch's responsibility. +- Run focused correctness checks for both modes/dtypes and one relevant autograd case during integration, plus one ROCm smoke case when AMD hardware is available. Keep this bounded; no exhaustive vendor-version matrix is needed to establish the prototype. +- Check representative overhead before claiming performance parity. Graph capture and broader multi-device coverage should be checked where required by the supported operator contract. + +## Prototype code + +The following is the existing standalone CUDA prototype, not the final production registration or generic-device-guard adaptation. The caller supplies validated buffers and zero-initializes outputs where empty groups need to remain zero. + +```cpp +// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +void check(AOTITorchError status) { + if (status != AOTI_TORCH_SUCCESS) + throw std::runtime_error("PyTorch C shim failed"); +} + +using Tensor = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_tensor_object)>; +using DeviceGuard = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_cuda_guard)>; + +// Only tensor metadata is created; the caller retains ownership of the buffer. +Tensor view(void* data, std::array sizes, + std::array strides, int32_t dtype, int32_t device) { + AtenTensorHandle tensor = nullptr; + check(aoti_torch_create_tensor_from_blob_v2( + data, 3, sizes.data(), strides.data(), 0, dtype, + aoti_torch_device_type_cuda(), device, &tensor, + aoti_torch_layout_strided(), nullptr, 0)); + return Tensor(tensor, aoti_torch_delete_tensor_object); +} + +} // namespace + +template +void group_gemm_cshim( + T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, + int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { + static_assert(std::is_same_v || std::is_same_v); + const int32_t dtype = std::is_same_v + ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); + + CUDAGuardHandle raw_guard = nullptr; + check(aoti_torch_create_cuda_guard(device, &raw_guard)); + DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); + + int64_t offset = 0; + for (int i = 0; i < num_groups; ++i) { + const int64_t n = ragged_counts[i]; + if (n == 0) continue; // Preserve the original empty-group behavior. + + if (ragged_inner == 0) { + // [batch, n, k] @ [batch, k, m] -> [batch, n, m] + auto input = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto weight = view(A + i * batch * m * k, + {batch, k, m}, {m * k, 1, k}, dtype, device); + auto output = view(C + offset * batch * m, + {batch, n, m}, {m, batch * m, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); + } else { + // [batch, m, n] @ [batch, n, k] -> [batch, m, k] + auto left = view(A + offset * batch * m, + {batch, m, n}, {m, 1, batch * m}, dtype, device); + auto right = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto output = view(C + i * batch * m * k, + {batch, m, k}, {m * k, k, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); + } + offset += n; + } +} + +// Small ctypes entry point for the smoke test, not production registration code. +extern "C" int oeq_group_gemm_cshim( + int dtype, void* A, void* B, void* C, const int64_t* counts, + int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { + try { + if (dtype == 0) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else if (dtype == 1) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else { + throw std::runtime_error("Expected float32 or float64"); + } + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "%s\n", error.what()); + return 1; + } +} +``` + +
+Reproduce the single CUDA smoke case + +Save the C++ above as `group_gemm.cpp` and this script beside it as `smoke_test.py`, then run `python3 smoke_test.py` with a CUDA-enabled Torch installation and a C++ compiler. The script builds only the standalone prototype and executes one configuration. + +```python +"""Compile the C++ prototype and run one float32 CUDA correctness case.""" +import ctypes +from pathlib import Path +import subprocess + +import torch + + +def main(): + root = Path(__file__).resolve().parent + torch_root = Path(torch.__file__).resolve().parent + library = root / "group_gemm.so" + subprocess.run( + [ + "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", + f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), + f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", + "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", + "-o", str(library), + ], + check=True, + ) + lib = ctypes.CDLL(str(library)) + call = lib.oeq_group_gemm_cshim + i64 = ctypes.c_int64 + call.argtypes = [ + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, + ctypes.c_int, ctypes.c_int32, + ] + call.restype = ctypes.c_int + + # One case includes differently sized groups, an empty group, interleaved + # batches, and a nondefault stream. The reference runs on the CPU. + counts = [2, 0, 5, 1] + batch, m, k = 3, 4, 5 + torch.manual_seed(123) + weights_cpu = torch.randn(len(counts), batch, m, k) + input_cpu = torch.randn(sum(counts), batch, k) + expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) + offset = 0 + for i, n in enumerate(counts): + expected[offset:offset + n] = torch.einsum( + "bmk,nbk->nbm", weights_cpu[i].double(), + input_cpu[offset:offset + n].double(), + ) + offset += n + + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + weights = weights_cpu.cuda() + inputs = input_cpu.cuda() + output = torch.full(expected.shape, float("nan"), device="cuda") + status = call( + 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), + (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, + ) + assert status == 0, f"C shim prototype returned {status}" + stream.synchronize() + actual = output.cpu().double() + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + print( + f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " + f"max absolute error={(actual - expected).abs().max().item():.3g}; " + f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" + ) + + +if __name__ == "__main__": + main() +``` + +
diff --git a/sandbox/cshim_group_gemm/group_gemm.cpp b/sandbox/cshim_group_gemm/group_gemm.cpp new file mode 100644 index 00000000..1f7462a7 --- /dev/null +++ b/sandbox/cshim_group_gemm/group_gemm.cpp @@ -0,0 +1,99 @@ +// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +void check(AOTITorchError status) { + if (status != AOTI_TORCH_SUCCESS) + throw std::runtime_error("PyTorch C shim failed"); +} + +using Tensor = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_tensor_object)>; +using DeviceGuard = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_cuda_guard)>; + +// Only tensor metadata is created; the caller retains ownership of the buffer. +Tensor view(void* data, std::array sizes, + std::array strides, int32_t dtype, int32_t device) { + AtenTensorHandle tensor = nullptr; + check(aoti_torch_create_tensor_from_blob_v2( + data, 3, sizes.data(), strides.data(), 0, dtype, + aoti_torch_device_type_cuda(), device, &tensor, + aoti_torch_layout_strided(), nullptr, 0)); + return Tensor(tensor, aoti_torch_delete_tensor_object); +} + +} // namespace + +template +void group_gemm_cshim( + T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, + int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { + static_assert(std::is_same_v || std::is_same_v); + const int32_t dtype = std::is_same_v + ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); + + CUDAGuardHandle raw_guard = nullptr; + check(aoti_torch_create_cuda_guard(device, &raw_guard)); + DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); + + int64_t offset = 0; + for (int i = 0; i < num_groups; ++i) { + const int64_t n = ragged_counts[i]; + if (n == 0) continue; // Preserve the original empty-group behavior. + + if (ragged_inner == 0) { + // [batch, n, k] @ [batch, k, m] -> [batch, n, m] + auto input = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto weight = view(A + i * batch * m * k, + {batch, k, m}, {m * k, 1, k}, dtype, device); + auto output = view(C + offset * batch * m, + {batch, n, m}, {m, batch * m, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); + } else { + // [batch, m, n] @ [batch, n, k] -> [batch, m, k] + auto left = view(A + offset * batch * m, + {batch, m, n}, {m, 1, batch * m}, dtype, device); + auto right = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto output = view(C + i * batch * m * k, + {batch, m, k}, {m * k, k, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); + } + offset += n; + } +} + +// Small ctypes entry point for the smoke test, not production registration code. +extern "C" int oeq_group_gemm_cshim( + int dtype, void* A, void* B, void* C, const int64_t* counts, + int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { + try { + if (dtype == 0) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else if (dtype == 1) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else { + throw std::runtime_error("Expected float32 or float64"); + } + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "%s\n", error.what()); + return 1; + } +} diff --git a/sandbox/cshim_group_gemm/smoke_test.py b/sandbox/cshim_group_gemm/smoke_test.py new file mode 100644 index 00000000..931d9698 --- /dev/null +++ b/sandbox/cshim_group_gemm/smoke_test.py @@ -0,0 +1,71 @@ +"""Compile the C++ prototype and run one float32 CUDA correctness case.""" +import ctypes +from pathlib import Path +import subprocess + +import torch + + +def main(): + root = Path(__file__).resolve().parent + torch_root = Path(torch.__file__).resolve().parent + library = root / "group_gemm.so" + subprocess.run( + [ + "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", + f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), + f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", + "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", + "-o", str(library), + ], + check=True, + ) + lib = ctypes.CDLL(str(library)) + call = lib.oeq_group_gemm_cshim + i64 = ctypes.c_int64 + call.argtypes = [ + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, + ctypes.c_int, ctypes.c_int32, + ] + call.restype = ctypes.c_int + + # One case includes differently sized groups, an empty group, interleaved + # batches, and a nondefault stream. The reference runs on the CPU. + counts = [2, 0, 5, 1] + batch, m, k = 3, 4, 5 + torch.manual_seed(123) + weights_cpu = torch.randn(len(counts), batch, m, k) + input_cpu = torch.randn(sum(counts), batch, k) + expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) + offset = 0 + for i, n in enumerate(counts): + expected[offset:offset + n] = torch.einsum( + "bmk,nbk->nbm", weights_cpu[i].double(), + input_cpu[offset:offset + n].double(), + ) + offset += n + + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + weights = weights_cpu.cuda() + inputs = input_cpu.cuda() + output = torch.full(expected.shape, float("nan"), device="cuda") + status = call( + 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), + (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, + ) + assert status == 0, f"C shim prototype returned {status}" + stream.synchronize() + actual = output.cpu().double() + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + print( + f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " + f"max absolute error={(actual - expected).abs().max().item():.3g}; " + f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" + ) + + +if __name__ == "__main__": + main() From 19635e83e019afd0c46f13d1ca7b5b32620cc7f8 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 00:02:12 -0700 Subject: [PATCH 02/18] Route Torch grouped GEMM through the stable BMM C shim --- CHANGELOG.md | 4 + openequivariance/CMakeLists.txt | 15 +- .../_torch/extlib/__init__.py | 8 +- .../openequivariance/extension/group_mm.hpp | 178 ++++------ .../extension/libtorch_tp_jit.cpp | 14 +- .../extension/libtorch_tp_jit_stable.cpp | 10 + .../extension/stubs/stream.cpp | 13 +- .../openequivariance/extension/torch_core.hpp | 42 ++- sandbox/cshim_group_gemm/DESIGN.md | 311 ++++-------------- tests/group_gemm_test.py | 173 ++++++++++ tests/import_test.py | 28 ++ 11 files changed, 391 insertions(+), 405 deletions(-) create mode 100644 tests/group_gemm_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c9f9847..22257be3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## Latest Changes +- Route Torch grouped GEMM through the stable BMM C shim on CUDA and ROCm, + using Torch's current device stream and precision settings. Remove OEQ's + direct cuBLAS and rocBLAS dependencies. + ### v0.7.0 (2026-09-10) **Added**: - Public XLA FFI registration provider diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7bfbd499..7bb26d4e 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -67,6 +67,7 @@ function(add_stable_extension target_name backend_define link_libraries) # Enforce CXX11 ABI to match LibTorch target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 + TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION ) @@ -97,6 +98,7 @@ function(add_stable_extension target_name backend_define link_libraries) target_compile_definitions(${aoti_target_name} PRIVATE ${backend_define}=1 + TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 ) @@ -117,7 +119,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,7 +139,6 @@ 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}") @@ -146,6 +146,7 @@ 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,18 +160,12 @@ 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}") + add_stable_extension(oeq_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() if(NOT CUDAToolkit_FOUND AND NOT hip_FOUND) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 2a114b3e..ed41bfa6 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/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") @@ -125,8 +125,10 @@ def load_jit_extension(): extra_cflags.append("-DCUDA_BACKEND") elif torch.version.hip: - torch_libs = library_paths("cuda")[0] - extra_link_args.append("-Wl,-rpath," + torch_libs) + hip_lib_dirs = library_paths("cuda") + extra_link_args.append("-Wl,-rpath," + hip_lib_dirs[0]) + extra_link_args.extend("-L" + directory for directory in hip_lib_dirs) + extra_link_args.extend(["-ltorch_hip", "-lhiprtc"]) extra_cflags.append("-DHIP_BACKEND") torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] diff --git a/openequivariance/openequivariance/extension/group_mm.hpp b/openequivariance/openequivariance/extension/group_mm.hpp index 19249c39..247b6b3d 100644 --- a/openequivariance/openequivariance/extension/group_mm.hpp +++ b/openequivariance/openequivariance/extension/group_mm.hpp @@ -1,141 +1,73 @@ #pragma once +#include #include #include #include +#include -#ifdef CUDA_BACKEND - #include "cublas_v2.h" - #include +#include - 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 - - 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 -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, + decltype(&aoti_torch_delete_tensor_object)>; + +inline GroupMMTensor group_mm_view( + AtenTensorHandle tensor, std::array sizes, + std::array 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(A_raw); - T* B_base = reinterpret_cast(B_raw); - T* C_base = reinterpret_cast(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(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(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::value) { - stat = cublasSgemmStridedBatched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else if (std::is_same::value) { - stat = cublasDgemmStridedBatched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(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::value) { - stat = rocblas_sgemm_strided_batched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else if (std::is_same::value) { - stat = rocblas_dgemm_strided_batched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(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; } } + +} diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index ddabd0bb..c2b71841 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -2,7 +2,7 @@ #include #ifdef CUDA_BACKEND - #include + #include #endif #ifdef HIP_BACKEND @@ -10,9 +10,11 @@ #endif #include +#include #include #include #include +#include #include using Tensor = torch::Tensor; @@ -29,6 +31,16 @@ constexpr Dtype kByte = torch::kByte; #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY TORCH_LIBRARY +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) { diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 6bf3d51f..1b15a97b 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -27,6 +27,16 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte; #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY +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) { diff --git a/openequivariance/openequivariance/extension/stubs/stream.cpp b/openequivariance/openequivariance/extension/stubs/stream.cpp index fd011c35..346493c4 100644 --- a/openequivariance/openequivariance/extension/stubs/stream.cpp +++ b/openequivariance/openequivariance/extension/stubs/stream.cpp @@ -1,8 +1,15 @@ +#define USE_CUDA + #include -#include +#include 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; } -} \ No newline at end of file +} diff --git a/openequivariance/openequivariance/extension/torch_core.hpp b/openequivariance/openequivariance/extension/torch_core.hpp index ab78d96a..52dd3331 100644 --- a/openequivariance/openequivariance/extension/torch_core.hpp +++ b/openequivariance/openequivariance/extension/torch_core.hpp @@ -621,13 +621,40 @@ inline tuple jit_conv_double_backward( inline Tensor group_gemm( Tensor A, Tensor B, Tensor ragged_counts, int64_t num_W, int64_t batch_size, int64_t m, int64_t k, int64_t ragged_inner) { + TCHECK(A.is_cuda() && B.is_cuda(), "group_gemm: A and B must be GPU tensors"); + TCHECK(A.get_device() == B.get_device(), "group_gemm: A and B must be on the same device"); TCHECK(A.scalar_type() == B.scalar_type(), "group_gemm: A and B must have the same dtype"); + TCHECK(A.scalar_type() == kFloat || A.scalar_type() == kDouble, + "group_gemm: unsupported dtype, expected float32 or float64"); + TCHECK(ragged_counts.is_cpu(), "group_gemm: ragged_counts must be on the CPU"); TCHECK(ragged_counts.scalar_type() == kLong, "group_gemm: ragged_counts must be int64"); + TCHECK(num_W >= 0 && batch_size >= 0 && m >= 0 && k >= 0, + "group_gemm: dimensions must be nonnegative"); + TCHECK(ragged_inner == 0 || ragged_inner == 1, "group_gemm: ragged_inner must be 0 or 1"); + TCHECK(ragged_counts.dim() == 1 && ragged_counts.size(0) == num_W, + "group_gemm: ragged_counts must contain num_W entries"); + TCHECK(B.dim() == 3, "group_gemm: B must have shape [rows, batch_size, k]"); + const int64_t rows = B.size(0); + check_tensor(B, {rows, batch_size, k}, A.scalar_type(), "group_gemm B"); + if (ragged_inner == 0) + check_tensor(A, {num_W, batch_size, m, k}, A.scalar_type(), "group_gemm A"); + else + check_tensor(A, {rows, batch_size, m}, A.scalar_type(), "group_gemm A"); + + Tensor rc_c = tensor_contiguous(ragged_counts); + const auto* rc_ptr = static_cast(data_ptr(rc_c)); + int64_t row_count = 0; + for (int64_t i = 0; i < num_W; ++i) { + TCHECK(rc_ptr[i] >= 0 && rc_ptr[i] <= rows - row_count, + "group_gemm: ragged_counts must be nonnegative and sum to the number of rows"); + row_count += rc_ptr[i]; + } + TCHECK(row_count == rows, + "group_gemm: ragged_counts must sum to the number of rows"); + TensorDeviceGuard device_guard(A); Tensor A_c = tensor_contiguous(A); Tensor B_c = tensor_contiguous(B); - Tensor rc_c = tensor_contiguous(ragged_counts); - int64_t* rc_ptr = reinterpret_cast(data_ptr(rc_c)); Tensor C; if (ragged_inner == 0) { @@ -637,15 +664,8 @@ inline Tensor group_gemm( C = tensor_zeros_like(A, make_sizes({num_W, batch_size, m, k})); } - if (A.scalar_type() == kFloat) { - group_gemm_blas(data_ptr(A_c), data_ptr(B_c), data_ptr(C), rc_ptr, - (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner); - } else if (A.scalar_type() == kDouble) { - group_gemm_blas(data_ptr(A_c), data_ptr(B_c), data_ptr(C), rc_ptr, - (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner); - } else { - throw std::logic_error("group_gemm: unsupported dtype, expected float32 or float64"); - } + oeq::group_gemm_torch(tensor_handle(A_c), tensor_handle(B_c), tensor_handle(C), + rc_ptr, num_W, batch_size, m, k, ragged_inner); return C; } diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index 3b919e07..9626260e 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -1,292 +1,95 @@ # Move Torch cuBLAS calls to the stable C shim -Replace OEQ's direct CUDA and ROCm BLAS calls for `libtorch_tp_jit::group_gemm` with PyTorch's `aoti_torch_cuda_bmm_out` C API. Keep the existing grouped operation, tensor layouts, output allocation, and custom autograd rules. PyTorch will own the underlying BLAS handles, streams, library calls, and backend selection. +The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). +GPU execution of this implementation is pending. GPU access is paused at the user's request; the earlier prototype results below are not results for this implementation. -## Problem and scope +## Implementation -OEQ's grouped GEMM is a host loop over nonempty ragged groups, issuing one strided batched GEMM per group. It does not require a vendor's heterogeneous grouped-GEMM API. The only direct BLAS computation calls found are float32/float64 strided batched GEMMs in `extension/group_mm.hpp`, called from the Torch backend in `extension/torch_core.hpp`. +The implementation lives in [group_mm.hpp](../../openequivariance/openequivariance/extension/group_mm.hpp), with input validation and output allocation in [torch_core.hpp](../../openequivariance/openequivariance/extension/torch_core.hpp). The algorithm is one BMM-out invocation per nonempty ragged group, using the same interleaved layout as the previous strided batched GEMMs. It does not need a vendor's heterogeneous grouped-GEMM API. -Using PyTorch's stable ABI to obtain a cuBLAS handle does not make that handle interchangeable with a separately loaded cuBLAS implementation. Earlier H100 experiments passed with matching handle/function owners, while some foreign-owner combinations returned incorrect streams, failed, or crashed. Successful version queries did not establish handle compatibility. Moving the GEMM itself behind PyTorch's C ABI removes this borrowed-handle boundary. +Inputs are made contiguous by the registered operator, then passed as borrowed `AtenTensorHandle` values to the shared helper. The stable extension gets these handles from `torch::stable::Tensor::get()`. The source/JIT extension uses Torch's `tensor_pointer_to_tensor_handle` utility to borrow a handle to its local `at::Tensor`; that bridge is compiled against the installed Torch, as the existing JIT extension already is. ATen C++ objects do not cross the stable extension's ABI boundary. -The current checkout at `dbe854415da7771eba33195534c171adbca5677b` creates its own static BLAS handle; the inspected PR snapshot at `b551d7469db6d7ac688859cec46a3bb2b1c71a5b` borrows Torch's handle. The replacement removes either form of OEQ-side handle management. The final patch must be based on the branch being merged. +The helper creates temporary views with `aoti_torch__reinterpret_tensor`. These views retain the original storage and add their element offsets to the original tensor's storage offset. Their handles are released through RAII after each call, including when an error is returned. Original input and output handles remain owned by the caller. This replaces the prototype's `from_blob` construction and avoids reconstructing storage or querying pointer devices for every group. [Torch reinterpretation API](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L435). -The stable wheels currently target PyTorch 2.10. Use official stable C declarations, with no dependency on ATen's C++ ABI in that build. The C entry point was also exercised successfully on the available PyTorch 2.7 installation; that observation does not establish support for every older Torch release. [PyTorch stable ABI documentation](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). +The operation checks GPU placement, matching devices and float32/float64 dtypes, tensor shapes, nonnegative dimensions, and mode 0 or 1. Ragged counts must be a CPU int64 vector with one entry per group. Counts may be noncontiguous; the operator makes a contiguous CPU copy before reading them. Counts must be nonnegative and sum to the input row count. Incremental bounds checking avoids overflowing the count sum. These checks establish the bounds needed by the storage-view helper. -## Public interface +Output allocation uses Torch and starts at zero. Empty groups are skipped; in mode 1 this leaves the corresponding weight-gradient block zero. Zero batch, output, or contraction dimensions return the zero-initialized output without invoking BMM. Dimensions and offsets remain int64 throughout the grouped-GEMM implementation. -Include `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`, which declares: +## Layout mapping -```cpp -AOTITorchError aoti_torch_cuda_bmm_out( - AtenTensorHandle out, - AtenTensorHandle self, - AtenTensorHandle mat2); -``` +Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Offsets are relative to the logical beginning of each original tensor; strides and offsets are in elements. -Use this generated API rather than the deprecated `aoti_torch_bmm_out` spelling. The API takes PyTorch tensor handles, not cuBLAS/hipBLAS handles. OEQ does not call a BLAS version getter or select a vendor library at runtime. [PyTorch 2.10 declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74). +| Mode | Tensor | Element offset | View shape | View strides | +| --- | --- | --- | --- | --- | +| 0 | B / input | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| 0 | A / weights | `i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | +| 0 | C / output | `o*b*m` | `[b,n,m]` | `[m,b*m,1]` | +| 1 | A / left | `o*b*m` | `[b,m,n]` | `[m,1,b*m]` | +| 1 | B / right | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| 1 | C / output | `i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | -Add a Torch-specific helper, such as `group_mm_torch.hpp`, taking the existing input/output pointers, dimensions, CPU ragged counts, dtype, and explicit device index. Both Torch extension variants should use it where their supported Torch versions expose the required C symbols. Keep this helper outside the framework-independent CUDA/HIP kernel backend. +The final operation for each group is `aoti_torch_cuda_bmm_out(output, left, right)`. The API overwrites the output with alpha 1 and beta 0. Torch may choose its own BLAS kernel or backend. The layout code requests views, but makes no unconditional promise about copies inside Torch's BMM implementation. [Torch BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). -Create temporary, non-owning tensor views with `aoti_torch_create_tensor_from_blob_v2`. Specify sizes and strides in elements, storage offset zero, and the actual input device. Release each temporary tensor handle with `aoti_torch_delete_tensor_object` through RAII, including error paths. Releasing metadata must not release the caller's input or output storage. [PyTorch's blob implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L525). +## Device, stream, precision, and autograd -## Layout mapping +The registered operator guards A's device before making GPU contiguous copies or allocating the output. The stable build uses `torch::stable::accelerator::DeviceGuard`; the JIT build uses `c10::DeviceGuard`. Each restores the caller's current device on return. -Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Pointer offsets and strides below are measured in elements. +BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. -| Mode | View | Base pointer | Shape | Strides | -| --- | --- | --- | --- | --- | -| `ragged_inner == 0` | Left / input | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| `ragged_inner == 0` | Right / weights | `A + i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | -| `ragged_inner == 0` | Output | `C + o*b*m` | `[b,n,m]` | `[m,b*m,1]` | -| `ragged_inner == 1` | Left | `A + o*b*m` | `[b,m,n]` | `[m,1,b*m]` | -| `ragged_inner == 1` | Right | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| `ragged_inner == 1` | Output | `C + i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | +Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. -Call BMM-out once for each nonempty group. Preserve the existing zero-initialized output and skip empty groups, including the all-empty case. This matters for the weight-gradient output blocks in mode 1. Use 64-bit dimensions and offsets throughout instead of the current narrowing casts to `int`. +## ABI and build paths -These views preserve OEQ's interleaved batch layout without explicitly transposing or copying the underlying buffers. PyTorch can handle BLAS-compatible strides directly, but some layouts can trigger its internal copy path; avoiding all copies is not an unconditional API guarantee. [BMM layout handling](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). +The official declaration is in `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`. The backend-specific spelling is used instead of the deprecated `aoti_torch_bmm_out`. [Declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74), [stable ABI policy](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). -## Device, stream, and autograd behavior +Both `aoti_torch_cuda_bmm_out` and `aoti_torch__reinterpret_tensor` are present in the inspected PyTorch 2.4 headers, matching OEQ's documented source/JIT baseline. The stable build explicitly targets the 2.10 ABI and continues to use the pinned LibTorch headers/libraries. This uses Torch's stated ABI guarantees; it does not try to select cuBLAS versions or infer handle ownership from version queries. -Guard the inputs' actual GPU device before making contiguous copies, allocating the output, constructing views, or invoking BMM. Restore the caller's device on return. The 2.10 stable implementation can use `aoti_torch_create_device_guard` / `aoti_torch_delete_device_guard`, whose implementation selects Torch's active accelerator; this avoids requiring a CUDA-specific guard implementation in the shared helper. The existing standalone prototype uses the older CUDA-named guard because it was tested on Torch 2.7. [Generic guard implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L1468). +The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` library supplying the BMM shim. CUDA also retains its driver/runtime/NVRTC dependencies; ROCm links hipRTC. All direct cuBLAS/rocBLAS calls, handle management, and corresponding OEQ link dependencies have been removed. -Use Torch's current stream for that device. Do not set a BLAS stream, create a stream, or synchronize inside the operator. Original tensors must remain valid through launch; asynchronous storage lifetime follows the usual PyTorch current-stream contract. No global or cached raw pointers or tensor views are needed. +The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. -Validate the raw-pointer helper's preconditions at the tensor boundary: same GPU device and dtype, supported float32/float64 types, contiguous CPU int64 counts, valid counts length and nonnegative values, compatible input shapes, and a valid mode. Check that counts describe the available rows before constructing views. Keep the existing custom operator schema, fake implementation, and backward formulas. BMM-out is an internal implementation detail; the custom autograd registration remains responsible for derivatives. +The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. -PyTorch's BMM implementation uses alpha 1 and beta 0, matching the existing GEMMs. It may choose a different BLAS backend or kernel and apply Torch's precision/determinism settings, so require numerical agreement rather than bitwise equivalence. [BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L680). +## ROCm evidence -## Why this should work on ROCm +PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -Source inspection of PyTorch **v2.10.0** supports using the same generated function on both GPU platforms: +This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution remains pending. -1. PyTorch's build defines `GENERATED_CXX_TORCH_CUDA` to contain `c_shim_cuda.cpp`, then explicitly adds that source to `torch_hip` under `USE_ROCM`. The generated shim is therefore part of the ROCm library too. [Generated source definition](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L331), [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941). -2. The shim generator names the API using the `CUDA` dispatch key and generates a call into the corresponding ATen backend. The exported name remains `aoti_torch_cuda_bmm_out`; do not invent an `aoti_torch_hip_bmm_out` symbol. [Shim generator](https://github.com/pytorch/pytorch/blob/v2.10.0/torchgen/gen_aoti_c_shim.py#L492). -3. BMM reaches Torch's GEMM/batched GEMM implementation. Its HIP conversion maps `cublasSgemmStridedBatched` and `cublasDgemmStridedBatched` to their `hipblas` equivalents. Torch also contains ROCm-specific backend selection, including a double-precision fallback when hipBLASLt cannot handle the operation. [HIP mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [float/double backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). +## JAX scope -Torch's ROCm tensors use the `cuda` device interface. The helper should use Torch's CUDA device-type value for both these Torch builds, rather than confusing that convention with a distinct external HIP tensor type. [ROCm semantics](https://docs.pytorch.org/docs/main/notes/hip.html). +JAX does not call this grouped-GEMM helper and has no BLAS dependency in its extension target. The replacement is entirely within the Torch frontend. The common CUDA/HIP kernel compilation backend gains no Torch dependency. -This is **source-level evidence**, not an AMD hardware test or a binary export audit of a ROCm wheel. Verify the exports in the supported wheel and run one AMD smoke case before marking ROCm complete. +JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). -## Build and packaging changes +## Validation and next GPU run -- Remove direct cuBLAS/rocBLAS includes, calls, and handle management from the Torch grouped-GEMM path. On the PR branch, also remove the now-unused borrowed-handle adapters. -- Remove `CUDA::cublas`, the grouped-GEMM-only `find_package(rocblas)` / `${HIP_BLAS_LIB}`, and the JIT loader's explicit `-lcublas` link flag after confirming there are no remaining consumers. Torch still brings its own BLAS dependencies. -- Resolve the generated BMM symbol through the installed Torch GPU library: `libtorch_cuda` for CUDA and `libtorch_hip` for ROCm. Common tensor/guard C APIs are supplied by Torch's common library. -- OEQ's wheel build currently downloads CPU LibTorch headers/libraries and creates GPU link stubs from `extension/stubs/stream.cpp`. Extend those build-only stubs with the exact generated BMM declaration/definition, sharing the header so signature drift causes a build error. If retaining the prototype's GPU-specific guards, add their GPU exports too; the proposed 2.10 generic guard avoids that need. Never package or execute the fake implementations. Check that installed extensions resolve against the real Torch libraries. -- Apply changes to both the Python extension and the AOTI shared library targets, and to source/JIT builds. The existing HIP CMake target is named `torch_stable_hip` while import/install expectations use `oeq_stable_hip`; resolve that naming mismatch as part of making the HIP artifact load correctly. -- Keep the existing CUDA/HIP runtime and runtime-compiler dependencies needed by OEQ-generated kernels. Removing the BLAS dependency does not make the entire extension independent of the GPU platform. +Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 30 GPU integration cases collect successfully; none has run on a GPU yet. -The stable-wheel baseline remains 2.10. Confirm the supported source/JIT baseline before sharing every C helper with that build: the prototype's helpers were exercised on 2.7, whereas the proposed generic guard needs its own baseline check. If older JIT support must be preserved, a JIT-only adapter using that installed Torch's ordinary BMM-out/device guard is acceptable; it still eliminates direct vendor BLAS calls and does not enter the stable wheel. +[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty groups/dimensions, backward gradients, the current stream, device guarding, and invalid counts. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. -## JAX scope +[The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. -Neither the inspected PR snapshot nor the current JAX extension calls `group_gemm_blas`; its CMake target links the runtime, driver, and NVRTC, without BLAS. This replacement therefore remains entirely within Torch. - -JAX's public FFI has buffers and a GPU stream, but no equivalent BMM shim or public BLAS-handle getter. If a future JAX operation needs this GEMM, express it in the compiled JAX graph where shapes permit, or provide a separate FFI implementation whose BLAS handles and calls come from the same library. [JAX GPU FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu), [public XLA FFI API](https://github.com/openxla/xla/blob/main/xla/ffi/api/c_api.h#L757). - -## Existing prototype and validation - -A standalone C++ prototype already implements both layout branches using only Torch C APIs. Its single float32 forward smoke case passed on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against an independent CPU float64 reference. - -Earlier ctypes calls to the same C shim passed 40 cases spanning both dtypes and modes. Four selected profiles showed one GEMM kernel with no observed copy or GPU allocation inside the BMM call. These measurements do not cover the CPU cost of temporary tensor metadata or establish end-to-end performance. - -The integrated extension, its autograd registration, the final wheel link setup, and ROCm execution remain to be validated. No additional GPU experiments were run for this design. - -## Acceptance criteria - -- Both grouped-GEMM modes use Torch BMM, with no direct vendor BLAS symbols or handles in OEQ's Torch extension. -- Existing layouts, empty groups, float32/float64 behavior, operator schema, fake implementation, and gradients are preserved. -- Device guarding and a nondefault current stream work; the caller's current device is restored. -- Stable CUDA and HIP artifacts resolve the official C symbols in the real installed Torch GPU library; link stubs do not ship. Audit direct dynamic dependencies to confirm BLAS is now Torch's responsibility. -- Run focused correctness checks for both modes/dtypes and one relevant autograd case during integration, plus one ROCm smoke case when AMD hardware is available. Keep this bounded; no exhaustive vendor-version matrix is needed to establish the prototype. -- Check representative overhead before claiming performance parity. Graph capture and broader multi-device coverage should be checked where required by the supported operator contract. - -## Prototype code - -The following is the existing standalone CUDA prototype, not the final production registration or generic-device-guard adaptation. The caller supplies validated buffers and zero-initializes outputs where empty groups need to remain zero. - -```cpp -// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. -#include - -#include -#include -#include -#include -#include -#include - -namespace { - -void check(AOTITorchError status) { - if (status != AOTI_TORCH_SUCCESS) - throw std::runtime_error("PyTorch C shim failed"); -} - -using Tensor = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_tensor_object)>; -using DeviceGuard = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_cuda_guard)>; - -// Only tensor metadata is created; the caller retains ownership of the buffer. -Tensor view(void* data, std::array sizes, - std::array strides, int32_t dtype, int32_t device) { - AtenTensorHandle tensor = nullptr; - check(aoti_torch_create_tensor_from_blob_v2( - data, 3, sizes.data(), strides.data(), 0, dtype, - aoti_torch_device_type_cuda(), device, &tensor, - aoti_torch_layout_strided(), nullptr, 0)); - return Tensor(tensor, aoti_torch_delete_tensor_object); -} - -} // namespace - -template -void group_gemm_cshim( - T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, - int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { - static_assert(std::is_same_v || std::is_same_v); - const int32_t dtype = std::is_same_v - ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); - - CUDAGuardHandle raw_guard = nullptr; - check(aoti_torch_create_cuda_guard(device, &raw_guard)); - DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); - - int64_t offset = 0; - for (int i = 0; i < num_groups; ++i) { - const int64_t n = ragged_counts[i]; - if (n == 0) continue; // Preserve the original empty-group behavior. - - if (ragged_inner == 0) { - // [batch, n, k] @ [batch, k, m] -> [batch, n, m] - auto input = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto weight = view(A + i * batch * m * k, - {batch, k, m}, {m * k, 1, k}, dtype, device); - auto output = view(C + offset * batch * m, - {batch, n, m}, {m, batch * m, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); - } else { - // [batch, m, n] @ [batch, n, k] -> [batch, m, k] - auto left = view(A + offset * batch * m, - {batch, m, n}, {m, 1, batch * m}, dtype, device); - auto right = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto output = view(C + i * batch * m * k, - {batch, m, k}, {m * k, k, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); - } - offset += n; - } -} - -// Small ctypes entry point for the smoke test, not production registration code. -extern "C" int oeq_group_gemm_cshim( - int dtype, void* A, void* B, void* C, const int64_t* counts, - int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { - try { - if (dtype == 0) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else if (dtype == 1) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else { - throw std::runtime_error("Expected float32 or float64"); - } - return 0; - } catch (const std::exception& error) { - std::fprintf(stderr, "%s\n", error.what()); - return 1; - } -} +After GPU access resumes, start with one case of the production operator: + +```sh +pytest -q 'tests/group_gemm_test.py::test_group_gemm_matches_reference[contiguous-0-dtype0]' ``` -
-Reproduce the single CUDA smoke case - -Save the C++ above as `group_gemm.cpp` and this script beside it as `smoke_test.py`, then run `python3 smoke_test.py` with a CUDA-enabled Torch installation and a C++ compiler. The script builds only the standalone prototype and executes one configuration. - -```python -"""Compile the C++ prototype and run one float32 CUDA correctness case.""" -import ctypes -from pathlib import Path -import subprocess - -import torch - - -def main(): - root = Path(__file__).resolve().parent - torch_root = Path(torch.__file__).resolve().parent - library = root / "group_gemm.so" - subprocess.run( - [ - "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", - f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), - f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", - "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", - "-o", str(library), - ], - check=True, - ) - lib = ctypes.CDLL(str(library)) - call = lib.oeq_group_gemm_cshim - i64 = ctypes.c_int64 - call.argtypes = [ - ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, - ctypes.c_int, ctypes.c_int32, - ] - call.restype = ctypes.c_int - - # One case includes differently sized groups, an empty group, interleaved - # batches, and a nondefault stream. The reference runs on the CPU. - counts = [2, 0, 5, 1] - batch, m, k = 3, 4, 5 - torch.manual_seed(123) - weights_cpu = torch.randn(len(counts), batch, m, k) - input_cpu = torch.randn(sum(counts), batch, k) - expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) - offset = 0 - for i, n in enumerate(counts): - expected[offset:offset + n] = torch.einsum( - "bmk,nbk->nbm", weights_cpu[i].double(), - input_cpu[offset:offset + n].double(), - ) - offset += n - - torch.cuda.set_device(0) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - weights = weights_cpu.cuda() - inputs = input_cpu.cuda() - output = torch.full(expected.shape, float("nan"), device="cuda") - status = call( - 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), - (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, - ) - assert status == 0, f"C shim prototype returned {status}" - stream.synchronize() - actual = output.cpu().double() - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) - print( - f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " - f"max absolute error={(actual - expected).abs().max().item():.3g}; " - f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" - ) - - -if __name__ == "__main__": - main() +Then run the focused integration suite against the stable build and the JIT build in separate processes: + +```sh +pytest -q tests/import_test.py tests/group_gemm_test.py +OEQ_JIT_EXTENSION=1 pytest -q tests/import_test.py tests/group_gemm_test.py ``` -
+Existing symmetric-contraction integration tests exercise the surrounding model and its higher-order derivatives when the optional MACE dependency is available. No end-to-end performance claim is made before measurement. + +## Earlier experiments + +The original [standalone prototype](group_gemm.cpp) and its [smoke script](smoke_test.py) remain as historical experiment artifacts. They use blob views rather than the production helper's views of existing storage. + +The prototype passed one float32 forward case on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against a CPU float64 reference. Earlier ctypes calls to the BMM C shim passed 40 small cases across both dtypes/modes. Four selected BMM-only profiles showed no copy or GPU allocation inside the call; they did not measure view creation overhead. + +The earlier direct-cuBLAS experiments explain the migration: matching handle/function owners worked, whereas some separately loaded foreign-owner combinations failed or crashed despite successful version queries. They do not establish interchangeability of private BLAS handles. Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py new file mode 100644 index 00000000..6e983446 --- /dev/null +++ b/tests/group_gemm_test.py @@ -0,0 +1,173 @@ +import importlib + +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="A CUDA or ROCm GPU is required" +) + + +@pytest.fixture(scope="module") +def group_gemm(): + importlib.import_module("openequivariance._torch.symmetric_contraction") + return torch.ops.libtorch_tp_jit.group_gemm + + +def reference(A, B, counts, inner): + pieces = [] + offset = 0 + for i, n in enumerate(counts): + if inner == 0: + pieces.append(torch.einsum("bmk,nbk->nbm", A[i], B[offset : offset + n])) + else: + pieces.append( + torch.einsum( + "nbm,nbk->bmk", A[offset : offset + n], B[offset : offset + n] + ) + ) + offset += n + if inner == 0: + return torch.cat(pieces, dim=0) + return torch.stack(pieces, dim=0) + + +def make_input(shape, dtype, layout="contiguous", device="cuda"): + values = torch.arange(1, 1 + torch.Size(shape).numel(), dtype=dtype, device="cpu") + values = (values.remainder(17) - 8).reshape(shape) / 8 + if layout == "offset": + storage = torch.empty(values.numel() + 7, dtype=dtype, device=device) + result = storage[3 : 3 + values.numel()].view(shape) + result.copy_(values) + return result + if layout == "noncontiguous": + storage = torch.empty((*shape, 2), dtype=dtype, device=device) + result = storage[..., 0] + result.copy_(values) + return result + return values.to(device) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize("layout", ["contiguous", "offset", "noncontiguous"]) +def test_group_gemm_matches_reference(group_gemm, dtype, inner, layout): + counts = [2, 0, 3] + batch, m, k = 2, 3, 4 + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, dtype, layout) + B = make_input((sum(counts), batch, k), dtype, layout) + expected = reference(A.cpu().double(), B.cpu().double(), counts, inner) + counts_tensor = torch.tensor([n for n in counts for _ in range(2)], device="cpu")[ + ::2 + ] + + actual = group_gemm(A, B, counts_tensor, len(counts), batch, m, k, inner) + + torch.testing.assert_close(actual.cpu().double(), expected, rtol=1e-5, atol=1e-5) + assert actual.dtype == dtype + assert actual.device == A.device + + +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize( + "counts,batch,m,k", + [ + ([0, 0], 2, 3, 4), + ([], 2, 3, 4), + ([2, 0, 3], 0, 3, 4), + ([2, 0, 3], 2, 0, 4), + ([2, 0, 3], 2, 3, 0), + ], +) +def test_group_gemm_empty_dimensions(group_gemm, inner, counts, batch, m, k): + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = torch.empty(A_shape, device="cuda", dtype=torch.float64) + B = torch.empty((sum(counts), batch, k), device="cuda", dtype=torch.float64) + actual = group_gemm( + A, + B, + torch.tensor(counts, dtype=torch.int64, device="cpu"), + len(counts), + batch, + m, + k, + inner, + ) + expected_shape = ( + (sum(counts), batch, m) if inner == 0 else (len(counts), batch, m, k) + ) + assert actual.shape == expected_shape + torch.testing.assert_close(actual, torch.zeros_like(actual)) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_backward(group_gemm, inner): + counts = [2, 0, 3] + batch, m, k = 2, 3, 4 + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, torch.float64, "noncontiguous").requires_grad_() + B = make_input((sum(counts), batch, k), torch.float64, "offset").requires_grad_() + A_ref = A.detach().cpu().requires_grad_() + B_ref = B.detach().cpu().requires_grad_() + expected = reference(A_ref, B_ref, counts, inner) + actual = group_gemm( + A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner + ) + grad = make_input(actual.shape, torch.float64, "noncontiguous") + + actual_grads = torch.autograd.grad(actual, (A, B), grad) + expected_grads = torch.autograd.grad(expected, (A_ref, B_ref), grad.cpu()) + + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close( + actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 + ) + + +def test_group_gemm_current_stream(group_gemm): + counts = torch.tensor([2, 0, 3], device="cpu") + A = torch.zeros((3, 2, 3, 4), device="cuda") + B = torch.zeros((5, 2, 4), device="cuda") + group_gemm(A, B, counts, 3, 2, 3, 4, 0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + if not torch.version.hip: + torch.cuda._sleep(1_000_000) + A.fill_(2) + B.fill_(3) + actual = group_gemm(A, B, counts, 3, 2, 3, 4, 0).clone() + stream.synchronize() + torch.testing.assert_close(actual, torch.full_like(actual, 24)) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs are required") +def test_group_gemm_device_guard(group_gemm): + with torch.cuda.device(0): + A = torch.full((1, 2, 3, 4), 2.0, device="cuda:1") + B = torch.full((2, 2, 4), 3.0, device="cuda:1") + actual = group_gemm(A, B, torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0) + assert torch.cuda.current_device() == 0 + assert actual.device == torch.device("cuda:1") + torch.testing.assert_close(actual, torch.full_like(actual, 24)) + with pytest.raises(RuntimeError, match="same device"): + group_gemm( + A, B.to("cuda:0"), torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0 + ) + + +@pytest.mark.parametrize("counts", [[-1, 0, 6], [2, 0, 2], [2, 0, 4]]) +def test_group_gemm_rejects_invalid_counts(group_gemm, counts): + A = torch.empty((3, 2, 3, 4), device="cuda") + B = torch.empty((5, 2, 4), device="cuda") + with pytest.raises(RuntimeError, match="ragged_counts"): + group_gemm(A, B, torch.tensor(counts, device="cpu"), 3, 2, 3, 4, 0) + + +def test_group_gemm_requires_cpu_counts(group_gemm): + A = torch.empty((3, 2, 3, 4), device="cuda") + B = torch.empty((5, 2, 4), device="cuda") + with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): + group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) diff --git a/tests/import_test.py b/tests/import_test.py index bf26af31..aa53630e 100644 --- a/tests/import_test.py +++ b/tests/import_test.py @@ -1,4 +1,9 @@ from importlib.metadata import version +import re +import shutil +import subprocess + +import pytest def test_import(): @@ -14,3 +19,26 @@ def test_extension_built(): assert BUILT_EXTENSION_ERROR is None assert BUILT_EXTENSION + + +def test_extension_has_no_direct_blas_dependency(): + if not shutil.which("readelf"): + pytest.skip("readelf is required to inspect ELF dependencies") + import openequivariance + from openequivariance._torch.extlib import extension_module + + paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} + for path in paths: + dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) + dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] + assert not any( + re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies + ) + symbols = subprocess.check_output( + ["readelf", "--dyn-syms", "--wide", path], text=True + ) + undefined = [line for line in symbols.splitlines() if " UND " in line] + assert not any( + re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) + for line in undefined + ) From bab7e4dba67525a53287e8991820390aaa43661e Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:09:32 -0700 Subject: [PATCH 03/18] Expand grouped GEMM integration coverage --- tests/group_gemm_test.py | 135 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 6e983446..6bf70330 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -1,4 +1,6 @@ import importlib +import subprocess +import sys import pytest import torch @@ -171,3 +173,136 @@ def test_group_gemm_requires_cpu_counts(group_gemm): B = torch.empty((5, 2, 4), device="cuda") with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) + + +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize( + "counts,batch,m,k", + [ + ([1, 0, 5], 1, 1, 7), + ([1, 0, 5], 3, 7, 1), + ([1, 0, 5], 3, 1, 1), + ([1, 7, 0, 13], 4, 17, 29), + ], +) +def test_group_gemm_varied_shapes(group_gemm, inner, counts, batch, m, k): + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, torch.float64, "offset") + B = make_input((sum(counts), batch, k), torch.float64, "offset") + actual = group_gemm( + A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner + ) + expected = reference(A.cpu(), B.cpu(), counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_double_backward(group_gemm, inner): + counts = torch.tensor([1, 0, 2], device="cpu") + A_shape = (3, 2, 2, 3) if inner == 0 else (3, 2, 2) + A = make_input(A_shape, torch.float64).requires_grad_() + B = make_input((3, 2, 3), torch.float64).requires_grad_() + + def operation(A, B): + return group_gemm(A, B, counts, 3, 2, 2, 3, inner) + + assert torch.autograd.gradgradcheck(operation, (A, B), fast_mode=True) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_graph_replay(group_gemm, inner): + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64) + B = make_input((5, 2, 4), torch.float64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) + + for scale in (2, 3): + A.mul_(scale) + B.add_(0.25) + graph.replay() + expected = reference(A.cpu(), B.cpu(), counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_compile(group_gemm, inner): + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64).requires_grad_() + B = make_input((5, 2, 4), torch.float64).requires_grad_() + + def operation(A, B, counts): + return group_gemm(A, B, counts, 3, 2, 3, 4, inner) + + compiled = torch.compile(operation, fullgraph=True) + actual = compiled(A, B, counts_tensor) + A_ref = A.detach().cpu().requires_grad_() + B_ref = B.detach().cpu().requires_grad_() + expected = reference(A_ref, B_ref, counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + actual_grads = torch.autograd.grad(actual.sum(), (A, B)) + expected_grads = torch.autograd.grad(expected.sum(), (A_ref, B_ref)) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close( + actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_aoti(group_gemm, inner, tmp_path): + import openequivariance + + class Model(torch.nn.Module): + def forward(self, A, B, counts): + return group_gemm(A, B, counts, 3, 2, 3, 4, inner) + + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64) + B = make_input((5, 2, 4), torch.float64) + exported = torch.export.export(Model(), (A, B, counts_tensor), strict=False) + package_path = torch._inductor.aoti_compile_and_package( + exported, package_path=str(tmp_path / "group_gemm.pt2") + ) + inputs_path = tmp_path / "inputs.pt" + torch.save( + (A.cpu(), B.cpu(), counts_tensor, reference(A.cpu(), B.cpu(), counts, inner)), + inputs_path, + ) + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +import torch + +torch.ops.load_library(sys.argv[1]) +model = torch._inductor.aoti_load_package(sys.argv[2]) +A, B, counts, expected = torch.load(sys.argv[3], weights_only=True) +actual = model(A.cuda(), B.cuda(), counts) +torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) +""", + openequivariance.torch_ext_so_path(), + package_path, + str(inputs_path), + ], + capture_output=True, + text=True, + timeout=180, + ) + assert result.returncode == 0, result.stdout + result.stderr From 387ecf74ec27a84f0ac7c5ee63ceea308c3b34a0 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:19:50 -0700 Subject: [PATCH 04/18] Package stable libraries in wheels and resolve editable library paths --- .github/workflows/verify_extension_build.yml | 3 ++- openequivariance/CMakeLists.txt | 2 +- .../openequivariance/_torch/extlib/__init__.py | 15 ++++++++++++--- sandbox/cshim_group_gemm/DESIGN.md | 6 ++++-- 4 files changed, 19 insertions(+), 7 deletions(-) diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index 6c903205..69217eaf 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -29,10 +29,11 @@ jobs: sudo apt-get update sudo apt install nvidia-cuda-toolkit pip install -r .github/workflows/requirements_cuda_ci.txt - pip install -e "./openequivariance[jax]" + pip install "./openequivariance[jax]" - name: Test CUDA extension build via import run: | + python -c "import openequivariance as oeq; assert oeq.USE_PRECOMPILED_EXTENSION" pytest tests/import_test.py export OEQ_JIT_EXTENSION=1 diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7bb26d4e..0aef37e0 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -52,7 +52,7 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) # Create nanobind extension diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index ed41bfa6..6785c230 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,6 +5,7 @@ import warnings import sysconfig import contextlib +import importlib.util from pathlib import Path from packaging.version import Version @@ -180,6 +181,16 @@ def load_precompiled_extension(): ) +def _has_precompiled_extension(): + backend = "hip" if IS_HIP else "cuda" + spec = importlib.util.find_spec(f"{__name__}.oeq_stable_{backend}") + return ( + spec is not None + and spec.origin is not None + and Path(spec.origin).with_name(f"liboeq_stable_{backend}_aoti.so").is_file() + ) + + USE_PRECOMPILED_EXTENSION = True WARNING_MESSAGE = "" @@ -195,9 +206,7 @@ def load_precompiled_extension(): WARNING_MESSAGE += "HIP does not support precompiled extension yet.\n" USE_PRECOMPILED_EXTENSION = False -if not os.path.exists( - os.path.join(os.path.dirname(__file__), "liboeq_stable_cuda_aoti.so") -): +if not _has_precompiled_extension(): WARNING_MESSAGE += "Precompiled extension shared object not found.\n" USE_PRECOMPILED_EXTENSION = False diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index 9626260e..c5c5ea2d 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -49,6 +49,8 @@ The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. +CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. + The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. ## ROCm evidence @@ -65,9 +67,9 @@ JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C ## Validation and next GPU run -Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 30 GPU integration cases collect successfully; none has run on a GPU yet. +Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 46 GPU integration cases collect successfully; none has run on a GPU yet. -[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty groups/dimensions, backward gradients, the current stream, device guarding, and invalid counts. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. +[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. [The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. From d2eac4ea9e2237b28bcb7b0339f9237e0fc002cc Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:30:29 -0700 Subject: [PATCH 05/18] Keep Torch adapter helpers local to each extension --- .../openequivariance/extension/libtorch_tp_jit.cpp | 4 ++++ .../openequivariance/extension/libtorch_tp_jit_stable.cpp | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index c2b71841..cf458306 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -31,6 +31,8 @@ constexpr Dtype kByte = torch::kByte; #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY TORCH_LIBRARY +namespace { + class TensorDeviceGuard { c10::DeviceGuard guard; public: @@ -41,6 +43,8 @@ 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) { diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 1b15a97b..642ab2b0 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -27,6 +27,8 @@ 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: @@ -37,6 +39,8 @@ AtenTensorHandle tensor_handle(Tensor& tensor) { return tensor.get(); } +} + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { From ee1829c99f36df44b6615cdcbf99a39164ecc9b7 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:36:54 -0700 Subject: [PATCH 06/18] Verify compiled grouped GEMM accepts changing counts --- sandbox/cshim_group_gemm/DESIGN.md | 2 ++ tests/group_gemm_test.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index c5c5ea2d..e1ca774b 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -37,6 +37,8 @@ The registered operator guards A's device before making GPU contiguous copies or BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. +GPU graph capture records the group sizes read by the host loop during capture. Direct graph replay requires those counts to remain unchanged; changing group sizes requires recapture. Ordinary calls and compiled execution without graph replay read the counts on each invocation. + Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. ## ABI and build paths diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 6bf70330..33f7945d 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -260,6 +260,11 @@ def operation(A, B, counts): actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 ) + next_counts = [0, 3, 2] + next_actual = compiled(A, B, torch.tensor(next_counts, device="cpu")) + next_expected = reference(A_ref, B_ref, next_counts, inner) + torch.testing.assert_close(next_actual.cpu(), next_expected, rtol=1e-10, atol=1e-10) + @pytest.mark.parametrize("inner", [0, 1]) def test_group_gemm_aoti(group_gemm, inner, tmp_path): From e4df36cb1f998af4f3589af75582e5a5631926a9 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:04:21 -0700 Subject: [PATCH 07/18] Initialize Inductor codecache for fresh-process AOTI validation --- tests/group_gemm_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 33f7945d..8df3d5bf 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -295,6 +295,7 @@ def forward(self, A, B, counts): """ import sys import torch +import torch._inductor.codecache torch.ops.load_library(sys.argv[1]) model = torch._inductor.aoti_load_package(sys.argv[2]) From 4e3dcc8f91a5e0f78900aca9fc53c7e2df5dcfe7 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:15:19 -0700 Subject: [PATCH 08/18] Document production CUDA validation of the stable BMM shim --- sandbox/cshim_group_gemm/DESIGN.md | 25 +++------ sandbox/cshim_group_gemm/VALIDATION.md | 73 ++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 18 deletions(-) create mode 100644 sandbox/cshim_group_gemm/VALIDATION.md diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index e1ca774b..ed88da3c 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -2,7 +2,7 @@ The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -GPU execution of this implementation is pending. GPU access is paused at the user's request; the earlier prototype results below are not results for this implementation. +The production implementation passed CUDA validation on an H100 with Torch 2.10's stable and JIT extensions and Torch 2.7's JIT extension. The complete matrix and its hardware limits are recorded in [the validation report](VALIDATION.md). ## Implementation @@ -59,7 +59,7 @@ The HIP CMake target is named `oeq_stable_hip`, matching its module entry point PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution remains pending. +This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution has not been validated. ## JAX scope @@ -67,28 +67,17 @@ JAX does not call this grouped-GEMM helper and has no BLAS dependency in its ext JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). -## Validation and next GPU run +## Validation -Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 46 GPU integration cases collect successfully; none has run on a GPU yet. +The stable wheel and JIT extensions compiled and loaded successfully. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stub is excluded. GPU validation produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [Results, environment corrections, and reproduction commands](VALIDATION.md). + +Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. [The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. [The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. -After GPU access resumes, start with one case of the production operator: - -```sh -pytest -q 'tests/group_gemm_test.py::test_group_gemm_matches_reference[contiguous-0-dtype0]' -``` - -Then run the focused integration suite against the stable build and the JIT build in separate processes: - -```sh -pytest -q tests/import_test.py tests/group_gemm_test.py -OEQ_JIT_EXTENSION=1 pytest -q tests/import_test.py tests/group_gemm_test.py -``` - -Existing symmetric-contraction integration tests exercise the surrounding model and its higher-order derivatives when the optional MACE dependency is available. No end-to-end performance claim is made before measurement. +The existing symmetric-contraction integration tests passed all 24 cases for both the stable and JIT extensions under Torch 2.10, including comparison with MACE and higher-order derivatives. No end-to-end performance claim is made by this correctness validation. ## Earlier experiments diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md new file mode 100644 index 00000000..e7b0d60f --- /dev/null +++ b/sandbox/cshim_group_gemm/VALIDATION.md @@ -0,0 +1,73 @@ +# Production validation + +The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. + +## Source and artifact + +- Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. +- Production source: `d2eac4e`; integration tests: `e4df36c`. +- Wheel: `openequivariance-0.7.0-cp310-cp310-linux_x86_64.whl`. +- Wheel SHA256: `f536c53f1ef2358b6f562931c30c3f55823c32123424c153c96cb808250ea97a`. +- Remote checkout and environments: `/tmp/oeq-cshim-production/`, separate from the workspace-allocation checkout and user-installed packages. + +## Build and packaging results + +All checks in this table ran with `CUDA_VISIBLE_DEVICES=""`. They establish successful compilation, loading, and packaging; they do not establish GPU numerical correctness. + +| Configuration | Result | +| --- | --- | +| Stable wheel, Torch 2.10.0+cu128, Python 3.10 | Both production libraries built; 3 import/dependency tests passed | +| Editable install, Torch 2.10.0+cu128 | Stable extension selected with Python sources and libraries in separate directories; 3 import/dependency tests passed | +| Source/JIT extension, Torch 2.10.0+cu128 | Compiled against the installed Torch; 3 import/dependency tests passed | +| Source/JIT extension, Torch 2.7.0, CUDA 12.8 | Compiled against the installed Torch; 3 import/dependency tests passed | + +The wheel contains its Python extension, companion AOTI library, and JIT sources. It contains no `libtorch_cuda.so` build stub. ELF checks found no direct vendor BLAS library dependency or undefined vendor BLAS function in either production library or either JIT build. Both installed Torch GPU libraries export `aoti_torch_cuda_bmm_out`. + +Validation found and fixed a packaging defect: the absolute CMake install destination placed libraries in the source tree, where wheel ignore rules excluded them. The relative install destination includes the libraries in the wheel. Resolving the extension's actual import location also fixes stable-library discovery for editable installations. + +## GPU validation + +Hardware: one NVIDIA H100 PCIe, driver 580.105.08. Both Torch environments report CUDA 12.8. The first production-wheel forward case passed before running the larger suites. + +| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | +| --- | --- | --- | +| Stable wheel, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | +| Source/JIT, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | +| Source/JIT, installed Torch 2.7.0 / CUDA 12.8 | 48 passed, 1 skipped | Not run in this environment | + +The stable-wheel result combines 46 initial passes with two successful AOTI reruns after correcting fresh-process loader initialization. The Torch 2.7 result is a complete rerun after correcting test-environment dependencies. The table reports final per-case outcomes; it does not count preliminary or repeated passes twice. + +The 46 grouped-GEMM cases cover layouts, empty groups and dimensions, both modes and dtypes, backward and double backward, the current stream, device guarding, invalid counts, graph capture/replay, compiled training with changing counts, and AOTI inference. Each AOTI case loads the exported OEQ library in a new Python process without importing the OEQ Python package, then compares against a CPU reference. Three additional tests check import, successful extension loading, and binary dependencies. + +The existing symmetric-contraction suite compares with MACE for float32/float64 forward, backward, and double backward across three configurations. It also checks compile and export. Both the stable and JIT extensions passed all 24 cases. + +## Issues resolved during validation + +The wheel and editable-discovery fixes above were production changes. GPU execution required no further changes to the GEMM implementation. + +The fresh-process AOTI test exposed a Torch 2.10 loader initialization issue: its package loader accesses `torch._inductor.codecache` before importing it. The subprocess now explicitly imports that module before loading the package. Both AOTI modes then passed for the stable wheel and both JIT builds. + +The Torch 2.10 model-test environment uses MACE 0.3.16 and its pinned e3nn 0.4.4. That e3nn version loads packaged constants containing Python `slice` objects. The test process allowed that specific built-in type with `torch.serialization.add_safe_globals([slice])`; no production code or global loading policy changed. NumPy was 1.26.4 and pytest was 9.1.1. + +The Torch 2.7 environment initially inherited NetworkX 2.4 from the system; its import failed on NumPy's removed `np.int` alias before compilation. Installing NetworkX 3.4.2, SciPy 1.15.3, and SymPy 1.13.3 inside the isolated environment corrected its compiler dependencies. Its final environment used e3nn 0.6.0 and NumPy 1.26.4. System packages and the other checkout were untouched. + +## Reproduction + +After installing the wheel and test dependencies in a matching GPU Torch environment, run: + +```sh +python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +OEQ_JIT_EXTENSION=1 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +``` + +The allowlist is only needed for the older e3nn dependency described above. The actual runs used separate JIT and Inductor cache directories for each Torch environment and ran the focused and model suites separately. The Torch 2.7 environment also set `PYTHONNOUSERSITE=1` to exclude the shared machine's user packages. + +## Scope limits + +- ROCm has source-level support evidence in [the design](DESIGN.md), but no AMD hardware run has been performed. +- Multi-GPU device guarding remains unverified on hardware because only one GPU was available. +- Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. +- Direct GPU graph replay requires fixed ragged counts, as described in the design. +- No end-to-end performance claim is made by these correctness and packaging checks. + +Detailed logs, environment snapshots, and JUnit XML are retained under `/tmp/oeq-cshim-production/logs/` on the validation machine. At the end of testing, the GPU reported zero memory use and no compute processes; its status reservation was released. From 7399be668a1aa9caa2c2e426f1f24bd24a6acf68 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:32:03 -0700 Subject: [PATCH 09/18] Document AMD hardware validation of the stable BMM shim --- sandbox/cshim_group_gemm/AMD_VALIDATION.md | 75 ++++++++++++++++++++++ sandbox/cshim_group_gemm/DESIGN.md | 8 +-- sandbox/cshim_group_gemm/VALIDATION.md | 3 +- 3 files changed, 81 insertions(+), 5 deletions(-) create mode 100644 sandbox/cshim_group_gemm/AMD_VALIDATION.md diff --git a/sandbox/cshim_group_gemm/AMD_VALIDATION.md b/sandbox/cshim_group_gemm/AMD_VALIDATION.md new file mode 100644 index 00000000..a928d44c --- /dev/null +++ b/sandbox/cshim_group_gemm/AMD_VALIDATION.md @@ -0,0 +1,75 @@ +# AMD production validation + +ROCm validation passed on 2026-09-13 without changing the GEMM implementation. The final matrix contains 192 successful checks and three skips for the device-guard test requiring two GPUs. The normal ROCm JIT path passed with Torch 2.10 and 2.7. The stable HIP artifacts also passed with the isolated loader override described below. + +## Source and environment + +- Branch: `move-bmm-calls-to-stable-cshim`, tested at `4e3dcc8`. +- GPU: one AMD Instinct MI300X VF, 192 GB, `gfx942`. +- Driver reported by `rocm-smi`: `6.19.14.31400000`; Linux kernel `6.8.0-138-generic`. +- Python: 3.12.3. +- Build SDK: `/opt/rocm/core-10.0`, reporting HIP `7.15.26333`; host compiler GCC 13.3.0. +- Test environments: official Torch `2.10.0+rocm7.1` and `2.7.0+rocm6.3` wheels in separate virtual environments. +- Checkout, environments, caches, and logs: `/tmp/oeq-cshim-amd/` on the AMD machine. + +The NVIDIA machine was not used during this validation. + +## Results + +The first production grouped-GEMM case passed before running the focused suites. + +| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | +| --- | --- | --- | +| Normal JIT path, Torch 2.10 / ROCm 7.1 | 48 passed, 1 skipped | 24 passed | +| Stable HIP artifacts, Torch 2.10 / ROCm 7.1, explicit test loader override | 48 passed, 1 skipped | 24 passed | +| Normal JIT path, Torch 2.7 / ROCm 6.3 | 48 passed, 1 skipped | Not run in this environment | + +The focused tests cover both GEMM modes, float32/float64, noncontiguous inputs and counts, storage offsets, empty groups and dimensions, backward and double backward, current-stream behavior, graph capture/replay, compiled training with changing counts, and AOTI inference in a fresh process. The model suite compares with MACE for forward, backward, and double backward across three configurations and both dtypes, and checks compile and export. + +Additional checks passed outside this matrix: the initial single GPU case, Torch 2.7 JIT compilation/import with GPU visibility disabled, and three editable-install import/dependency tests with GPU visibility disabled. They are not counted again in the 192-check total. + +## Build, packaging, and BLAS dependency checks + +The HIP wheel built successfully using the branch's CMake configuration: + +- Artifact: `openequivariance-0.7.0-cp312-cp312-linux_x86_64.whl`. +- SHA256: `8ac6dc8cb3313298e613e7425e255fc369f89c797fbdefa1126778932fe76fb8`. +- Packaged libraries: `oeq_stable_hip.cpython-312-x86_64-linux-gnu.so` and `liboeq_stable_hip_aoti.so`. +- The wheel includes the JIT sources and excludes the `libtorch_hip.so` build stub. + +ELF inspection of both stable libraries and both JIT libraries found no direct cuBLAS, hipBLAS, or rocBLAS library dependency and no undefined vendor BLAS functions. All four libraries reference `aoti_torch_cuda_bmm_out`. Both installed ROCm Torch libraries export that symbol and `aoti_torch_get_current_cuda_stream`, retaining the CUDA spelling on AMD. + +The stable libraries depend on `libtorch_hip.so`, `libhiprtc.so.7`, and `libamdhip64.so.7`, along with CPU Torch and standard system libraries. The JIT builds link the Torch GPU library and the HIPRTC library found for their environments: `.so.7` for Torch 2.10 and `.so.6` for Torch 2.7. BLAS selection and handle ownership remain inside Torch. + +The actual editable installation also passed. Its Python source resolved to the checkout, its compiled HIP extension resolved to the virtual environment, and `_has_precompiled_extension()` found the compiled artifacts. The existing ROCm policy selected JIT as intended. The normal wheel was restored afterward. + +## Stable HIP loader policy + +The production loader still unconditionally selects JIT on ROCm. This validation does not enable automatic stable HIP loading. + +To exercise the compiled stable implementation, the wheel was extracted into `/tmp/oeq-cshim-amd/stable210-package`. Only the unconditional HIP precompiled-disable block was removed from that copy of the Python loader. Both shared libraries were verified byte-for-byte against the wheel. The tests asserted that the stable extension was selected and exercised the real registered operator, its existing autograd registration, and the companion AOTI library. No numerical operations were mocked. + +This establishes hardware correctness for the tested stable artifacts while keeping the production loader policy unchanged. + +## Setup and reproduction + +The new machine lacked Python packaging bootstrap support and Python development headers. Pip was bootstrapped inside the isolated environments, and `python3.12-dev` was installed for extension compilation. The initial CMake attempt stopped at the missing Python headers; compilation succeeded after installing them. No GPU test failed, and no production source changes were needed. + +The model-test environment used MACE 0.3.16, its pinned e3nn 0.4.4, NumPy 1.26.4, and pytest 9.1.1. As in the CUDA run, the test process allowed the built-in `slice` type for e3nn's packaged constants. The Torch 2.7 environment used e3nn 0.6.0 and NumPy 1.26.4. + +With the wheel and dependencies installed, the normal ROCm path is exercised by: + +```sh +ROCM_HOME=/opt/rocm/core-10.0 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +``` + +The actual runs used separate JIT and Inductor caches for each Torch version, and ran the focused and model suites separately. For stable-artifact validation, `PYTHONPATH` selected the isolated wheel copy described above. The older Torch environment ran the focused suite without the optional MACE dependency. Full environment snapshots, build logs, binary inspection output, and JUnit XML are retained in `/tmp/oeq-cshim-amd/logs/`. + +## Scope limits + +- Only one AMD GPU was available, so multi-GPU device guarding remains unverified. +- Automatic stable HIP loading remains disabled by existing policy. +- These are correctness and packaging results, with no performance claim. +- Direct GPU graph replay requires fixed ragged counts. + +At the end of GPU testing, `rocm-smi` reported zero GPU utilization, zero allocated VRAM percentage, and no KFD processes. The GPU reservation was released; remaining installation checks used disabled GPU visibility. diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index ed88da3c..19be961f 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -2,7 +2,7 @@ The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -The production implementation passed CUDA validation on an H100 with Torch 2.10's stable and JIT extensions and Torch 2.7's JIT extension. The complete matrix and its hardware limits are recorded in [the validation report](VALIDATION.md). +The production implementation passed CUDA validation on an H100 and ROCm validation on an MI300X. Both runs covered Torch 2.10's stable and JIT implementations and Torch 2.7's JIT implementation; stable HIP artifact testing used an isolated loader override because production still selects JIT on ROCm. The complete matrices and hardware limits are recorded in the [CUDA report](VALIDATION.md) and [AMD report](AMD_VALIDATION.md). ## Implementation @@ -53,13 +53,13 @@ The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](. CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. -The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. +The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. The HIP wheel built successfully, and its stable artifacts passed hardware tests with the existing loader restriction removed only in an isolated test copy. ## ROCm evidence PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution has not been validated. +Hardware validation confirmed these entry points in Torch 2.10 / ROCm 7.1 and Torch 2.7 / ROCm 6.3. On an MI300X, both normal JIT configurations and the stable HIP artifacts passed the focused integration suite; both Torch 2.10 implementations also passed all model tests. The AMD matrix contains 192 successful checks and three skips requiring a second GPU. Binary inspection confirmed no direct vendor BLAS dependency in either stable library or either JIT library. [AMD validation and loader-policy details](AMD_VALIDATION.md). ## JAX scope @@ -69,7 +69,7 @@ JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C ## Validation -The stable wheel and JIT extensions compiled and loaded successfully. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stub is excluded. GPU validation produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [Results, environment corrections, and reproduction commands](VALIDATION.md). +The stable wheels and JIT extensions compiled and loaded successfully on CUDA and ROCm. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stubs are excluded. Each platform's validation matrix produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [CUDA results](VALIDATION.md), [AMD results](AMD_VALIDATION.md). Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md index e7b0d60f..c45a0565 100644 --- a/sandbox/cshim_group_gemm/VALIDATION.md +++ b/sandbox/cshim_group_gemm/VALIDATION.md @@ -2,6 +2,8 @@ The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. +Subsequent ROCm hardware validation is recorded in [the AMD report](AMD_VALIDATION.md). + ## Source and artifact - Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. @@ -64,7 +66,6 @@ The allowlist is only needed for the older e3nn dependency described above. The ## Scope limits -- ROCm has source-level support evidence in [the design](DESIGN.md), but no AMD hardware run has been performed. - Multi-GPU device guarding remains unverified on hardware because only one GPU was available. - Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. - Direct GPU graph replay requires fixed ragged counts, as described in the design. From 99d06eade512235c635d84ede366b71fb9c193ae Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:13:33 -0700 Subject: [PATCH 10/18] human changelog --- CHANGELOG.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 22257be3..221418d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,6 @@ ## Latest Changes -- Route Torch grouped GEMM through the stable BMM C shim on CUDA and ROCm, - using Torch's current device stream and precision settings. Remove OEQ's - direct cuBLAS and rocBLAS dependencies. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's c shim to instead. ### v0.7.0 (2026-09-10) **Added**: From 4314273282c2f3a256cfac958bfef1eaba147243 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:20:30 -0700 Subject: [PATCH 11/18] minimize diff, remove files that should not be tracked --- .github/workflows/verify_extension_build.yml | 3 +- sandbox/cshim_group_gemm/AMD_VALIDATION.md | 75 --------------- sandbox/cshim_group_gemm/DESIGN.md | 88 ----------------- sandbox/cshim_group_gemm/VALIDATION.md | 74 --------------- sandbox/cshim_group_gemm/group_gemm.cpp | 99 -------------------- sandbox/cshim_group_gemm/smoke_test.py | 71 -------------- tests/extension_linkage_test.py | 28 ++++++ tests/import_test.py | 28 ------ 8 files changed, 29 insertions(+), 437 deletions(-) delete mode 100644 sandbox/cshim_group_gemm/AMD_VALIDATION.md delete mode 100644 sandbox/cshim_group_gemm/DESIGN.md delete mode 100644 sandbox/cshim_group_gemm/VALIDATION.md delete mode 100644 sandbox/cshim_group_gemm/group_gemm.cpp delete mode 100644 sandbox/cshim_group_gemm/smoke_test.py create mode 100644 tests/extension_linkage_test.py diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index 69217eaf..6c903205 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -29,11 +29,10 @@ jobs: sudo apt-get update sudo apt install nvidia-cuda-toolkit pip install -r .github/workflows/requirements_cuda_ci.txt - pip install "./openequivariance[jax]" + pip install -e "./openequivariance[jax]" - name: Test CUDA extension build via import run: | - python -c "import openequivariance as oeq; assert oeq.USE_PRECOMPILED_EXTENSION" pytest tests/import_test.py export OEQ_JIT_EXTENSION=1 diff --git a/sandbox/cshim_group_gemm/AMD_VALIDATION.md b/sandbox/cshim_group_gemm/AMD_VALIDATION.md deleted file mode 100644 index a928d44c..00000000 --- a/sandbox/cshim_group_gemm/AMD_VALIDATION.md +++ /dev/null @@ -1,75 +0,0 @@ -# AMD production validation - -ROCm validation passed on 2026-09-13 without changing the GEMM implementation. The final matrix contains 192 successful checks and three skips for the device-guard test requiring two GPUs. The normal ROCm JIT path passed with Torch 2.10 and 2.7. The stable HIP artifacts also passed with the isolated loader override described below. - -## Source and environment - -- Branch: `move-bmm-calls-to-stable-cshim`, tested at `4e3dcc8`. -- GPU: one AMD Instinct MI300X VF, 192 GB, `gfx942`. -- Driver reported by `rocm-smi`: `6.19.14.31400000`; Linux kernel `6.8.0-138-generic`. -- Python: 3.12.3. -- Build SDK: `/opt/rocm/core-10.0`, reporting HIP `7.15.26333`; host compiler GCC 13.3.0. -- Test environments: official Torch `2.10.0+rocm7.1` and `2.7.0+rocm6.3` wheels in separate virtual environments. -- Checkout, environments, caches, and logs: `/tmp/oeq-cshim-amd/` on the AMD machine. - -The NVIDIA machine was not used during this validation. - -## Results - -The first production grouped-GEMM case passed before running the focused suites. - -| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | -| --- | --- | --- | -| Normal JIT path, Torch 2.10 / ROCm 7.1 | 48 passed, 1 skipped | 24 passed | -| Stable HIP artifacts, Torch 2.10 / ROCm 7.1, explicit test loader override | 48 passed, 1 skipped | 24 passed | -| Normal JIT path, Torch 2.7 / ROCm 6.3 | 48 passed, 1 skipped | Not run in this environment | - -The focused tests cover both GEMM modes, float32/float64, noncontiguous inputs and counts, storage offsets, empty groups and dimensions, backward and double backward, current-stream behavior, graph capture/replay, compiled training with changing counts, and AOTI inference in a fresh process. The model suite compares with MACE for forward, backward, and double backward across three configurations and both dtypes, and checks compile and export. - -Additional checks passed outside this matrix: the initial single GPU case, Torch 2.7 JIT compilation/import with GPU visibility disabled, and three editable-install import/dependency tests with GPU visibility disabled. They are not counted again in the 192-check total. - -## Build, packaging, and BLAS dependency checks - -The HIP wheel built successfully using the branch's CMake configuration: - -- Artifact: `openequivariance-0.7.0-cp312-cp312-linux_x86_64.whl`. -- SHA256: `8ac6dc8cb3313298e613e7425e255fc369f89c797fbdefa1126778932fe76fb8`. -- Packaged libraries: `oeq_stable_hip.cpython-312-x86_64-linux-gnu.so` and `liboeq_stable_hip_aoti.so`. -- The wheel includes the JIT sources and excludes the `libtorch_hip.so` build stub. - -ELF inspection of both stable libraries and both JIT libraries found no direct cuBLAS, hipBLAS, or rocBLAS library dependency and no undefined vendor BLAS functions. All four libraries reference `aoti_torch_cuda_bmm_out`. Both installed ROCm Torch libraries export that symbol and `aoti_torch_get_current_cuda_stream`, retaining the CUDA spelling on AMD. - -The stable libraries depend on `libtorch_hip.so`, `libhiprtc.so.7`, and `libamdhip64.so.7`, along with CPU Torch and standard system libraries. The JIT builds link the Torch GPU library and the HIPRTC library found for their environments: `.so.7` for Torch 2.10 and `.so.6` for Torch 2.7. BLAS selection and handle ownership remain inside Torch. - -The actual editable installation also passed. Its Python source resolved to the checkout, its compiled HIP extension resolved to the virtual environment, and `_has_precompiled_extension()` found the compiled artifacts. The existing ROCm policy selected JIT as intended. The normal wheel was restored afterward. - -## Stable HIP loader policy - -The production loader still unconditionally selects JIT on ROCm. This validation does not enable automatic stable HIP loading. - -To exercise the compiled stable implementation, the wheel was extracted into `/tmp/oeq-cshim-amd/stable210-package`. Only the unconditional HIP precompiled-disable block was removed from that copy of the Python loader. Both shared libraries were verified byte-for-byte against the wheel. The tests asserted that the stable extension was selected and exercised the real registered operator, its existing autograd registration, and the companion AOTI library. No numerical operations were mocked. - -This establishes hardware correctness for the tested stable artifacts while keeping the production loader policy unchanged. - -## Setup and reproduction - -The new machine lacked Python packaging bootstrap support and Python development headers. Pip was bootstrapped inside the isolated environments, and `python3.12-dev` was installed for extension compilation. The initial CMake attempt stopped at the missing Python headers; compilation succeeded after installing them. No GPU test failed, and no production source changes were needed. - -The model-test environment used MACE 0.3.16, its pinned e3nn 0.4.4, NumPy 1.26.4, and pytest 9.1.1. As in the CUDA run, the test process allowed the built-in `slice` type for e3nn's packaged constants. The Torch 2.7 environment used e3nn 0.6.0 and NumPy 1.26.4. - -With the wheel and dependencies installed, the normal ROCm path is exercised by: - -```sh -ROCM_HOME=/opt/rocm/core-10.0 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -``` - -The actual runs used separate JIT and Inductor caches for each Torch version, and ran the focused and model suites separately. For stable-artifact validation, `PYTHONPATH` selected the isolated wheel copy described above. The older Torch environment ran the focused suite without the optional MACE dependency. Full environment snapshots, build logs, binary inspection output, and JUnit XML are retained in `/tmp/oeq-cshim-amd/logs/`. - -## Scope limits - -- Only one AMD GPU was available, so multi-GPU device guarding remains unverified. -- Automatic stable HIP loading remains disabled by existing policy. -- These are correctness and packaging results, with no performance claim. -- Direct GPU graph replay requires fixed ragged counts. - -At the end of GPU testing, `rocm-smi` reported zero GPU utilization, zero allocated VRAM percentage, and no KFD processes. The GPU reservation was released; remaining installation checks used disabled GPU visibility. diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md deleted file mode 100644 index 19be961f..00000000 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ /dev/null @@ -1,88 +0,0 @@ -# Move Torch cuBLAS calls to the stable C shim - -The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. - -The production implementation passed CUDA validation on an H100 and ROCm validation on an MI300X. Both runs covered Torch 2.10's stable and JIT implementations and Torch 2.7's JIT implementation; stable HIP artifact testing used an isolated loader override because production still selects JIT on ROCm. The complete matrices and hardware limits are recorded in the [CUDA report](VALIDATION.md) and [AMD report](AMD_VALIDATION.md). - -## Implementation - -The implementation lives in [group_mm.hpp](../../openequivariance/openequivariance/extension/group_mm.hpp), with input validation and output allocation in [torch_core.hpp](../../openequivariance/openequivariance/extension/torch_core.hpp). The algorithm is one BMM-out invocation per nonempty ragged group, using the same interleaved layout as the previous strided batched GEMMs. It does not need a vendor's heterogeneous grouped-GEMM API. - -Inputs are made contiguous by the registered operator, then passed as borrowed `AtenTensorHandle` values to the shared helper. The stable extension gets these handles from `torch::stable::Tensor::get()`. The source/JIT extension uses Torch's `tensor_pointer_to_tensor_handle` utility to borrow a handle to its local `at::Tensor`; that bridge is compiled against the installed Torch, as the existing JIT extension already is. ATen C++ objects do not cross the stable extension's ABI boundary. - -The helper creates temporary views with `aoti_torch__reinterpret_tensor`. These views retain the original storage and add their element offsets to the original tensor's storage offset. Their handles are released through RAII after each call, including when an error is returned. Original input and output handles remain owned by the caller. This replaces the prototype's `from_blob` construction and avoids reconstructing storage or querying pointer devices for every group. [Torch reinterpretation API](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L435). - -The operation checks GPU placement, matching devices and float32/float64 dtypes, tensor shapes, nonnegative dimensions, and mode 0 or 1. Ragged counts must be a CPU int64 vector with one entry per group. Counts may be noncontiguous; the operator makes a contiguous CPU copy before reading them. Counts must be nonnegative and sum to the input row count. Incremental bounds checking avoids overflowing the count sum. These checks establish the bounds needed by the storage-view helper. - -Output allocation uses Torch and starts at zero. Empty groups are skipped; in mode 1 this leaves the corresponding weight-gradient block zero. Zero batch, output, or contraction dimensions return the zero-initialized output without invoking BMM. Dimensions and offsets remain int64 throughout the grouped-GEMM implementation. - -## Layout mapping - -Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Offsets are relative to the logical beginning of each original tensor; strides and offsets are in elements. - -| Mode | Tensor | Element offset | View shape | View strides | -| --- | --- | --- | --- | --- | -| 0 | B / input | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| 0 | A / weights | `i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | -| 0 | C / output | `o*b*m` | `[b,n,m]` | `[m,b*m,1]` | -| 1 | A / left | `o*b*m` | `[b,m,n]` | `[m,1,b*m]` | -| 1 | B / right | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| 1 | C / output | `i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | - -The final operation for each group is `aoti_torch_cuda_bmm_out(output, left, right)`. The API overwrites the output with alpha 1 and beta 0. Torch may choose its own BLAS kernel or backend. The layout code requests views, but makes no unconditional promise about copies inside Torch's BMM implementation. [Torch BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). - -## Device, stream, precision, and autograd - -The registered operator guards A's device before making GPU contiguous copies or allocating the output. The stable build uses `torch::stable::accelerator::DeviceGuard`; the JIT build uses `c10::DeviceGuard`. Each restores the caller's current device on return. - -BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. - -GPU graph capture records the group sizes read by the host loop during capture. Direct graph replay requires those counts to remain unchanged; changing group sizes requires recapture. Ordinary calls and compiled execution without graph replay read the counts on each invocation. - -Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. - -## ABI and build paths - -The official declaration is in `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`. The backend-specific spelling is used instead of the deprecated `aoti_torch_bmm_out`. [Declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74), [stable ABI policy](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). - -Both `aoti_torch_cuda_bmm_out` and `aoti_torch__reinterpret_tensor` are present in the inspected PyTorch 2.4 headers, matching OEQ's documented source/JIT baseline. The stable build explicitly targets the 2.10 ABI and continues to use the pinned LibTorch headers/libraries. This uses Torch's stated ABI guarantees; it does not try to select cuBLAS versions or infer handle ownership from version queries. - -The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` library supplying the BMM shim. CUDA also retains its driver/runtime/NVRTC dependencies; ROCm links hipRTC. All direct cuBLAS/rocBLAS calls, handle management, and corresponding OEQ link dependencies have been removed. - -The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. - -CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. - -The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. The HIP wheel built successfully, and its stable artifacts passed hardware tests with the existing loader restriction removed only in an isolated test copy. - -## ROCm evidence - -PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). - -Hardware validation confirmed these entry points in Torch 2.10 / ROCm 7.1 and Torch 2.7 / ROCm 6.3. On an MI300X, both normal JIT configurations and the stable HIP artifacts passed the focused integration suite; both Torch 2.10 implementations also passed all model tests. The AMD matrix contains 192 successful checks and three skips requiring a second GPU. Binary inspection confirmed no direct vendor BLAS dependency in either stable library or either JIT library. [AMD validation and loader-policy details](AMD_VALIDATION.md). - -## JAX scope - -JAX does not call this grouped-GEMM helper and has no BLAS dependency in its extension target. The replacement is entirely within the Torch frontend. The common CUDA/HIP kernel compilation backend gains no Torch dependency. - -JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). - -## Validation - -The stable wheels and JIT extensions compiled and loaded successfully on CUDA and ROCm. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stubs are excluded. Each platform's validation matrix produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [CUDA results](VALIDATION.md), [AMD results](AMD_VALIDATION.md). - -Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. - -[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. - -[The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. - -The existing symmetric-contraction integration tests passed all 24 cases for both the stable and JIT extensions under Torch 2.10, including comparison with MACE and higher-order derivatives. No end-to-end performance claim is made by this correctness validation. - -## Earlier experiments - -The original [standalone prototype](group_gemm.cpp) and its [smoke script](smoke_test.py) remain as historical experiment artifacts. They use blob views rather than the production helper's views of existing storage. - -The prototype passed one float32 forward case on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against a CPU float64 reference. Earlier ctypes calls to the BMM C shim passed 40 small cases across both dtypes/modes. Four selected BMM-only profiles showed no copy or GPU allocation inside the call; they did not measure view creation overhead. - -The earlier direct-cuBLAS experiments explain the migration: matching handle/function owners worked, whereas some separately loaded foreign-owner combinations failed or crashed despite successful version queries. They do not establish interchangeability of private BLAS handles. Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md deleted file mode 100644 index c45a0565..00000000 --- a/sandbox/cshim_group_gemm/VALIDATION.md +++ /dev/null @@ -1,74 +0,0 @@ -# Production validation - -The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. - -Subsequent ROCm hardware validation is recorded in [the AMD report](AMD_VALIDATION.md). - -## Source and artifact - -- Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. -- Production source: `d2eac4e`; integration tests: `e4df36c`. -- Wheel: `openequivariance-0.7.0-cp310-cp310-linux_x86_64.whl`. -- Wheel SHA256: `f536c53f1ef2358b6f562931c30c3f55823c32123424c153c96cb808250ea97a`. -- Remote checkout and environments: `/tmp/oeq-cshim-production/`, separate from the workspace-allocation checkout and user-installed packages. - -## Build and packaging results - -All checks in this table ran with `CUDA_VISIBLE_DEVICES=""`. They establish successful compilation, loading, and packaging; they do not establish GPU numerical correctness. - -| Configuration | Result | -| --- | --- | -| Stable wheel, Torch 2.10.0+cu128, Python 3.10 | Both production libraries built; 3 import/dependency tests passed | -| Editable install, Torch 2.10.0+cu128 | Stable extension selected with Python sources and libraries in separate directories; 3 import/dependency tests passed | -| Source/JIT extension, Torch 2.10.0+cu128 | Compiled against the installed Torch; 3 import/dependency tests passed | -| Source/JIT extension, Torch 2.7.0, CUDA 12.8 | Compiled against the installed Torch; 3 import/dependency tests passed | - -The wheel contains its Python extension, companion AOTI library, and JIT sources. It contains no `libtorch_cuda.so` build stub. ELF checks found no direct vendor BLAS library dependency or undefined vendor BLAS function in either production library or either JIT build. Both installed Torch GPU libraries export `aoti_torch_cuda_bmm_out`. - -Validation found and fixed a packaging defect: the absolute CMake install destination placed libraries in the source tree, where wheel ignore rules excluded them. The relative install destination includes the libraries in the wheel. Resolving the extension's actual import location also fixes stable-library discovery for editable installations. - -## GPU validation - -Hardware: one NVIDIA H100 PCIe, driver 580.105.08. Both Torch environments report CUDA 12.8. The first production-wheel forward case passed before running the larger suites. - -| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | -| --- | --- | --- | -| Stable wheel, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | -| Source/JIT, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | -| Source/JIT, installed Torch 2.7.0 / CUDA 12.8 | 48 passed, 1 skipped | Not run in this environment | - -The stable-wheel result combines 46 initial passes with two successful AOTI reruns after correcting fresh-process loader initialization. The Torch 2.7 result is a complete rerun after correcting test-environment dependencies. The table reports final per-case outcomes; it does not count preliminary or repeated passes twice. - -The 46 grouped-GEMM cases cover layouts, empty groups and dimensions, both modes and dtypes, backward and double backward, the current stream, device guarding, invalid counts, graph capture/replay, compiled training with changing counts, and AOTI inference. Each AOTI case loads the exported OEQ library in a new Python process without importing the OEQ Python package, then compares against a CPU reference. Three additional tests check import, successful extension loading, and binary dependencies. - -The existing symmetric-contraction suite compares with MACE for float32/float64 forward, backward, and double backward across three configurations. It also checks compile and export. Both the stable and JIT extensions passed all 24 cases. - -## Issues resolved during validation - -The wheel and editable-discovery fixes above were production changes. GPU execution required no further changes to the GEMM implementation. - -The fresh-process AOTI test exposed a Torch 2.10 loader initialization issue: its package loader accesses `torch._inductor.codecache` before importing it. The subprocess now explicitly imports that module before loading the package. Both AOTI modes then passed for the stable wheel and both JIT builds. - -The Torch 2.10 model-test environment uses MACE 0.3.16 and its pinned e3nn 0.4.4. That e3nn version loads packaged constants containing Python `slice` objects. The test process allowed that specific built-in type with `torch.serialization.add_safe_globals([slice])`; no production code or global loading policy changed. NumPy was 1.26.4 and pytest was 9.1.1. - -The Torch 2.7 environment initially inherited NetworkX 2.4 from the system; its import failed on NumPy's removed `np.int` alias before compilation. Installing NetworkX 3.4.2, SciPy 1.15.3, and SymPy 1.13.3 inside the isolated environment corrected its compiler dependencies. Its final environment used e3nn 0.6.0 and NumPy 1.26.4. System packages and the other checkout were untouched. - -## Reproduction - -After installing the wheel and test dependencies in a matching GPU Torch environment, run: - -```sh -python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -OEQ_JIT_EXTENSION=1 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -``` - -The allowlist is only needed for the older e3nn dependency described above. The actual runs used separate JIT and Inductor cache directories for each Torch environment and ran the focused and model suites separately. The Torch 2.7 environment also set `PYTHONNOUSERSITE=1` to exclude the shared machine's user packages. - -## Scope limits - -- Multi-GPU device guarding remains unverified on hardware because only one GPU was available. -- Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. -- Direct GPU graph replay requires fixed ragged counts, as described in the design. -- No end-to-end performance claim is made by these correctness and packaging checks. - -Detailed logs, environment snapshots, and JUnit XML are retained under `/tmp/oeq-cshim-production/logs/` on the validation machine. At the end of testing, the GPU reported zero memory use and no compute processes; its status reservation was released. diff --git a/sandbox/cshim_group_gemm/group_gemm.cpp b/sandbox/cshim_group_gemm/group_gemm.cpp deleted file mode 100644 index 1f7462a7..00000000 --- a/sandbox/cshim_group_gemm/group_gemm.cpp +++ /dev/null @@ -1,99 +0,0 @@ -// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. -#include - -#include -#include -#include -#include -#include -#include - -namespace { - -void check(AOTITorchError status) { - if (status != AOTI_TORCH_SUCCESS) - throw std::runtime_error("PyTorch C shim failed"); -} - -using Tensor = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_tensor_object)>; -using DeviceGuard = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_cuda_guard)>; - -// Only tensor metadata is created; the caller retains ownership of the buffer. -Tensor view(void* data, std::array sizes, - std::array strides, int32_t dtype, int32_t device) { - AtenTensorHandle tensor = nullptr; - check(aoti_torch_create_tensor_from_blob_v2( - data, 3, sizes.data(), strides.data(), 0, dtype, - aoti_torch_device_type_cuda(), device, &tensor, - aoti_torch_layout_strided(), nullptr, 0)); - return Tensor(tensor, aoti_torch_delete_tensor_object); -} - -} // namespace - -template -void group_gemm_cshim( - T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, - int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { - static_assert(std::is_same_v || std::is_same_v); - const int32_t dtype = std::is_same_v - ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); - - CUDAGuardHandle raw_guard = nullptr; - check(aoti_torch_create_cuda_guard(device, &raw_guard)); - DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); - - int64_t offset = 0; - for (int i = 0; i < num_groups; ++i) { - const int64_t n = ragged_counts[i]; - if (n == 0) continue; // Preserve the original empty-group behavior. - - if (ragged_inner == 0) { - // [batch, n, k] @ [batch, k, m] -> [batch, n, m] - auto input = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto weight = view(A + i * batch * m * k, - {batch, k, m}, {m * k, 1, k}, dtype, device); - auto output = view(C + offset * batch * m, - {batch, n, m}, {m, batch * m, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); - } else { - // [batch, m, n] @ [batch, n, k] -> [batch, m, k] - auto left = view(A + offset * batch * m, - {batch, m, n}, {m, 1, batch * m}, dtype, device); - auto right = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto output = view(C + i * batch * m * k, - {batch, m, k}, {m * k, k, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); - } - offset += n; - } -} - -// Small ctypes entry point for the smoke test, not production registration code. -extern "C" int oeq_group_gemm_cshim( - int dtype, void* A, void* B, void* C, const int64_t* counts, - int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { - try { - if (dtype == 0) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else if (dtype == 1) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else { - throw std::runtime_error("Expected float32 or float64"); - } - return 0; - } catch (const std::exception& error) { - std::fprintf(stderr, "%s\n", error.what()); - return 1; - } -} diff --git a/sandbox/cshim_group_gemm/smoke_test.py b/sandbox/cshim_group_gemm/smoke_test.py deleted file mode 100644 index 931d9698..00000000 --- a/sandbox/cshim_group_gemm/smoke_test.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Compile the C++ prototype and run one float32 CUDA correctness case.""" -import ctypes -from pathlib import Path -import subprocess - -import torch - - -def main(): - root = Path(__file__).resolve().parent - torch_root = Path(torch.__file__).resolve().parent - library = root / "group_gemm.so" - subprocess.run( - [ - "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", - f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), - f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", - "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", - "-o", str(library), - ], - check=True, - ) - lib = ctypes.CDLL(str(library)) - call = lib.oeq_group_gemm_cshim - i64 = ctypes.c_int64 - call.argtypes = [ - ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, - ctypes.c_int, ctypes.c_int32, - ] - call.restype = ctypes.c_int - - # One case includes differently sized groups, an empty group, interleaved - # batches, and a nondefault stream. The reference runs on the CPU. - counts = [2, 0, 5, 1] - batch, m, k = 3, 4, 5 - torch.manual_seed(123) - weights_cpu = torch.randn(len(counts), batch, m, k) - input_cpu = torch.randn(sum(counts), batch, k) - expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) - offset = 0 - for i, n in enumerate(counts): - expected[offset:offset + n] = torch.einsum( - "bmk,nbk->nbm", weights_cpu[i].double(), - input_cpu[offset:offset + n].double(), - ) - offset += n - - torch.cuda.set_device(0) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - weights = weights_cpu.cuda() - inputs = input_cpu.cuda() - output = torch.full(expected.shape, float("nan"), device="cuda") - status = call( - 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), - (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, - ) - assert status == 0, f"C shim prototype returned {status}" - stream.synchronize() - actual = output.cpu().double() - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) - print( - f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " - f"max absolute error={(actual - expected).abs().max().item():.3g}; " - f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" - ) - - -if __name__ == "__main__": - main() diff --git a/tests/extension_linkage_test.py b/tests/extension_linkage_test.py new file mode 100644 index 00000000..77fbb8bc --- /dev/null +++ b/tests/extension_linkage_test.py @@ -0,0 +1,28 @@ +import re +import shutil +import subprocess + +import pytest + + +def test_extension_has_no_direct_blas_dependency(): + if not shutil.which("readelf"): + pytest.skip("readelf is required to inspect ELF dependencies") + import openequivariance + from openequivariance._torch.extlib import extension_module + + paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} + for path in paths: + dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) + dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] + assert not any( + re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies + ) + symbols = subprocess.check_output( + ["readelf", "--dyn-syms", "--wide", path], text=True + ) + undefined = [line for line in symbols.splitlines() if " UND " in line] + assert not any( + re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) + for line in undefined + ) diff --git a/tests/import_test.py b/tests/import_test.py index aa53630e..bf26af31 100644 --- a/tests/import_test.py +++ b/tests/import_test.py @@ -1,9 +1,4 @@ from importlib.metadata import version -import re -import shutil -import subprocess - -import pytest def test_import(): @@ -19,26 +14,3 @@ def test_extension_built(): assert BUILT_EXTENSION_ERROR is None assert BUILT_EXTENSION - - -def test_extension_has_no_direct_blas_dependency(): - if not shutil.which("readelf"): - pytest.skip("readelf is required to inspect ELF dependencies") - import openequivariance - from openequivariance._torch.extlib import extension_module - - paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} - for path in paths: - dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) - dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] - assert not any( - re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies - ) - symbols = subprocess.check_output( - ["readelf", "--dyn-syms", "--wide", path], text=True - ) - undefined = [line for line in symbols.splitlines() if " UND " in line] - assert not any( - re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) - for line in undefined - ) From cc02c70732ba32a1273226c673fac56c8054b30b Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:23:11 -0700 Subject: [PATCH 12/18] specify cuda c shim --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 221418d0..958161f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## Latest Changes -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's c shim to instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim to instead. ### v0.7.0 (2026-09-10) **Added**: From 62eae5ced074efc3db40e87a058e61cc3fca6395 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:01:51 -0700 Subject: [PATCH 13/18] remove target version to minimize PR, this will go on a separate PR --- openequivariance/CMakeLists.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 0aef37e0..c5a1142f 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -67,7 +67,6 @@ function(add_stable_extension target_name backend_define link_libraries) # Enforce CXX11 ABI to match LibTorch target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 - TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION ) @@ -98,7 +97,6 @@ function(add_stable_extension target_name backend_define link_libraries) target_compile_definitions(${aoti_target_name} PRIVATE ${backend_define}=1 - TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 ) From 15f766a404fa6087fbe3e648e64e961f4e9ddcac Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:19:58 -0700 Subject: [PATCH 14/18] simply PR my moving build issues to a separate pr --- openequivariance/CMakeLists.txt | 2 +- .../openequivariance/_torch/extlib/__init__.py | 15 +++------------ 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index c5a1142f..3e799a82 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -52,7 +52,7 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) # Create nanobind extension diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 6785c230..ed41bfa6 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,7 +5,6 @@ import warnings import sysconfig import contextlib -import importlib.util from pathlib import Path from packaging.version import Version @@ -181,16 +180,6 @@ def load_precompiled_extension(): ) -def _has_precompiled_extension(): - backend = "hip" if IS_HIP else "cuda" - spec = importlib.util.find_spec(f"{__name__}.oeq_stable_{backend}") - return ( - spec is not None - and spec.origin is not None - and Path(spec.origin).with_name(f"liboeq_stable_{backend}_aoti.so").is_file() - ) - - USE_PRECOMPILED_EXTENSION = True WARNING_MESSAGE = "" @@ -206,7 +195,9 @@ def _has_precompiled_extension(): WARNING_MESSAGE += "HIP does not support precompiled extension yet.\n" USE_PRECOMPILED_EXTENSION = False -if not _has_precompiled_extension(): +if not os.path.exists( + os.path.join(os.path.dirname(__file__), "liboeq_stable_cuda_aoti.so") +): WARNING_MESSAGE += "Precompiled extension shared object not found.\n" USE_PRECOMPILED_EXTENSION = False From e644be8a723ab6882689e49817d42d81e7e78186 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:31:45 -0700 Subject: [PATCH 15/18] simplify pr and move unrelated fixes to separate prs --- openequivariance/CMakeLists.txt | 2 +- openequivariance/openequivariance/_torch/extlib/__init__.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 3e799a82..44caad4a 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -163,7 +163,7 @@ if(hip_FOUND) hip::host hiprtc::hiprtc ) - add_stable_extension(oeq_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") + add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() if(NOT CUDAToolkit_FOUND AND NOT hip_FOUND) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index ed41bfa6..b38080ac 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -125,9 +125,9 @@ def load_jit_extension(): extra_cflags.append("-DCUDA_BACKEND") elif torch.version.hip: - hip_lib_dirs = library_paths("cuda") - extra_link_args.append("-Wl,-rpath," + hip_lib_dirs[0]) - extra_link_args.extend("-L" + directory for directory in hip_lib_dirs) + 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") From 157de6513e847544b4c58fa05d758c4d29b431df Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:53:57 -0700 Subject: [PATCH 16/18] simplify pr. We're not exposing gemm as an external interface, just test it through symmetric tensor product. The test that we don't rely on blas is over testing. --- tests/extension_linkage_test.py | 28 --- tests/group_gemm_test.py | 314 -------------------------------- 2 files changed, 342 deletions(-) delete mode 100644 tests/extension_linkage_test.py delete mode 100644 tests/group_gemm_test.py diff --git a/tests/extension_linkage_test.py b/tests/extension_linkage_test.py deleted file mode 100644 index 77fbb8bc..00000000 --- a/tests/extension_linkage_test.py +++ /dev/null @@ -1,28 +0,0 @@ -import re -import shutil -import subprocess - -import pytest - - -def test_extension_has_no_direct_blas_dependency(): - if not shutil.which("readelf"): - pytest.skip("readelf is required to inspect ELF dependencies") - import openequivariance - from openequivariance._torch.extlib import extension_module - - paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} - for path in paths: - dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) - dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] - assert not any( - re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies - ) - symbols = subprocess.check_output( - ["readelf", "--dyn-syms", "--wide", path], text=True - ) - undefined = [line for line in symbols.splitlines() if " UND " in line] - assert not any( - re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) - for line in undefined - ) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py deleted file mode 100644 index 8df3d5bf..00000000 --- a/tests/group_gemm_test.py +++ /dev/null @@ -1,314 +0,0 @@ -import importlib -import subprocess -import sys - -import pytest -import torch - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="A CUDA or ROCm GPU is required" -) - - -@pytest.fixture(scope="module") -def group_gemm(): - importlib.import_module("openequivariance._torch.symmetric_contraction") - return torch.ops.libtorch_tp_jit.group_gemm - - -def reference(A, B, counts, inner): - pieces = [] - offset = 0 - for i, n in enumerate(counts): - if inner == 0: - pieces.append(torch.einsum("bmk,nbk->nbm", A[i], B[offset : offset + n])) - else: - pieces.append( - torch.einsum( - "nbm,nbk->bmk", A[offset : offset + n], B[offset : offset + n] - ) - ) - offset += n - if inner == 0: - return torch.cat(pieces, dim=0) - return torch.stack(pieces, dim=0) - - -def make_input(shape, dtype, layout="contiguous", device="cuda"): - values = torch.arange(1, 1 + torch.Size(shape).numel(), dtype=dtype, device="cpu") - values = (values.remainder(17) - 8).reshape(shape) / 8 - if layout == "offset": - storage = torch.empty(values.numel() + 7, dtype=dtype, device=device) - result = storage[3 : 3 + values.numel()].view(shape) - result.copy_(values) - return result - if layout == "noncontiguous": - storage = torch.empty((*shape, 2), dtype=dtype, device=device) - result = storage[..., 0] - result.copy_(values) - return result - return values.to(device) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize("layout", ["contiguous", "offset", "noncontiguous"]) -def test_group_gemm_matches_reference(group_gemm, dtype, inner, layout): - counts = [2, 0, 3] - batch, m, k = 2, 3, 4 - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, dtype, layout) - B = make_input((sum(counts), batch, k), dtype, layout) - expected = reference(A.cpu().double(), B.cpu().double(), counts, inner) - counts_tensor = torch.tensor([n for n in counts for _ in range(2)], device="cpu")[ - ::2 - ] - - actual = group_gemm(A, B, counts_tensor, len(counts), batch, m, k, inner) - - torch.testing.assert_close(actual.cpu().double(), expected, rtol=1e-5, atol=1e-5) - assert actual.dtype == dtype - assert actual.device == A.device - - -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize( - "counts,batch,m,k", - [ - ([0, 0], 2, 3, 4), - ([], 2, 3, 4), - ([2, 0, 3], 0, 3, 4), - ([2, 0, 3], 2, 0, 4), - ([2, 0, 3], 2, 3, 0), - ], -) -def test_group_gemm_empty_dimensions(group_gemm, inner, counts, batch, m, k): - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = torch.empty(A_shape, device="cuda", dtype=torch.float64) - B = torch.empty((sum(counts), batch, k), device="cuda", dtype=torch.float64) - actual = group_gemm( - A, - B, - torch.tensor(counts, dtype=torch.int64, device="cpu"), - len(counts), - batch, - m, - k, - inner, - ) - expected_shape = ( - (sum(counts), batch, m) if inner == 0 else (len(counts), batch, m, k) - ) - assert actual.shape == expected_shape - torch.testing.assert_close(actual, torch.zeros_like(actual)) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_backward(group_gemm, inner): - counts = [2, 0, 3] - batch, m, k = 2, 3, 4 - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, torch.float64, "noncontiguous").requires_grad_() - B = make_input((sum(counts), batch, k), torch.float64, "offset").requires_grad_() - A_ref = A.detach().cpu().requires_grad_() - B_ref = B.detach().cpu().requires_grad_() - expected = reference(A_ref, B_ref, counts, inner) - actual = group_gemm( - A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner - ) - grad = make_input(actual.shape, torch.float64, "noncontiguous") - - actual_grads = torch.autograd.grad(actual, (A, B), grad) - expected_grads = torch.autograd.grad(expected, (A_ref, B_ref), grad.cpu()) - - for actual_grad, expected_grad in zip(actual_grads, expected_grads): - torch.testing.assert_close( - actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 - ) - - -def test_group_gemm_current_stream(group_gemm): - counts = torch.tensor([2, 0, 3], device="cpu") - A = torch.zeros((3, 2, 3, 4), device="cuda") - B = torch.zeros((5, 2, 4), device="cuda") - group_gemm(A, B, counts, 3, 2, 3, 4, 0) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - if not torch.version.hip: - torch.cuda._sleep(1_000_000) - A.fill_(2) - B.fill_(3) - actual = group_gemm(A, B, counts, 3, 2, 3, 4, 0).clone() - stream.synchronize() - torch.testing.assert_close(actual, torch.full_like(actual, 24)) - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs are required") -def test_group_gemm_device_guard(group_gemm): - with torch.cuda.device(0): - A = torch.full((1, 2, 3, 4), 2.0, device="cuda:1") - B = torch.full((2, 2, 4), 3.0, device="cuda:1") - actual = group_gemm(A, B, torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0) - assert torch.cuda.current_device() == 0 - assert actual.device == torch.device("cuda:1") - torch.testing.assert_close(actual, torch.full_like(actual, 24)) - with pytest.raises(RuntimeError, match="same device"): - group_gemm( - A, B.to("cuda:0"), torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0 - ) - - -@pytest.mark.parametrize("counts", [[-1, 0, 6], [2, 0, 2], [2, 0, 4]]) -def test_group_gemm_rejects_invalid_counts(group_gemm, counts): - A = torch.empty((3, 2, 3, 4), device="cuda") - B = torch.empty((5, 2, 4), device="cuda") - with pytest.raises(RuntimeError, match="ragged_counts"): - group_gemm(A, B, torch.tensor(counts, device="cpu"), 3, 2, 3, 4, 0) - - -def test_group_gemm_requires_cpu_counts(group_gemm): - A = torch.empty((3, 2, 3, 4), device="cuda") - B = torch.empty((5, 2, 4), device="cuda") - with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): - group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) - - -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize( - "counts,batch,m,k", - [ - ([1, 0, 5], 1, 1, 7), - ([1, 0, 5], 3, 7, 1), - ([1, 0, 5], 3, 1, 1), - ([1, 7, 0, 13], 4, 17, 29), - ], -) -def test_group_gemm_varied_shapes(group_gemm, inner, counts, batch, m, k): - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, torch.float64, "offset") - B = make_input((sum(counts), batch, k), torch.float64, "offset") - actual = group_gemm( - A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner - ) - expected = reference(A.cpu(), B.cpu(), counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_double_backward(group_gemm, inner): - counts = torch.tensor([1, 0, 2], device="cpu") - A_shape = (3, 2, 2, 3) if inner == 0 else (3, 2, 2) - A = make_input(A_shape, torch.float64).requires_grad_() - B = make_input((3, 2, 3), torch.float64).requires_grad_() - - def operation(A, B): - return group_gemm(A, B, counts, 3, 2, 2, 3, inner) - - assert torch.autograd.gradgradcheck(operation, (A, B), fast_mode=True) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_graph_replay(group_gemm, inner): - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64) - B = make_input((5, 2, 4), torch.float64) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) - torch.cuda.current_stream().wait_stream(stream) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - actual = group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) - - for scale in (2, 3): - A.mul_(scale) - B.add_(0.25) - graph.replay() - expected = reference(A.cpu(), B.cpu(), counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_compile(group_gemm, inner): - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64).requires_grad_() - B = make_input((5, 2, 4), torch.float64).requires_grad_() - - def operation(A, B, counts): - return group_gemm(A, B, counts, 3, 2, 3, 4, inner) - - compiled = torch.compile(operation, fullgraph=True) - actual = compiled(A, B, counts_tensor) - A_ref = A.detach().cpu().requires_grad_() - B_ref = B.detach().cpu().requires_grad_() - expected = reference(A_ref, B_ref, counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - actual_grads = torch.autograd.grad(actual.sum(), (A, B)) - expected_grads = torch.autograd.grad(expected.sum(), (A_ref, B_ref)) - for actual_grad, expected_grad in zip(actual_grads, expected_grads): - torch.testing.assert_close( - actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 - ) - - next_counts = [0, 3, 2] - next_actual = compiled(A, B, torch.tensor(next_counts, device="cpu")) - next_expected = reference(A_ref, B_ref, next_counts, inner) - torch.testing.assert_close(next_actual.cpu(), next_expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_aoti(group_gemm, inner, tmp_path): - import openequivariance - - class Model(torch.nn.Module): - def forward(self, A, B, counts): - return group_gemm(A, B, counts, 3, 2, 3, 4, inner) - - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64) - B = make_input((5, 2, 4), torch.float64) - exported = torch.export.export(Model(), (A, B, counts_tensor), strict=False) - package_path = torch._inductor.aoti_compile_and_package( - exported, package_path=str(tmp_path / "group_gemm.pt2") - ) - inputs_path = tmp_path / "inputs.pt" - torch.save( - (A.cpu(), B.cpu(), counts_tensor, reference(A.cpu(), B.cpu(), counts, inner)), - inputs_path, - ) - result = subprocess.run( - [ - sys.executable, - "-c", - """ -import sys -import torch -import torch._inductor.codecache - -torch.ops.load_library(sys.argv[1]) -model = torch._inductor.aoti_load_package(sys.argv[2]) -A, B, counts, expected = torch.load(sys.argv[3], weights_only=True) -actual = model(A.cuda(), B.cuda(), counts) -torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) -""", - openequivariance.torch_ext_so_path(), - package_path, - str(inputs_path), - ], - capture_output=True, - text=True, - timeout=180, - ) - assert result.returncode == 0, result.stdout + result.stderr From 9ab016f6a5b94e5bb2bf7e07109086920af2c63d Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 19:12:02 -0700 Subject: [PATCH 17/18] typo --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 958161f6..ace42549 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## Latest Changes -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim to instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. ### v0.7.0 (2026-09-10) **Added**: From 64af3194056ccf2f1a046c3e1babc125df9120ea Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 19:14:33 -0700 Subject: [PATCH 18/18] fmt --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ace42549..f3eabca6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## Latest Changes -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. ### v0.7.0 (2026-09-10) **Added**: