From 08378a45a3622c05cd8417eb3039b13197f86a27 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 5 Jul 2026 18:18:33 -0700 Subject: [PATCH 01/11] test in place --- tests/stream_ordering_test.py | 177 ++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 tests/stream_ordering_test.py diff --git a/tests/stream_ordering_test.py b/tests/stream_ordering_test.py new file mode 100644 index 0000000..deb2970 --- /dev/null +++ b/tests/stream_ordering_test.py @@ -0,0 +1,177 @@ +# ruff: noqa: E731 +"""Stream-ordering and device-mismatch tests for every GPU-launching API: +the raw group_gemm op, TensorProduct, and TensorProductConv. + +Technique: zero one input, stall a stream with torch.cuda._sleep, enqueue the +real input write behind the stall, then run the op with that stream current. +An op that enqueues on the correct stream waits for the write; an op that uses +the wrong stream (e.g. the legacy default stream, or another device's stream) +always reads stale zeros, so the failure is deterministic rather than a race. + +Expected against the current implementation: + - test_ordering_current_device: FAILS for group_gemm (cuBLAS work is issued + with no stream set, i.e. the legacy default stream); PASSES for tp/conv, + which are stream-correct on a single device. + - test_cuda_graph_capture: RAISES for group_gemm (legacy-stream work aborts + stream capture); PASSES for tp/conv. + - test_ordering_nondefault_device (2+ GPUs): FAILS for all three cases, since + every op derives device and stream from the host thread, not its inputs. + tp/conv may fail as an async CUDA error or process abort rather than a + clean assert, because kernel-launch errors currently call exit(1). +""" + +import pytest +import torch + +from e3nn import o3 + +from openequivariance import TensorProduct, TensorProductConv, TPProblem +from openequivariance._torch import extlib # noqa: F401 loads the extension + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA device required" +) + +# ~50-100ms spin on modern GPUs: much longer than any launch latency here. +SLEEP_CYCLES = int(2e8) + +NUM_W, NUM_FEATURES, M, K, N_PER_W = 4, 8, 16, 32, 256 +N_BATCH = 1000 + + +def _tpp(): + return TPProblem( + o3.Irreps("1x2e"), + o3.Irreps("1x3e"), + o3.Irreps("1x2e"), + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + ) + + +class Case: + """Callable op plus its input tensors; tensors[0] is the delayed input.""" + + def __init__(self, fn, tensors): + self.fn = fn + self.tensors = tensors + + def __call__(self): + return self.fn(*self.tensors) + + +def _build_group_gemm(device, gen): + counts = torch.full((NUM_W,), N_PER_W, dtype=torch.int64) # CPU tensor + A = torch.randn(NUM_W, NUM_FEATURES, M, K, device=device, generator=gen) + B = torch.randn( + NUM_W * N_PER_W, NUM_FEATURES, K, device=device, generator=gen + ) + fn = lambda A, B: torch.ops.libtorch_tp_jit.group_gemm( + A, B, counts, NUM_W, NUM_FEATURES, M, K, 0 + ) + return Case(fn, [A, B]) + + +def _build_tp(device, gen): + tpp = _tpp() + tp = TensorProduct(tpp) + X = torch.rand(N_BATCH, tpp.irreps_in1.dim, device=device, generator=gen) + Y = torch.rand(N_BATCH, tpp.irreps_in2.dim, device=device, generator=gen) + W = torch.rand(N_BATCH, tpp.weight_numel, device=device, generator=gen) + return Case(tp, [X, Y, W]) + + +def _build_conv_atomic(device, gen): + tpp = _tpp() + conv = TensorProductConv(tpp, torch_op=True, deterministic=False) + receivers = torch.tensor([0, 1, 1, 2], device=device, dtype=torch.long) + senders = torch.tensor([1, 0, 2, 1], device=device, dtype=torch.long) + X = torch.rand(3, tpp.irreps_in1.dim, device=device, generator=gen) + Y = torch.rand(4, tpp.irreps_in2.dim, device=device, generator=gen) + W = torch.rand(4, tpp.weight_numel, device=device, generator=gen) + return Case(conv, [X, Y, W, receivers, senders]) + + +BUILDERS = { + "group_gemm": _build_group_gemm, + "tp": _build_tp, + "conv_atomic": _build_conv_atomic, +} + + +@pytest.fixture(params=list(BUILDERS)) +def case_name(request): + return request.param + + +def _build(case_name, device): + gen = torch.Generator(device=device) + gen.manual_seed(0) + return BUILDERS[case_name](device, gen) + + +def _delayed_run(case, tensor_device, stream): + """Re-write the delayed input on a stalled stream, then run the op with + that stream current on the tensors' device. The op must wait for the + write; reading early yields zeros.""" + delayed = case.tensors[0] + src = delayed.clone() + torch.cuda.synchronize(tensor_device) + delayed.zero_() + torch.cuda.synchronize(tensor_device) + + with torch.cuda.device(tensor_device), torch.cuda.stream(stream): + torch.cuda._sleep(SLEEP_CYCLES) + delayed.copy_(src) + + # Note: torch.cuda.stream() selects the stream on *its* device without + # changing the current device, so in the cross-device test the host + # thread's current device stays cuda:0 here. + with torch.cuda.stream(stream): + out = case() + torch.cuda.synchronize(tensor_device) + return out + + +def test_ordering_current_device(case_name): + case = _build(case_name, "cuda") + torch.cuda.synchronize() + ref = case() + torch.cuda.synchronize() + + out = _delayed_run(case, 0, torch.cuda.Stream()) + torch.testing.assert_close(out, ref) + + +def test_cuda_graph_capture(case_name): + case = _build(case_name, "cuda") + warmup = torch.cuda.Stream() + with torch.cuda.stream(warmup): + for _ in range(3): + case() + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = case() + g.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(out, case()) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires 2+ GPUs") +def test_ordering_nondefault_device(case_name): + """All inputs on cuda:1 while the host thread stays on cuda:0. Ops must + derive device and stream from their inputs, not from thread state.""" + assert torch.cuda.current_device() == 0 + + case = _build(case_name, "cuda:1") + torch.cuda.synchronize(1) + ref = case() # already the report's mismatch scenario + torch.cuda.synchronize(1) + assert ref.device == case.tensors[0].device + + out = _delayed_run(case, 1, torch.cuda.Stream(device=1)) + torch.testing.assert_close(out, ref) From fb534011cf63c0f462fb51d57c9034c5c2964742 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 5 Jul 2026 18:18:55 -0700 Subject: [PATCH 02/11] check that inputs are on the same device --- .../openequivariance/extension/torch_core.hpp | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/openequivariance/openequivariance/extension/torch_core.hpp b/openequivariance/openequivariance/extension/torch_core.hpp index ab78d96..cb7d96f 100644 --- a/openequivariance/openequivariance/extension/torch_core.hpp +++ b/openequivariance/openequivariance/extension/torch_core.hpp @@ -128,6 +128,21 @@ inline void check_tensor(const Tensor &tensor, ". Got: ", static_cast(tensor.scalar_type())); } +inline void check_same_device( + std::initializer_list> tensors) { + auto it = tensors.begin(); + const Tensor *first = it->first; + const char *first_name = it->second; + for (++it; it != tensors.end(); ++it) { + TCHECK(it->first->get_device() == first->get_device(), + "Tensor '", it->second, "' is on device ", + static_cast(it->first->get_device()), + ", but tensor '", first_name, "' is on device ", + static_cast(first->get_device()), + ". All tensors must be on the same device."); + } +} + inline std::unordered_map parse_json_config(const json &j_obj) { std::unordered_map result; for (const auto &kv : j_obj.object_items()) { @@ -286,6 +301,8 @@ inline Tensor jit_tp_forward( else check_tensor(W, {num_batch, k.weight_numel}, k.weight_dtype, "W"); + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}}); + Tensor L3_out = tensor_empty_like(L1_in, make_sizes({num_batch, k.L3_dim})); Tensor L1_contig = tensor_contiguous(L1_in); @@ -325,6 +342,9 @@ inline tuple jit_tp_backward( else check_tensor(W, {num_batch, k.weight_numel}, k.weight_dtype, "W"); + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}, + {&L3_grad, "L3_grad"}}); + Tensor L1_grad = tensor_empty_like(L1_in, tensor_sizes_vec(L1_in)); Tensor L2_grad = tensor_empty_like(L2_in, tensor_sizes_vec(L2_in)); Tensor W_grad = tensor_empty_like(W, tensor_sizes_vec(W)); @@ -378,6 +398,10 @@ inline tuple jit_tp_double_backward( check_tensor(W_dgrad, {num_batch, k.weight_numel}, k.weight_dtype, "W_dgrad"); } + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}, + {&L3_grad, "L3_grad"}, {&L1_dgrad, "L1_dgrad"}, + {&L2_dgrad, "L2_dgrad"}, {&W_dgrad, "W_dgrad"}}); + Tensor L1_grad = tensor_empty_like(L1_in, tensor_sizes_vec(L1_in)); Tensor L2_grad = tensor_empty_like(L2_in, tensor_sizes_vec(L2_in)); Tensor W_grad = tensor_empty_like(W, tensor_sizes_vec(W)); @@ -447,6 +471,11 @@ inline Tensor jit_conv_forward( else check_tensor(W, {nnz, k.weight_numel}, k.weight_dtype, "W"); + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}, + {&rows, "rows"}, {&cols, "cols"}, {&workspace, "workspace"}}); + if (k.deterministic) + check_same_device({{&L1_in, "L1_in"}, {&transpose_perm, "transpose_perm"}}); + Tensor L3_out = tensor_zeros_like(L1_in, make_sizes({node_count, k.L3_dim})); Tensor L1_contig = tensor_contiguous(L1_in); @@ -505,6 +534,12 @@ inline tuple jit_conv_backward( else check_tensor(W, {nnz, k.weight_numel}, k.weight_dtype, "W"); + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}, + {&L3_grad, "L3_grad"}, {&rows, "rows"}, {&cols, "cols"}, + {&workspace, "workspace"}}); + if (k.deterministic) + check_same_device({{&L1_in, "L1_in"}, {&transpose_perm, "transpose_perm"}}); + Tensor L1_grad = tensor_zeros_like(L1_in, tensor_sizes_vec(L1_in)); Tensor L2_grad = tensor_zeros_like(L2_in, tensor_sizes_vec(L2_in)); Tensor W_grad = tensor_empty_like(W, tensor_sizes_vec(W)); @@ -579,6 +614,13 @@ inline tuple jit_conv_double_backward( check_tensor(W_dgrad, {nnz, k.weight_numel}, k.weight_dtype, "W_dgrad"); } + check_same_device({{&L1_in, "L1_in"}, {&L2_in, "L2_in"}, {&W, "W"}, + {&L3_grad, "L3_grad"}, {&L1_dgrad, "L1_dgrad"}, + {&L2_dgrad, "L2_dgrad"}, {&W_dgrad, "W_dgrad"}, + {&rows, "rows"}, {&cols, "cols"}, {&workspace, "workspace"}}); + if (k.deterministic) + check_same_device({{&L1_in, "L1_in"}, {&transpose_perm, "transpose_perm"}}); + Tensor L1_grad = tensor_zeros_like(L1_in, tensor_sizes_vec(L1_in)); Tensor L2_grad = tensor_zeros_like(L2_in, tensor_sizes_vec(L2_in)); Tensor W_grad = tensor_empty_like(W, tensor_sizes_vec(W)); @@ -623,6 +665,11 @@ inline Tensor group_gemm( int64_t num_W, int64_t batch_size, int64_t m, int64_t k, int64_t ragged_inner) { TCHECK(A.scalar_type() == B.scalar_type(), "group_gemm: A and B must have the same dtype"); TCHECK(ragged_counts.scalar_type() == kLong, "group_gemm: ragged_counts must be int64"); + TCHECK(A.is_cuda(), "group_gemm: A must be a GPU tensor"); + TCHECK(B.is_cuda(), "group_gemm: B must be a GPU tensor"); + check_same_device({{&A, "A"}, {&B, "B"}}); + TCHECK(ragged_counts.is_cpu(), + "group_gemm: ragged_counts must be a CPU tensor; its values are read on the host"); Tensor A_c = tensor_contiguous(A); Tensor B_c = tensor_contiguous(B); From 5eda2990b432fdadf54a6eb0c20e9f79d8a373d7 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 11 Jul 2026 23:57:43 -0700 Subject: [PATCH 03/11] more stubs --- .../extension/stubs/stream.cpp | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/openequivariance/openequivariance/extension/stubs/stream.cpp b/openequivariance/openequivariance/extension/stubs/stream.cpp index fd011c3..c3c0e34 100644 --- a/openequivariance/openequivariance/extension/stubs/stream.cpp +++ b/openequivariance/openequivariance/extension/stubs/stream.cpp @@ -1,8 +1,27 @@ #include #include +// Link-time stand-ins for symbols that live in the real libtorch_cuda / +// libtorch_hip, which is already loaded at runtime and wins symbol +// resolution. These bodies never execute in production. extern "C" { AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) { return 0; } + + AOTITorchError aoti_torch_create_device_guard(int32_t device_index, DeviceGuardHandle* ret_guard) { + return 0; + } + + AOTITorchError aoti_torch_delete_device_guard(DeviceGuardHandle guard) { + return 0; + } + + AOTITorchError aoti_torch_device_guard_set_index(DeviceGuardHandle guard, int32_t device_index) { + return 0; + } + + AOTITorchError torch_get_current_cuda_blas_handle(void** ret_handle) { + return 0; + } } \ No newline at end of file From 2b3d458c72f2d582b18b27ba0ce15f0b955d0581 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 11 Jul 2026 23:58:08 -0700 Subject: [PATCH 04/11] get device specific streams --- .../openequivariance/extension/torch_core.hpp | 46 +++++++++++++++---- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/openequivariance/openequivariance/extension/torch_core.hpp b/openequivariance/openequivariance/extension/torch_core.hpp index cb7d96f..97411e9 100644 --- a/openequivariance/openequivariance/extension/torch_core.hpp +++ b/openequivariance/openequivariance/extension/torch_core.hpp @@ -42,7 +42,11 @@ Tensor tensor_zeros_like(const Tensor &ref, const std::vector &sizes); void tensor_zero_(Tensor &tensor); void alert_not_deterministic(const char *name); -Stream get_current_stream(); + +// Current PyTorch stream for the given device, NOT the thread's current +// device. Each op must guard to its input tensor's device (DeviceGuard, +// aliased per-backend) before launching work, then take that device's stream. +Stream get_current_stream(int32_t device_index); const uint8_t *tensor_data_ptr_u8(const Tensor &tensor); void *data_ptr(const Tensor &tensor); @@ -289,7 +293,10 @@ inline Tensor jit_tp_forward( int64_t L3_dim) { auto [jit_kernel, k] = compile_tp_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t num_batch = L1_in.size(0); @@ -329,7 +336,10 @@ inline tuple jit_tp_backward( Tensor L3_grad) { auto [jit_kernel, k] = compile_tp_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t num_batch = L1_in.size(0); @@ -380,7 +390,10 @@ inline tuple jit_tp_double_backward( Tensor W_dgrad) { auto [jit_kernel, k] = compile_tp_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t num_batch = L1_in.size(0); @@ -450,7 +463,10 @@ inline Tensor jit_conv_forward( Tensor transpose_perm) { auto [jit_kernel, k] = compile_conv_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t nnz = rows.size(0); const int64_t node_count = L1_in.size(0); @@ -511,7 +527,10 @@ inline tuple jit_conv_backward( Tensor transpose_perm) { auto [jit_kernel, k] = compile_conv_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t nnz = rows.size(0); const int64_t node_count = L1_in.size(0); @@ -586,7 +605,10 @@ inline tuple jit_conv_double_backward( Tensor transpose_perm) { auto [jit_kernel, k] = compile_conv_with_caching(json_bytes, hash); - Stream stream = get_current_stream(); + + const int32_t device_index = static_cast(L1_in.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); const int64_t nnz = rows.size(0); const int64_t node_count = L1_in.size(0); @@ -671,6 +693,10 @@ inline Tensor group_gemm( TCHECK(ragged_counts.is_cpu(), "group_gemm: ragged_counts must be a CPU tensor; its values are read on the host"); + const int32_t device_index = static_cast(A.get_device()); + DeviceGuard device_guard(device_index); + Stream stream = get_current_stream(device_index); + Tensor A_c = tensor_contiguous(A); Tensor B_c = tensor_contiguous(B); Tensor rc_c = tensor_contiguous(ragged_counts); @@ -686,10 +712,12 @@ inline Tensor group_gemm( 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); + (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner, + device_index, stream); } 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); + (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner, + device_index, stream); } else { throw std::logic_error("group_gemm: unsupported dtype, expected float32 or float64"); } From fca71ac2579764216615b1da6e1cc7fb9b16b2e5 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 11 Jul 2026 23:58:35 -0700 Subject: [PATCH 05/11] back to hipblas - becuase it's cublas compatible --- openequivariance/CMakeLists.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7bfbd49..acf86da 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -117,7 +117,7 @@ endfunction() find_package(CUDAToolkit QUIET) find_package(hip QUIET) -find_package(rocblas QUIET) +find_package(hipblas QUIET) if(CUDAToolkit_FOUND) message(STATUS "Building stable extension with CUDA backend.") @@ -159,10 +159,10 @@ if(hip_FOUND) CXX_STANDARD 17 ) - if(TARGET roc::rocblas) - set(HIP_BLAS_LIB roc::rocblas) + if(TARGET roc::hipblas) + set(HIP_BLAS_LIB roc::hipblas) else() - set(HIP_BLAS_LIB rocblas) + set(HIP_BLAS_LIB hipblas) endif() set(HIP_LINK_LIBS From 8b708e4599b8b9379c4dfaa95c4ead79b0810b40 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:06:39 -0700 Subject: [PATCH 06/11] use pytorch owned blas handles --- .../openequivariance/extension/group_mm.hpp | 55 ++++++++----------- 1 file changed, 23 insertions(+), 32 deletions(-) diff --git a/openequivariance/openequivariance/extension/group_mm.hpp b/openequivariance/openequivariance/extension/group_mm.hpp index 19249c3..d7372d9 100644 --- a/openequivariance/openequivariance/extension/group_mm.hpp +++ b/openequivariance/openequivariance/extension/group_mm.hpp @@ -8,38 +8,29 @@ #include "cublas_v2.h" #include - struct BlasHandle { - cublasHandle_t handle; - BlasHandle() { - if (cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS) - throw std::logic_error("CUBLAS initialization failed"); - } - ~BlasHandle() { cublasDestroy(handle); } - }; + using BlasStream = cudaStream_t; + using BlasHandleT = cublasHandle_t; #elif defined(HIP_BACKEND) - #include "rocblas/rocblas.h" + #include #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); } - }; + using BlasStream = hipStream_t; + using BlasHandleT = hipblasHandle_t; #endif -inline BlasHandle& get_blas_handle() { - static BlasHandle handle; - return handle; -} +// Returns a BLAS handle ready to launch on (device_index, stream). Defined +// per-build (libtorch_tp_jit.cpp / libtorch_tp_jit_stable.cpp): both return +// PyTorch's pooled handle (cuBLAS on CUDA, hipBLAS on ROCm), which arrives +// with the stream, workspace, and math mode already set for the current +// device and stream. The caller must hold a device guard for device_index. +BlasHandleT get_op_blas_handle(int32_t device_index, BlasStream stream); 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) { + int64_t* ragged_counts, int num_W, int batch_size, int m, int k, int ragged_inner, + int32_t device_index, BlasStream stream) { - auto& blas = get_blas_handle(); + BlasHandleT handle = get_op_blas_handle(device_index, stream); T alpha = 1.0, beta = 0.0; T* A_base = reinterpret_cast(A_raw); T* B_base = reinterpret_cast(B_raw); @@ -52,7 +43,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, #ifdef CUDA_BACKEND cublasOperation_t transa, transb; #elif defined(HIP_BACKEND) - rocblas_operation transa, transb; + hipblasOperation_t transa, transb; #endif if (ragged_inner == 0) { @@ -66,7 +57,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, #ifdef CUDA_BACKEND transa = CUBLAS_OP_T; transb = CUBLAS_OP_N; #elif defined(HIP_BACKEND) - transa = rocblas_operation_transpose; transb = rocblas_operation_none; + transa = HIPBLAS_OP_T; transb = HIPBLAS_OP_N; #endif } else { M = k; K = static_cast(ragged_counts[i]); N = m; @@ -79,7 +70,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, #ifdef CUDA_BACKEND transa = CUBLAS_OP_N; transb = CUBLAS_OP_T; #elif defined(HIP_BACKEND) - transa = rocblas_operation_none; transb = rocblas_operation_transpose; + transa = HIPBLAS_OP_N; transb = HIPBLAS_OP_T; #endif } ragged_offset += ragged_counts[i]; @@ -88,7 +79,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, #ifdef CUDA_BACKEND cublasStatus_t stat; if (std::is_same::value) { - stat = cublasSgemmStridedBatched(blas.handle, + stat = cublasSgemmStridedBatched(handle, transa, transb, M, N, K, reinterpret_cast(&alpha), reinterpret_cast(A), lda, strideA, @@ -97,7 +88,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, reinterpret_cast(C), ldc, strideC, batch_size); } else if (std::is_same::value) { - stat = cublasDgemmStridedBatched(blas.handle, + stat = cublasDgemmStridedBatched(handle, transa, transb, M, N, K, reinterpret_cast(&alpha), reinterpret_cast(A), lda, strideA, @@ -111,9 +102,9 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, if (stat != CUBLAS_STATUS_SUCCESS) throw std::logic_error("Grouped GEMM failed!"); #elif defined(HIP_BACKEND) - rocblas_status stat; + hipblasStatus_t stat; if (std::is_same::value) { - stat = rocblas_sgemm_strided_batched(blas.handle, + stat = hipblasSgemmStridedBatched(handle, transa, transb, M, N, K, reinterpret_cast(&alpha), reinterpret_cast(A), lda, strideA, @@ -122,7 +113,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, reinterpret_cast(C), ldc, strideC, batch_size); } else if (std::is_same::value) { - stat = rocblas_dgemm_strided_batched(blas.handle, + stat = hipblasDgemmStridedBatched(handle, transa, transb, M, N, K, reinterpret_cast(&alpha), reinterpret_cast(A), lda, strideA, @@ -133,7 +124,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw, } else { throw std::logic_error("Unsupported datatype for grouped GEMM!"); } - if (stat != rocblas_status_success) + if (stat != HIPBLAS_STATUS_SUCCESS) throw std::logic_error("Grouped GEMM failed!"); #endif } From f7022a54cae67bbcf71053053625c3194def89c2 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:07:38 -0700 Subject: [PATCH 07/11] pytorch blas handles --- .../extension/libtorch_tp_jit_stable.cpp | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 6bf3d51..e9d8e7d 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -11,6 +11,7 @@ #include #include #include +#include using Tensor = torch::stable::Tensor; @@ -27,6 +28,8 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte; #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY +using DeviceGuard = torch::stable::accelerator::DeviceGuard; + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { @@ -66,16 +69,23 @@ void *data_ptr(const Tensor &tensor) { return tensor.data_ptr(); } -Stream get_current_stream() { - auto device_idx = torch::stable::accelerator::getCurrentDeviceIndex(); +Stream get_current_stream(int32_t device_index) { void* stream_ptr = nullptr; - TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_idx, &stream_ptr)); + TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); + return static_cast(stream_ptr); +} - #ifdef CUDA_BACKEND - return static_cast(stream_ptr); - #elif defined(HIP_BACKEND) - return static_cast(stream_ptr); - #endif +BlasHandleT get_op_blas_handle(int32_t device_index, BlasStream stream) { + // The caller's device guard makes device_index current, and `stream` is + // that device's current stream, so PyTorch's handle arrives configured + // for exactly this (device, stream) with its workspace and math mode + // managed by PyTorch. On ROCm builds of libtorch this shim symbol keeps + // its name and returns a hipblasHandle_t. + (void)device_index; + (void)stream; + void* handle = nullptr; + TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&handle)); + return static_cast(handle); } #ifdef CUDA_BACKEND From f465d4ada205018872923e28885fc6625792f073 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:07:53 -0700 Subject: [PATCH 08/11] pytorch blas handles --- .../extension/libtorch_tp_jit.cpp | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index ddabd0b..51de6a3 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -3,10 +3,13 @@ #ifdef CUDA_BACKEND #include + #include #endif #ifdef HIP_BACKEND + #include #include + #include #endif #include @@ -29,6 +32,13 @@ constexpr Dtype kByte = torch::kByte; #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY TORCH_LIBRARY +#ifdef CUDA_BACKEND + using DeviceGuard = c10::cuda::CUDAGuard; +#endif +#ifdef HIP_BACKEND + using DeviceGuard = c10::hip::HIPGuard; +#endif + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { @@ -74,15 +84,26 @@ void *data_ptr(const Tensor &tensor) { throw std::logic_error("Unsupported tensor datatype!"); } -Stream get_current_stream() { +Stream get_current_stream(int32_t device_index) { #ifdef CUDA_BACKEND - return c10::cuda::getCurrentCUDAStream(); + return c10::cuda::getCurrentCUDAStream(device_index); #endif #ifdef HIP_BACKEND - return c10::hip::getCurrentHIPStream(); + return c10::hip::getCurrentHIPStream(device_index); #endif } +BlasHandleT get_op_blas_handle(int32_t device_index, BlasStream stream) { + // The caller's device guard makes device_index current, and `stream` is + // that device's current stream, so PyTorch's handle arrives configured + // for exactly this (device, stream) with its workspace and math mode + // managed by PyTorch. On ROCm the same-named function returns a + // hipblasHandle_t. + (void)device_index; + (void)stream; + return at::cuda::getCurrentCUDABlasHandle(); +} + namespace py=pybind11; PYBIND11_MODULE(libtorch_tp_jit, m) { py::class_(m, "DeviceProp") From 7ff963f871cc459f6e19058e0233a4a4b990c300 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:08:14 -0700 Subject: [PATCH 09/11] tests for device specific stream --- tests/stream_ordering_test.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/stream_ordering_test.py b/tests/stream_ordering_test.py index deb2970..484af9a 100644 --- a/tests/stream_ordering_test.py +++ b/tests/stream_ordering_test.py @@ -146,14 +146,17 @@ def test_ordering_current_device(case_name): def test_cuda_graph_capture(case_name): case = _build(case_name, "cuda") - warmup = torch.cuda.Stream() - with torch.cuda.stream(warmup): + # Warm up on the same stream we capture on: JIT compilation and cuBLAS + # workspace allocation for this (handle, stream) pair happen here, outside + # the capture region. + s = torch.cuda.Stream() + with torch.cuda.stream(s): for _ in range(3): case() torch.cuda.synchronize() g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): + with torch.cuda.graph(g, stream=s): out = case() g.replay() torch.cuda.synchronize() From d511ce16b0a58a349414732d9d56b283b6812b40 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:14:58 -0700 Subject: [PATCH 10/11] loudly error if stubs attempt to be used --- .../extension/stubs/stream.cpp | 42 +++++++++++++++---- 1 file changed, 33 insertions(+), 9 deletions(-) diff --git a/openequivariance/openequivariance/extension/stubs/stream.cpp b/openequivariance/openequivariance/extension/stubs/stream.cpp index c3c0e34..078f5db 100644 --- a/openequivariance/openequivariance/extension/stubs/stream.cpp +++ b/openequivariance/openequivariance/extension/stubs/stream.cpp @@ -1,27 +1,51 @@ #include +#include #include -// Link-time stand-ins for symbols that live in the real libtorch_cuda / -// libtorch_hip, which is already loaded at runtime and wins symbol -// resolution. These bodies never execute in production. +// The stable extension is linked on a machine with CPU-only libtorch, where +// no libtorch_cuda/libtorch_hip exists. This file builds a stand-in library +// with the real one's SONAME so the link succeeds; at runtime the DT_NEEDED +// entry is satisfied by the real library, already loaded via `import torch`, +// and this file is never opened. It is not shipped in the wheel. +// +// These bodies can therefore only execute in a misconfigured process (e.g. +// the extension loaded without torch). Fail through the shim's error-code +// contract - TORCH_ERROR_CODE_CHECK at the call site turns this into a +// catchable exception - rather than returning success with garbage outputs. +namespace { + +AOTITorchError oeq_stub_called(const char *name) { + fprintf(stderr, + "OpenEquivariance: link-time stub '%s' executed; the real " + "libtorch_cuda/libtorch_hip is not loaded (import torch before " + "loading the OpenEquivariance extension).\n", + name); + return AOTI_TORCH_FAILURE; +} + +} // namespace + extern "C" { AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) { - return 0; + if (ret_stream) *ret_stream = nullptr; + return oeq_stub_called("aoti_torch_get_current_cuda_stream"); } AOTITorchError aoti_torch_create_device_guard(int32_t device_index, DeviceGuardHandle* ret_guard) { - return 0; + if (ret_guard) *ret_guard = nullptr; + return oeq_stub_called("aoti_torch_create_device_guard"); } AOTITorchError aoti_torch_delete_device_guard(DeviceGuardHandle guard) { - return 0; + return oeq_stub_called("aoti_torch_delete_device_guard"); } AOTITorchError aoti_torch_device_guard_set_index(DeviceGuardHandle guard, int32_t device_index) { - return 0; + return oeq_stub_called("aoti_torch_device_guard_set_index"); } AOTITorchError torch_get_current_cuda_blas_handle(void** ret_handle) { - return 0; + if (ret_handle) *ret_handle = nullptr; + return oeq_stub_called("torch_get_current_cuda_blas_handle"); } -} \ No newline at end of file +} From ca8c0834aa11e62be5cfbcdcfa1d876087b743ff Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 12 Jul 2026 20:45:59 -0700 Subject: [PATCH 11/11] ruff fmt --- tests/stream_ordering_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/stream_ordering_test.py b/tests/stream_ordering_test.py index 484af9a..ae1d444 100644 --- a/tests/stream_ordering_test.py +++ b/tests/stream_ordering_test.py @@ -64,9 +64,7 @@ def __call__(self): def _build_group_gemm(device, gen): counts = torch.full((NUM_W,), N_PER_W, dtype=torch.int64) # CPU tensor A = torch.randn(NUM_W, NUM_FEATURES, M, K, device=device, generator=gen) - B = torch.randn( - NUM_W * N_PER_W, NUM_FEATURES, K, device=device, generator=gen - ) + B = torch.randn(NUM_W * N_PER_W, NUM_FEATURES, K, device=device, generator=gen) fn = lambda A, B: torch.ops.libtorch_tp_jit.group_gemm( A, B, counts, NUM_W, NUM_FEATURES, M, K, 0 )