Merge branch 'main' into issue/500

Author: voltjia

include/infinicore/tensor.hpp

@@ -27,6 +27,7 @@ struct TensorMetaData {
     infiniopTensorDescriptor_t desc;
 
     TensorMetaData(const Shape &shape, const Strides &strides, const DataType &dtype);
+    ~TensorMetaData();
 };
 
 struct TensorData {
@@ -121,6 +122,10 @@ class TensorImpl : public std::enable_shared_from_this<TensorImpl> {
 
     std::string info() const;
 
+    void debug(const std::string &filename) const;
+
+    void debug() const;
+
     ///
     /// Data Transfer APIs
     ///

include/infiniop.h

@@ -6,6 +6,7 @@
 #include "infiniop/ops/attention.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"
+#include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/conv.h"
 #include "infiniop/ops/dequantize_awq.h"
 #include "infiniop/ops/gemm.h"
@@ -19,6 +20,8 @@
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/ops/topkrouter.h"
+#include "infiniop/ops/topksoftmax.h"
+#include "infiniop/ops/sigmoid.h"
 #include "infiniop/tensor_descriptor.h"
 
 #endif // __INFINIOP_API_H__

include/infiniop/ops/logsoftmax.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_LOGSOFTMAX_API_H__
+#define __INFINIOP_LOGSOFTMAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogSoftmaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogSoftmaxDescriptor(infiniopHandle_t handle,
+                                                                  infiniopLogSoftmaxDescriptor_t *desc_ptr,
+                                                                  infiniopTensorDescriptor_t y_desc,
+                                                                  infiniopTensorDescriptor_t x_desc);
+
+__C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogSoftmax(infiniopLogSoftmaxDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *y,
+                                               const void *x,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc);
+
+#endif

include/infiniop/ops/sigmoid.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_SIGMOID_API_H__
+#define __INFINIOP_SIGMOID_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopSigmoidDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateSigmoidDescriptor(infiniopHandle_t handle,
+                                                            infiniopSigmoidDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t y,
+                                                            infiniopTensorDescriptor_t x);
+
+__C __export infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopSigmoid(infiniopSigmoidDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *y,
+                                            const void *x,
+                                            void *stream);
+
+__C __export infiniStatus_t infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc);
+
+#endif

include/infiniop/ops/topkrouter.h

@@ -5,16 +5,23 @@
 
 typedef struct InfiniopDescriptor *infiniopTopkrouterDescriptor_t;
 
-__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(
-    infiniopHandle_t handle,
-    infiniopTopkrouterDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t correction_bias_desc);
+__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t handle,
+                                                               infiniopTopkrouterDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t x_desc,
+                                                               infiniopTensorDescriptor_t correction_bias_desc);
 
 __C __export infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterDescriptor_t desc, size_t *size);
 
-__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc, void *workspace, size_t workspace_size,
-                                               void *values, void *indices, void *x, void *correction_bias, float routed_scaling_factor, size_t topk, void *stream);
+__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *values,
+                                               void *indices,
+                                               const void *x,
+                                               const void *correction_bias,
+                                               const float routed_scaling_factor,
+                                               const size_t topk,
+                                               void *stream);
 
 __C __export infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouterDescriptor_t desc);
 

include/infiniop/ops/topksoftmax.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_TOPKSOFTMAX_API_H__
+#define __INFINIOP_TOPKSOFTMAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTopksoftmaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t handle,
+                                                                infiniopTopksoftmaxDescriptor_t *desc_ptr,
+                                                                infiniopTensorDescriptor_t x_desc);
+
+__C __export infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftmaxDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t desc,
+                                                void *workspace,
+                                                size_t workspace_size,
+                                                void *values,
+                                                void *indices,
+                                                const void *x,
+                                                const size_t topk,
+                                                const int norm,
+                                                void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftmaxDescriptor_t desc);
+
+#endif

python/infinicore/tensor.py

@@ -75,6 +75,17 @@ def permute(self, dims):
     def view(self, shape):
         return Tensor(self._underlying.view(shape))
 
+    def debug(self, filename=None):
+        """Print tensor data or save to file for debugging
+        
+        Args:
+            filename: Optional filename to save raw binary data. If None, prints to stdout.
+        """
+        if filename is None:
+            self._underlying.debug()
+        else:
+            self._underlying.debug(filename)
+
 
 def empty(size, *, dtype=None, device=None, pin_memory=False):
     return Tensor(

scripts/python_test.py

@@ -17,6 +17,7 @@ def run_tests(args):
         "causal_softmax.py",
         "clip.py",
         "gemm.py",
+        "logsoftmax.py",
         "mul.py",
         "random_sample.py",
         "rearrange.py",
@@ -25,6 +26,9 @@ def run_tests(args):
         "sub.py",
         "swiglu.py",
         "softplus.py",
+        "sigmoid.py",
+        "topkrouter.py",
+        "topksoftmax.py",
     ]:
         result = subprocess.run(
             f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True