cuequivariance-ops-cu12 0.8.1__py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

cuequivariance_ops/equivariance/fused_tensor_product.cuh

@@ -0,0 +1,297 @@
+/*
+ * Copyright (c) 2021-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
+
+#include "../common/common.hpp"
+
+#include <algorithm>
+#include <limits>
+
+namespace kernelcatcher::tensor_product {
+
+struct __attribute__((aligned(16))) tp_data_sizes {
+  int64_t batch_size;
+  bool shared_a;
+  bool shared_b;
+  bool shared_w;
+  int32_t stride_a;
+  int32_t stride_b;
+  int32_t stride_w;
+  int32_t stride_o;
+};  // struct tp_data_sizes
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_fwd(DataOutT* out,
+                              const DataAT* in_a,
+                              const DataBT* in_b,
+                              const DataWeightT* weight,
+                              ConnectionModeT mode,
+                              const tp_info<MathT>& info,
+                              const tp_data_sizes& sizes,
+                              cudaStream_t stream);
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_bwd(DataAT* grad_in_a,
+                              DataBT* grad_in_b,
+                              DataWeightT* grad_weight,
+                              const DataOutT* grad_out,
+                              const DataAT* in_a,
+                              const DataBT* in_b,
+                              const DataWeightT* weight,
+                              ConnectionModeT mode,
+                              const tp_info<MathT>& info_bwd_dgrad_a,
+                              const tp_info<MathT>& info_bwd_dgrad_b,
+                              const tp_info<MathT>& info_bwd_dgrad_w,
+                              const tp_data_sizes& sizes,
+                              cudaStream_t stream);
+
+template <typename DataAT, typename DataBT, typename DataWeightT, typename DataOutT, typename MathT>
+void fused_tensor_product_bwd_bwd(DataAT* grad_in_a,
+                                  DataBT* grad_in_b,
+                                  DataWeightT* grad_weight,
+                                  DataOutT* grad_grad_out,
+                                  const DataAT* grad_grad_in_a,
+                                  const DataBT* grad_grad_in_b,
+                                  const DataWeightT* grad_grad_weight,
+                                  const DataOutT* grad_out,
+                                  const DataAT* in_a,
+                                  const DataBT* in_b,
+                                  const DataWeightT* weight,
+                                  ConnectionModeT mode,
+                                  const tp_info<MathT>& info_fwd,
+                                  const tp_info<MathT>& info_bwd_dgrad_a,
+                                  const tp_info<MathT>& info_bwd_dgrad_b,
+                                  const tp_info<MathT>& info_bwd_dgrad_w,
+                                  const tp_data_sizes& sizes,
+                                  cudaStream_t stream);
+
+extern template void fused_tensor_product_bwd_bwd<float, float, float, float, float>(
+  float*,
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<float, float, float, float, double>(
+  float*,
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<double, double, double, double, double>(
+  double*,
+  double*,
+  double*,
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd_bwd<__half, __half, __half, __half, float>(
+  __half*,
+  __half*,
+  __half*,
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_bwd_bwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<float, float, float, float, float>(
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<float, float, float, float, double>(
+  float*,
+  float*,
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<double, double, double, double, double>(
+  double*,
+  double*,
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_bwd<__half, __half, __half, __half, float>(
+  __half*,
+  __half*,
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_bwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_fwd<float, float, float, float, float>(
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void fused_tensor_product_fwd<float, float, float, float, double>(
+  float*,
+  const float*,
+  const float*,
+  const float*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void fused_tensor_product_fwd<double, double, double, double, double>(
+  double*,
+  const double*,
+  const double*,
+  const double*,
+  ConnectionModeT,
+  const tp_info<double>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+extern template void fused_tensor_product_fwd<__half, __half, __half, __half, float>(
+  __half*,
+  const __half*,
+  const __half*,
+  const __half*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+extern template void
+fused_tensor_product_fwd<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16, __nv_bfloat16, float>(
+  __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  const __nv_bfloat16*,
+  ConnectionModeT,
+  const tp_info<float>&,
+  const tp_data_sizes&,
+  cudaStream_t);
+
+}  // namespace kernelcatcher::tensor_product

cuequivariance_ops/equivariance/indexed_linear.hh

@@ -0,0 +1,41 @@
+/*
+ * 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
+
+#ifndef CUEQUIVARIANCE_OPS_WITH_CUBLAS
+#pragma message( \
+    "indexed_linear functions require cuBLAS support. Set CUEQUIVARIANCE_OPS_WITH_CUBLAS=1 environment variable before building to enable indexed_linear functionality. When cuBLAS is disabled, indexed_linear functions will return error codes.")
+#endif
+
+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,   \
+    int compute_type, 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, compute_type, \
+    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