cuequivariance-ops-cu12 0.4.0__py3-none-manylinux_2_39_aarch64.whl → 0.5.0__py3-none-manylinux_2_39_aarch64.whl

cuequivariance_ops/VERSION

@@ -1 +1 @@
-0.4.0
+0.5.0

cuequivariance_ops/__init__.py

@@ -10,6 +10,7 @@
 
 from ._version import __version__, __git_commit__
 import os
+import sys
 import ctypes
 
 PREFERRED_LOAD_FLAG = ctypes.RTLD_LOCAL
@@ -32,8 +33,8 @@ def load_library():
         ctypes.CDLL(
             os.path.join(root_dir(), "lib/libcue_ops.so"), mode=PREFERRED_LOAD_FLAG
         )
-    except Exception:
-        pass
+    except Exception as e:
+        print(f"Error while loading libcue_ops.so: {e}", file=sys.stderr)
 
 
 load_library()

cuequivariance_ops/cache_manager.py

@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+import hashlib
+import json
+import math
+import os
+from typing import Any
+
+import pynvml
+import torch
+
+
+def get_gpu_information():
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available")
+    current_device = torch.cuda.current_device()
+    props = torch.cuda.get_device_properties(current_device)
+
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(current_device)
+    power_limit = pynvml.nvmlDeviceGetPowerManagementLimit(handle)
+    max_clock_rate = pynvml.nvmlDeviceGetMaxClockInfo(
+        handle, pynvml.NVML_CLOCK_GRAPHICS
+    )
+    pynvml.nvmlShutdown()
+
+    return {
+        "name": props.name,
+        "total_memory": math.ceil(props.total_memory / (1024**3)),
+        "multi_processor_count": props.multi_processor_count,
+        "power_limit": power_limit // 1000,
+        "clock_rate": max_clock_rate,
+        "major": props.major,
+        "minor": props.minor,
+    }
+
+
+def gpu_information_to_key(information: dict) -> str:
+    information.pop("name", None)
+    key_string = "_".join(f"{value}" for value in information.values()).replace(
+        " ", "_"
+    )
+    hash_object = hashlib.sha256(key_string.encode())
+    hast_str = hash_object.hexdigest()
+    return hast_str
+
+
+class CacheManager:
+    """Singleton managing the cache"""
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.cache = {}
+            cls._instance.gpu_information = get_gpu_information()
+            cls._instance.gpu_key = gpu_information_to_key(
+                cls._instance.gpu_information
+            )
+        return cls._instance
+
+    def load_cache(
+        self, fn_key: str, json_path: str = os.path.dirname(__file__)
+    ) -> dict:
+        # load the json file and store it in the cache-dict
+        json_file = os.path.join(json_path, f"{fn_key}.json")
+        if os.path.exists(json_file):
+            with open(json_file, "r") as f:
+                fn_cache = json.load(f)
+                self.cache[fn_key] = fn_cache
+        else:
+            # if the file does not exist, create an empty dict for the specified function
+            fn_cache = self.cache[fn_key] = {}
+
+        return fn_cache
+
+    def save_cache(
+        self, fn_key: str, json_path: str = os.path.dirname(__file__)
+    ) -> None:
+        # save cache-dict to json file
+        json_file = os.path.join(json_path, f"{fn_key}.json")
+        # Load existing data from the file if it exists
+        if os.path.exists(json_file):
+            with open(json_file, "r") as f:
+                existing_data = json.load(f)
+        else:
+            existing_data = {}
+
+        # Merge the existing data with the current cache
+        merged_data = {**existing_data, **self.cache[fn_key]}
+
+        # Save the merged data back to the file
+        with open(json_file, "w") as f:
+            json.dump(merged_data, f, indent=4)
+
+        # Update the cache with the merged data
+        self.cache[fn_key] = merged_data
+
+    def get(
+        self, fn_key: str, inp_key: str, json_path: str = os.path.dirname(__file__)
+    ) -> Any:
+        # get value from cache
+        # if necessary, load json first
+        fn_cache = (
+            self.load_cache(fn_key) if fn_key not in self.cache else self.cache[fn_key]
+        )
+        # check if fn_key and inp_key exist in cache
+        if self.gpu_key in fn_cache and inp_key in fn_cache[self.gpu_key]:
+            return self.cache[fn_key][self.gpu_key][inp_key]
+        else:
+            # if not found, return None or raise an error
+            return None
+
+    def set(self, fn_key: str, inp_key: str, value: Any) -> None:
+        # write value to cache-dict
+        # and write-back to json
+        if fn_key not in self.cache:
+            self.cache[fn_key] = {}
+        if self.gpu_key not in self.cache[fn_key]:
+            self.cache[fn_key][self.gpu_key] = {"gpu_information": self.gpu_information}
+        self.cache[fn_key][self.gpu_key][inp_key] = value

cuequivariance_ops/equivariance/dtypes.hh

@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+namespace kernelcatcher::utils {
+
+enum class Datatype : int {
+  kFloat32  = 0,
+  kFloat64  = 1,
+  kFloat16  = 2,
+  kBFloat16 = 3,
+  kInt32    = 4,
+  kInt64    = 5
+};
+
+}  // namespace kernelcatcher::utils

cuequivariance_ops/equivariance/indexed_linear.hh

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+ *
+ * This source code and/or documentation ("Licensed Deliverables") are
+ * subject to NVIDIA intellectual property rights under U.S. and
+ * international Copyright laws.
+ */
+
+#pragma once
+
+#include "dtypes.hh"  // for Datatype
+
+namespace kernelcatcher::equivariance::indexed_linear {
+using namespace kernelcatcher::utils;  // for Datatype
+
+#define KC_INDEXED_LINEAR_DECL_ARGUMENTS                                                  \
+  const void *ptr_A, const void *ptr_B, const int *counts, void *ptr_C, Datatype dtype_A, \
+    Datatype dtype_B, Datatype dtype_D, int Z, int C, int u, int v, double coefficient,   \
+    Datatype math_dtype, void *workspace, size_t workspace_size, void *stream
+
+#define KC_INDEXED_LINEAR_ARGUMENTS                                                            \
+  ptr_A, ptr_B, counts, ptr_C, dtype_A, dtype_B, dtype_D, Z, C, u, v, coefficient, math_dtype, \
+    workspace, workspace_size, stream
+
+int run_indexed_linear_B(  // ptr_A  Zu
+                           // ptr_B  Cuv or Cvu
+                           // ptr_C  Zv
+  bool transpose_B,
+  KC_INDEXED_LINEAR_DECL_ARGUMENTS);
+
+int run_indexed_linear_C(  // ptr_A  Zu
+                           // ptr_B  Zv
+                           // ptr_C  Cuv
+  KC_INDEXED_LINEAR_DECL_ARGUMENTS);
+
+}  // namespace kernelcatcher::equivariance::indexed_linear

cuequivariance_ops/equivariance/run_fmha.h

@@ -0,0 +1,192 @@
+#ifndef CUDNN_FMHA_RUN_FMHA_H
+#define CUDNN_FMHA_RUN_FMHA_H
+
+#include <cstdint>  // for uint32_t
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+
+namespace cudnn_fmha {
+using DataType_TriBias = float;
+
+/**
+ * @brief Performs Flash Multi-Head Attention computation on GPU using cuDNN
+ * @tparam DType Data type for the computation (float, __half, or __nv_bfloat16)
+ */
+template <typename DType>
+void run_fmha(DType* q_ptr,
+              DType* k_ptr,
+              DType* v_ptr,
+              DType* o_ptr,
+              bool* mask_bias_ptr,
+              DataType_TriBias* triangle_bias_ptr,
+              float* softmax_lse_ptr,
+              float* softmax_max_ptr,
+              const uint32_t B,
+              const uint32_t I,
+              const uint32_t H,
+              const uint32_t S_qo,
+              const uint32_t S_kv,
+              const uint32_t D,
+              const float bmm_scale,
+              bool use_tf32,
+              void* stream = nullptr);
+
+/**
+ * @brief Performs the backward pass of Flash Multi-Head Attention computation on GPU using cuDNN
+ * Note: Backward pass remains in float before fp16/bf16 integration
+ */
+template <typename DType>
+void run_fmha_bwd(DType* do_ptr,              // [B, N, H, S_qo, D]
+                  DType* o_ptr,               // [B, N, H, S_qo, D]
+                  float* softmax_lse_ptr,     // [B, N, H, S_qo, 1]
+                  DType* q_ptr,               // [B, N, H, S_qo, D]
+                  DType* k_ptr,               // [B, N, H, S_kv, D]
+                  DType* v_ptr,               // [B, N, H, S_kv, D]
+                  bool* mask_bias_ptr,        // [B, N, 1, 1, S_kv]
+                  float* triangle_bias_ptr,   // [B, 1, H, S_qo, S_kv]
+                  DType* dq_ptr,              // [B, N, H, S_qo, D] output
+                  DType* dk_ptr,              // [B, N, H, S_kv, D] output
+                  DType* dv_ptr,              // [B, N, H, S_kv, D] output
+                  float* triangle_dbias_ptr,  // [B, 1, H, S_qo, S_kv] output
+                  float* do_o_dot_ptr,
+                  float* dq_fp32_buf,  // [B, N, H, S_qo, D] worspace
+                  const uint32_t B,
+                  const uint32_t I,
+                  const uint32_t H,
+                  const uint32_t S_qo,
+                  const uint32_t S_kv,
+                  const uint32_t D,
+                  const float bmm_scale,
+                  bool use_tf32,
+                  void* stream = nullptr);
+
+// Explicit template declarations for supported types
+extern template void run_fmha<float>(float*,
+                                     float*,
+                                     float*,
+                                     float*,
+                                     bool*,
+                                     DataType_TriBias*,
+                                     float*,
+                                     float*,
+                                     uint32_t,
+                                     uint32_t,
+                                     uint32_t,
+                                     uint32_t,
+                                     uint32_t,
+                                     uint32_t,
+                                     float,
+                                     bool,
+                                     void*);
+
+extern template void run_fmha<__half>(__half*,
+                                      __half*,
+                                      __half*,
+                                      __half*,
+                                      bool*,
+                                      DataType_TriBias*,
+                                      float*,
+                                      float*,
+                                      uint32_t,
+                                      uint32_t,
+                                      uint32_t,
+                                      uint32_t,
+                                      uint32_t,
+                                      uint32_t,
+                                      float,
+                                      bool,
+                                      void*);
+
+extern template void run_fmha<__nv_bfloat16>(__nv_bfloat16*,
+                                             __nv_bfloat16*,
+                                             __nv_bfloat16*,
+                                             __nv_bfloat16*,
+                                             bool*,
+                                             DataType_TriBias*,
+                                             float*,
+                                             float*,
+                                             uint32_t,
+                                             uint32_t,
+                                             uint32_t,
+                                             uint32_t,
+                                             uint32_t,
+                                             uint32_t,
+                                             float,
+                                             bool,
+                                             void*);
+
+extern template void run_fmha_bwd<__half>(__half* do_ptr,
+                                          __half* o_ptr,
+                                          float* softmax_lse_ptr,
+                                          __half* q_ptr,
+                                          __half* k_ptr,
+                                          __half* v_ptr,
+                                          bool* mask_bias_ptr,
+                                          float* triangle_bias_ptr,
+                                          __half* dq_ptr,
+                                          __half* dk_ptr,
+                                          __half* dv_ptr,
+                                          float* triangle_dbias_ptr,
+                                          float* do_o_dot_ptr,
+                                          float* dq_fp32_buf,
+                                          const uint32_t B,
+                                          const uint32_t I,
+                                          const uint32_t H,
+                                          const uint32_t S_qo,
+                                          const uint32_t S_kv,
+                                          const uint32_t D,
+                                          const float bmm_scale,
+                                          bool use_tf32,
+                                          void* stream);
+
+extern template void run_fmha_bwd<__nv_bfloat16>(__nv_bfloat16* do_ptr,
+                                                 __nv_bfloat16* o_ptr,
+                                                 float* softmax_lse_ptr,
+                                                 __nv_bfloat16* q_ptr,
+                                                 __nv_bfloat16* k_ptr,
+                                                 __nv_bfloat16* v_ptr,
+                                                 bool* mask_bias_ptr,
+                                                 float* triangle_bias_ptr,
+                                                 __nv_bfloat16* dq_ptr,
+                                                 __nv_bfloat16* dk_ptr,
+                                                 __nv_bfloat16* dv_ptr,
+                                                 float* triangle_dbias_ptr,
+                                                 float* do_o_dot_ptr,
+                                                 float* dq_fp32_buf,
+                                                 const uint32_t B,
+                                                 const uint32_t I,
+                                                 const uint32_t H,
+                                                 const uint32_t S_qo,
+                                                 const uint32_t S_kv,
+                                                 const uint32_t D,
+                                                 const float bmm_scale,
+                                                 bool use_tf32,
+                                                 void* stream);
+
+extern template void run_fmha_bwd<float>(float* do_ptr,
+                                         float* o_ptr,
+                                         float* softmax_lse_ptr,
+                                         float* q_ptr,
+                                         float* k_ptr,
+                                         float* v_ptr,
+                                         bool* mask_bias_ptr,
+                                         float* triangle_bias_ptr,
+                                         float* dq_ptr,
+                                         float* dk_ptr,
+                                         float* dv_ptr,
+                                         float* triangle_dbias_ptr,
+                                         float* do_o_dot_ptr,
+                                         float* dq_fp32_buf,
+                                         const uint32_t B,
+                                         const uint32_t I,
+                                         const uint32_t H,
+                                         const uint32_t S_qo,
+                                         const uint32_t S_kv,
+                                         const uint32_t D,
+                                         const float bmm_scale,
+                                         bool use_tf32,
+                                         void* stream);
+
+}  // namespace cudnn_fmha
+
+#endif  // CUDNN_FMHA_RUN_FMHA_H

cuequivariance_ops/equivariance/run_fmha_cudafree.h

@@ -0,0 +1,77 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+ *
+ * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+ * property and proprietary rights in and to this material, related
+ * documentation and any modifications thereto. Any use, reproduction,
+ * disclosure or distribution of this material and related documentation
+ * without an express license agreement from NVIDIA CORPORATION or
+ * its affiliates is strictly prohibited.
+ */
+#ifndef CUDNN_FMHA_RUN_FMHA_CUDAFREE_H
+#define CUDNN_FMHA_RUN_FMHA_CUDAFREE_H
+
+#include <cstdint>  // for uint32_t
+
+namespace cudnn_fmha {
+using DType = void;
+
+enum class Datatype : uint32_t {
+  kFloat32  = 0,
+  kFloat64  = 1,
+  kFloat16  = 2,
+  kBFloat16 = 3,
+  kInt32    = 4,
+  kInt64    = 5
+};
+
+__attribute__((visibility("default"))) void run_fmha_for_dtype(
+  Datatype dtype,
+  DType* q_ptr,              // [B, N, H, S_qo, D]
+  DType* k_ptr,              // [B, N, H, S_kv, D]
+  DType* v_ptr,              // [B, N, H, S_kv, D]
+  DType* o_ptr,              // [B, N, H, S_qo, D] output
+  bool* mask_bias_ptr,       // [B, N, 1, 1, S_kv]
+  float* triangle_bias_ptr,  // [B, 1, H, S_qo, S_kv]
+  float* softmax_lse_ptr,    // [B, N, H, S_qo, 1] output
+  float* softmax_max_ptr,    // [B, N, H, S_qo, 1] output
+  const uint32_t B,
+  const uint32_t I,
+  const uint32_t H,
+  const uint32_t S_qo,
+  const uint32_t S_kv,
+  const uint32_t D,
+  const float bmm_scale,
+  bool use_tf32,
+  void* stream = nullptr);
+
+__attribute__((visibility("default"))) void run_fmha_bwd_for_dtype(
+  Datatype dtype,
+  DType* do_ptr,              // [B, N, H, S_qo, D]
+  DType* o_ptr,               // [B, N, H, S_qo, D]
+  float* softmax_lse_ptr,     // [B, N, H, S_qo, 1]
+  DType* q_ptr,               // [B, N, H, S_qo, D]
+  DType* k_ptr,               // [B, N, H, S_kv, D]
+  DType* v_ptr,               // [B, N, H, S_kv, D]
+  bool* mask_bias_ptr,        // [B, N, 1, 1, S_kv]
+  float* triangle_bias_ptr,   // [B, 1, H, S_qo, S_kv]
+  DType* dq_ptr,              // [B, N, H, S_qo, D] output
+  DType* dk_ptr,              // [B, N, H, S_kv, D] output
+  DType* dv_ptr,              // [B, N, H, S_kv, D] output
+  float* triangle_dbias_ptr,  // [B, 1, H, S_qo, S_kv] output
+  float* do_o_dot_ptr,        // [B, N, H, S_qo, 1] worspace
+  float* dq_fp32_buf_ptr,     // [B, N, H, S_qo, D] workspace
+  const uint32_t B,
+  const uint32_t I,
+  const uint32_t H,
+  const uint32_t S_qo,
+  const uint32_t S_kv,
+  const uint32_t D,
+  const float bmm_scale,
+  bool use_tf32,
+  void* stream);
+
+}  // namespace cudnn_fmha
+
+#endif  // CUDNN_FMHA_RUN_FMHA_CUDAFREE_H

cuequivariance_ops/equivariance/tensor_product_uniform_1d_jit.hh

@@ -8,49 +8,31 @@
 
 #pragma once
 
+#include "dtypes.hh"  // for Datatype
 #include <cstdint>
 #include <string>
 #include <vector>
 
 namespace kernelcatcher::equivariance::tensor_product_uniform_1d_jit {
+using namespace kernelcatcher::utils;
 
-enum class Datatype : int {
-  kFloat32  = 0,
-  kFloat64  = 1,
-  kFloat16  = 2,
-  kBFloat16 = 3,
-  kInt32    = 4,
-  kInt64    = 5
-};
 enum class Dimension : int { kScalar = 0, kOneDimensional = 1 };
 enum class BatchDimension : int { kBatched = 0, kShared = 1, kIndexed = 2 };
 
-extern int run_tensor_product_uniform_1d_jit(
-  std::string const& name,
-  Datatype math_dtype,
-  int operand_extent,
-  int num_inputs,
-  int num_outputs,
-  int num_index,
-  std::vector<Dimension> const& buffer_dim,     // num_inputs + num_outputs
-  std::vector<int> const& buffer_num_segments,  // num_inputs + num_outputs
-  std::vector<std::vector<BatchDimension>> const&
-    batch_dim,  // (num_inputs + num_outputs + num_index) x num_batch_axes
-  std::vector<std::vector<int>> const&
-    index_buffer,  // (num_inputs + num_outputs + num_index) x num_batch_axes
-  // ^^^^^ ignored unless batch_dim[i][j] is indexed
-  std::vector<int> const& index_extent,  // num_index
-  std::vector<Datatype> const& dtypes,   // num_inputs + num_outputs + num_index
-  std::vector<std::vector<int>> const& operations,
-  std::vector<int> num_paths,                // num_operations
-  std::vector<int> path_indices_start,       // num_operations
-  std::vector<int> path_coefficients_start,  // num_operations
-  std::vector<int> const& path_indices,
-  std::vector<double> const& path_coefficients,
-  std::vector<int> const& batch_sizes,  // num_batch_axes
-  std::vector<void*> const& buffers,
-  std::vector<size_t> const& buffer_bytes,
-  bool zero_output_buffers,
-  void* stream);
+#define KC_UNIFORM_1D_DECL_ARGUMENTS                                                         \
+  std::string const &name, Datatype math_dtype, int operand_extent, int num_inputs,          \
+    int num_outputs, int num_index, std::vector<Dimension> const &buffer_dim,                \
+    std::vector<int> const &buffer_num_segments,                                             \
+    std::vector<std::vector<BatchDimension>> const &batch_dim,                               \
+    std::vector<std::vector<int>> const &index_buffer, std::vector<int> const &index_extent, \
+    std::vector<Datatype> const &dtypes, std::vector<std::vector<int>> const &operations,    \
+    std::vector<int> num_paths, std::vector<int> path_indices_start,                         \
+    std::vector<int> path_coefficients_start, std::vector<int> const &path_indices,          \
+    std::vector<double> const &path_coefficients, std::vector<int> const &batch_sizes,       \
+    std::vector<void*> const &buffers, std::vector<size_t> const &buffer_bytes,              \
+    bool zero_output_buffers
+
+extern int run_tensor_product_uniform_1d_jit(KC_UNIFORM_1D_DECL_ARGUMENTS, void* stream);
+extern int run_tensor_product_uniform_1d_cpu(KC_UNIFORM_1D_DECL_ARGUMENTS);
 
 }  // namespace kernelcatcher::equivariance::tensor_product_uniform_1d_jit