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

cuequivariance_ops/equivariance/run_fmha_cudafree.h

@@ -0,0 +1,176 @@
+/*
+ * 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
+#include <optional>
+
+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,
+  bool zero_init_dbias_dq_buf = true);
+
+// Shared sanity checking functions for kernel requirements
+// These can be called from bindings before invoking kernel functions
+__attribute__((visibility("default"))) void validate_fmha_fwd_params(
+  int bits,        // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  uint32_t D,      // Head dimension
+  bool use_tf32);  // Whether TF32 is requested
+
+__attribute__((visibility("default"))) void validate_fmha_bwd_params(
+  int bits,        // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  bool use_tf32);  // Whether TF32 is requested
+
+// Shared kernel selection helpers
+// Determine whether to use SM100 kernels based on conditions
+__attribute__((visibility("default"))) bool should_use_sm100f_fwd(
+  int bits,  // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  bool has_triangle_bias_same_type,  // Whether triangle_bias has same dtype as q/k/v
+  uint32_t D,                        // Head dimension
+  uint32_t S_kv,                     // Key/value sequence length
+  bool mask_consistent,  // Whether mask_bias_ptr and actual_s_kv_ptr are both null or both non-null
+  void* stream,          // CUDA stream (for device capability check)
+  const std::optional<std::vector<int>> device_cc =
+    std::nullopt);  // Device capability (compute capability version)
+
+__attribute__((visibility("default"))) bool should_use_sm100f_bwd(
+  int bits,  // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  bool has_triangle_bias_same_type,  // Whether triangle_bias has same dtype as q/k/v
+  uint32_t D,                        // Head dimension
+  bool has_dbias_fp32_buf,           // Whether dbias_fp32_buf_ptr is provided
+  void* stream,                      // CUDA stream (for device capability check)
+  const std::optional<std::vector<int>> device_cc =
+    std::nullopt);  // Device capability (compute capability version)
+
+// CUDA-free wrappers for SM100 kernels (can be called from non-CUDA code)
+__attribute__((visibility("default"))) void run_fmha_sm100_for_dtype(
+  Datatype dtype,
+  void* q_ptr,              // [B, N, H, S_qo, D]
+  void* k_ptr,              // [B, N, H, S_kv, D]
+  void* v_ptr,              // [B, N, H, S_kv, D]
+  void* o_ptr,              // [B, N, H, S_qo, D] output
+  bool* mask_bias_ptr,      // [B, N, 1, 1, S_kv] (can be nullptr)
+  int* actual_s_kv_ptr,     // [B, N] (can be nullptr, must match mask_bias_ptr)
+  void* triangle_bias_ptr,  // [B, 1, H, S_qo, S_kv] (same dtype as q/k/v)
+  float* softmax_lse_ptr,   // [B, N, H, S_qo, 1] output
+  float* softmax_max_ptr,   // [B, N, H, S_qo, 1] output
+  uint32_t B,
+  uint32_t I,
+  uint32_t H,
+  uint32_t S_qo,
+  uint32_t S_kv,
+  uint32_t D,
+  float bmm_scale,
+  void* stream);
+
+__attribute__((visibility("default"))) void run_fmha_bwd_sm100_for_dtype(
+  Datatype dtype,
+  void* q_ptr,                // [B, N, H, S_qo, D]
+  void* k_ptr,                // [B, N, H, S_kv, D]
+  void* v_ptr,                // [B, N, H, S_kv, D]
+  void* o_ptr,                // [B, N, H, S_qo, D]
+  bool* mask_bias_ptr,        // [B, N, 1, 1, S_kv] (can be nullptr)
+  void* triangle_bias_ptr,    // [B, 1, H, S_qo, S_kv] (same dtype as q/k/v)
+  float* softmax_lse_ptr,     // [B, N, H, S_qo, 1]
+  void* do_ptr,               // [B, N, H, S_qo, D]
+  void* dq_ptr,               // [B, N, H, S_qo, D] output
+  void* dk_ptr,               // [B, N, H, S_kv, D] output
+  void* dv_ptr,               // [B, N, H, S_kv, D] output
+  void* triangle_dbias_ptr,   // [B, 1, H, S_qo, S_kv] output (same dtype as q/k/v)
+  float* do_o_dot_ptr,        // [B, N, H, S_qo, 1] workspace
+  float* dq_fp32_buf_ptr,     // [B, N, H, S_qo, D] workspace
+  float* dbias_fp32_buf_ptr,  // [B, 1, H, padded_S_qo, padded_S_kv] workspace
+  uint32_t B,
+  uint32_t I,
+  uint32_t H,
+  uint32_t S_qo,
+  uint32_t S_kv,
+  uint32_t D,
+  float bmm_scale,
+  bool zero_init_dbias_dq_buf,  // Whether to zero initialize dq_fp32_buf and dbias_fp32_buf
+  void* stream);
+
+// Returns true if FP32 bias is needed for forward pass, false if same-dtype bias is used (SM100
+// case) Takes minimal arguments needed to determine bias dtype requirement for forward pass
+__attribute__((visibility("default"))) bool needs_fp32_bias_fwd(
+  int bits,       // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  uint32_t D,     // Head dimension
+  uint32_t S_kv,  // Key/value sequence length (for forward pass constraint)
+  void* stream);  // CUDA stream (for device capability check)
+
+// Returns true if FP32 bias input is needed for backward pass, false if same-dtype bias is used
+// (SM100 case) Takes minimal arguments needed to determine bias input dtype requirement for
+// backward pass.
+__attribute__((visibility("default"))) bool needs_fp32_bias_bwd(
+  int bits,                 // Bit width of dtype (16 for FP16/BF16, 32 for FP32, 64 for FP64)
+  uint32_t D,               // Head dimension
+  bool has_dbias_fp32_buf,  // Whether dbias FP32 workspace buffer is provided (required for SM100)
+  void* stream);            // CUDA stream (for device capability check)
+
+}  // namespace cudnn_fmha
+
+#endif  // CUDNN_FMHA_RUN_FMHA_CUDAFREE_H

cuequivariance_ops/equivariance/run_fmha_sm100.h

@@ -0,0 +1,135 @@
+#pragma once
+
+#include <cuda_fp16.h>
+
+namespace cudnn_fmha {
+
+template <typename DType>
+void run_fmha_sm100(DType* q_ptr,
+                    DType* k_ptr,
+                    DType* v_ptr,
+                    DType* o_ptr,
+                    bool* mask_bias_ptr,
+                    int* actual_s_kv_ptr,
+                    DType* 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,
+                    void* stream_ = nullptr);
+
+extern template void run_fmha_sm100<__nv_bfloat16>(__nv_bfloat16*,
+                                                   __nv_bfloat16*,
+                                                   __nv_bfloat16*,
+                                                   __nv_bfloat16*,
+                                                   bool*,
+                                                   int*,
+                                                   __nv_bfloat16*,
+                                                   float*,
+                                                   float*,
+                                                   const uint32_t,
+                                                   const uint32_t,
+                                                   const uint32_t,
+                                                   const uint32_t,
+                                                   const uint32_t,
+                                                   const uint32_t,
+                                                   const float,
+                                                   void*);
+
+extern template void run_fmha_sm100<__half>(__half*,
+                                            __half*,
+                                            __half*,
+                                            __half*,
+                                            bool*,
+                                            int*,
+                                            __half*,
+                                            float*,
+                                            float*,
+                                            const uint32_t,
+                                            const uint32_t,
+                                            const uint32_t,
+                                            const uint32_t,
+                                            const uint32_t,
+                                            const uint32_t,
+                                            const float,
+                                            void*);
+
+template <typename DType>
+void run_fmha_bwd_sm100(DType* q_ptr,
+                        DType* k_ptr,
+                        DType* v_ptr,
+                        DType* o_ptr,
+                        bool* mask_bias_ptr,
+                        DType* triangle_bias_ptr,
+                        float* softmax_lse_ptr,
+                        DType* do_ptr,
+                        DType* dq_ptr,
+                        DType* dk_ptr,
+                        DType* dv_ptr,
+                        DType* dtriangle_bias_ptr,
+                        float* do_o_dot_tmp_ptr,
+                        float* dq_tmp_ptr,              // must be zero initialized
+                        float* dtriangle_bias_tmp_ptr,  // must be zero initialized
+                        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,
+                        void* stream_ = nullptr);
+
+extern template void run_fmha_bwd_sm100<__nv_bfloat16>(__nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       bool*,
+                                                       __nv_bfloat16*,
+                                                       float*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       __nv_bfloat16*,
+                                                       float*,
+                                                       float*,
+                                                       float*,
+                                                       const uint32_t,
+                                                       const uint32_t,
+                                                       const uint32_t,
+                                                       const uint32_t,
+                                                       const uint32_t,
+                                                       const uint32_t,
+                                                       const float,
+                                                       void* stream_);
+
+extern template void run_fmha_bwd_sm100<__half>(__half*,
+                                                __half*,
+                                                __half*,
+                                                __half*,
+                                                bool*,
+                                                __half*,
+                                                float*,
+                                                __half*,
+                                                __half*,
+                                                __half*,
+                                                __half*,
+                                                __half*,
+                                                float*,
+                                                float*,
+                                                float*,
+                                                const uint32_t,
+                                                const uint32_t,
+                                                const uint32_t,
+                                                const uint32_t,
+                                                const uint32_t,
+                                                const uint32_t,
+                                                const float,
+                                                void* stream_);
+
+}  // namespace cudnn_fmha

cuequivariance_ops/equivariance/segmented_transpose.cuh

@@ -0,0 +1,40 @@
+/*
+ * 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"
+
+namespace kernelcatcher::tensor_product {
+
+template <typename DataT>
+void segmented_transpose(DataT* tensor_transpose,
+                         const DataT* tensor,
+                         const int32_t* segment_info,
+                         int32_t num_segments,
+                         int64_t batch_size,
+                         int64_t stride,
+                         bool input_contiguous_as_info,
+                         cudaStream_t stream);
+
+extern template void segmented_transpose<float>(
+  float*, const float*, const int32_t*, int32_t, int64_t, int64_t, bool, cudaStream_t);
+extern template void segmented_transpose<double>(
+  double*, const double*, const int32_t*, int32_t, int64_t, int64_t, bool, cudaStream_t);
+extern template void segmented_transpose<__nv_bfloat16>(__nv_bfloat16*,
+                                                        const __nv_bfloat16*,
+                                                        const int32_t*,
+                                                        int32_t,
+                                                        int64_t,
+                                                        int64_t,
+                                                        bool,
+                                                        cudaStream_t);
+extern template void segmented_transpose<__half>(
+  __half*, const __half*, const int32_t*, int32_t, int64_t, int64_t, bool, cudaStream_t);
+
+}  // namespace kernelcatcher::tensor_product

cuequivariance_ops/equivariance/tensor_product_uniform_1d_jit.hh

@@ -0,0 +1,38 @@
+/*
+ * 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 "batch_dimension.hh"  // for BatchDimension
+#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 Dimension : int { kScalar = 0, kOneDimensional = 1 };
+
+#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> const &num_paths, std::vector<int> const &path_indices_start,           \
+    std::vector<int> const &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

cuequivariance_ops/gpu_timing_kernels.hh

@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+/*
+ * GPU Events Implementation Attribution
+ *
+ * The CUDA event recording and timing functionality (record_event, event_elapsed)
+ * has been adapted from JAX's GPU events implementation that was removed in version 0.7.2.
+ *
+ * Original source: https://github.com/jax-ml/jax/
+ * License: Apache License 2.0
+ *
+ * JAX Copyright 2018 The JAX Authors.
+ * Licensed under the Apache License, Version 2.0.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace kernelcatcher::gpu_utilities {
+
+int run_sleep(float* seconds, int64_t* elapsed_ticks, void* stream);
+int run_synchronize(float* elapsed_seconds, void* stream);
+
+// Record a CUDA event on the given stream
+// Returns 0 on success, non-zero on error
+int record_event(uint64_t* event_handle, void* stream, bool copy_before);
+
+// Calculate elapsed time between two events
+// Returns 0 on success, non-zero on error
+int event_elapsed(const uint64_t* start_event_handle,
+                  const uint64_t* end_event_handle,
+                  float* elapsed_ms,
+                  void* stream);
+
+}  // namespace kernelcatcher::gpu_utilities

cuequivariance_ops/sleep.hh

@@ -0,0 +1,40 @@
+/*
+ * 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
+
+/*
+ * DEPRECATED HEADER: This header has been renamed to gpu_timing_kernels.hh
+ *
+ * This compatibility header provides backward compatibility for code that includes
+ * the old sleep.hh header name. The header has been renamed to better reflect
+ * its content (GPU timing and synchronization kernels).
+ *
+ * For new code, please use:
+ *     #include "gpu_timing_kernels.hh"
+ *
+ * OLD (deprecated but still works):
+ *     #include "sleep.hh"
+ *
+ * NEW (recommended):
+ *     #include "gpu_timing_kernels.hh"
+ *
+ * This compatibility header will be removed in a future version.
+ */
+
+#warning \
+  "The 'sleep.hh' header has been renamed to 'gpu_timing_kernels.hh'. Please update your #include statements. This compatibility header will be removed in a future version."
+
+// Include the new header to maintain compatibility
+#include "gpu_timing_kernels.hh"
+
+// Provide backward compatibility aliases
+namespace kernelcatcher::sleep {
+using kernelcatcher::gpu_utilities::run_sleep;
+using kernelcatcher::gpu_utilities::run_synchronize;
+}  // namespace kernelcatcher::sleep

cuequivariance_ops/triton/__init__.py

@@ -0,0 +1,66 @@
+# 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.
+
+from .fused_layer_norm_triton import (
+    Layout,
+    layer_norm_transpose_backward_kernel,
+    layer_norm_transpose_backward_single_pass_kernel,
+    layer_norm_transpose_forward_kernel,
+    layer_norm_transpose_forward_single_pass_kernel,
+)
+
+from .gated_gemm_triton import (
+    fused_sigmoid_gated_dual_gemm_backward_pregemm_kernel,
+    fused_sigmoid_gated_dual_gemm_forward_kernel,
+)
+from .utils import Precision
+from .tuning_decorator import autotune_aot
+from .cache_manager import get_cache_manager
+
+cached_kernels = [ "fused_sigmoid_gated_dual_gemm_backward_pregemm_kernel",
+                   "fused_sigmoid_gated_dual_gemm_forward_kernel",
+                  ]
+
+def init_triton_cache():
+    """
+    Initializes Triton cache manager by pre-loading cache for all available kernels.
+    This function is useful to initialize cache in eager mode before running torch.compile()'d methods
+    that cannot handle cache initialization code
+    """
+    mgr = get_cache_manager()
+    for kernel in cached_kernels:
+        mgr.load_cache(kernel+'_wrapper')
+
+from .utils import (
+    Precision,
+)
+
+from .pair_bias import (
+    pair_bias_norm_linear_mask_forward_kernel,
+    pair_bias_linear_mask_forward_kernel,
+    pair_bias_mask_forward_kernel,
+)
+
+
+__all__ = [
+    "Precision",
+    "fused_sigmoid_gated_dual_gemm_backward_pregemm_kernel",
+    "fused_sigmoid_gated_dual_gemm_forward_kernel",
+    "layer_norm_transpose_backward_kernel",
+    "layer_norm_transpose_backward_single_pass_kernel",
+    "layer_norm_transpose_forward_kernel",
+    "layer_norm_transpose_forward_single_pass_kernel",
+    "pair_bias_norm_linear_mask_forward_kernel",
+    "pair_bias_linear_mask_forward_kernel",
+    "pair_bias_mask_forward_kernel",
+    "autotune_aot",
+    "get_cache_manager",
+    "init_triton_cache"
+] + cached_kernels