T2-1-1/T2-1-2: 吃kfc

include/infiniccl.h

@@ -11,6 +11,12 @@ typedef enum {
     INFINICCL_AVG = 4,
 } infinicclReduceOp_t;
 
+typedef enum {
+    INFINICCL_ALLREDUCE_BACKEND_AUTO = 0,
+    INFINICCL_ALLREDUCE_BACKEND_NCCL = 1,
+    INFINICCL_ALLREDUCE_BACKEND_CUSTOM = 2,
+} infinicclAllReduceBackend_t;
+
 struct InfinicclComm;
 
 typedef struct InfinicclComm *infinicclComm_t;
@@ -23,6 +29,30 @@ __INFINI_C __export infiniStatus_t infinicclCommInitAll(
 
 __INFINI_C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm);
 
+__INFINI_C __export infiniStatus_t infinicclCommSetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t backend);
+
+__INFINI_C __export infiniStatus_t infinicclCommGetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t *backend);
+
+__INFINI_C __export infiniStatus_t infinicclCommRegisterAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes);
+
+__INFINI_C __export infiniStatus_t infinicclCommRegisterAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes);
+
+__INFINI_C __export infiniStatus_t infinicclCommClearAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice);
+
 __INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm);
 
 __INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm);

include/infinicore/ops.hpp

@@ -26,6 +26,7 @@
 #include "ops/cross_entropy.hpp"
 #include "ops/deepseek_moe.hpp"
 #include "ops/embedding.hpp"
+#include "ops/ernie45_rope.hpp"
 #include "ops/flash_attention.hpp"
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"

include/infinicore/ops/ernie45_rope.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Ernie45MRoPE, Tensor, Tensor, const Tensor &, double, size_t, size_t, size_t);
+INFINICORE_GRAPH_OP_CLASS(Ernie45VisionRoPE, Tensor, Tensor, const Tensor &, double);
+
+Tensor ernie45_mrope_(Tensor q,
+                      Tensor k,
+                      const Tensor &positions,
+                      double rope_theta,
+                      size_t section_h,
+                      size_t section_w,
+                      size_t section_t);
+
+Tensor ernie45_vision_rope_(Tensor q,
+                            Tensor k,
+                            const Tensor &positions,
+                            double rope_theta);
+
+} // namespace infinicore::op

include/infiniop/ops/ernie45_rope.h

@@ -0,0 +1,56 @@
+#ifndef __INFINIOP_ERNIE45_ROPE_API_H__
+#define __INFINIOP_ERNIE45_ROPE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopErnie45MropeDescriptor_t;
+typedef struct InfiniopDescriptor *infiniopErnie45VisionRopeDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateErnie45MropeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45MropeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta,
+    size_t section_h,
+    size_t section_w,
+    size_t section_t);
+
+__INFINI_C __export infiniStatus_t infiniopGetErnie45MropeWorkspaceSize(infiniopErnie45MropeDescriptor_t desc,
+                                                                        size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopErnie45Mrope(
+    infiniopErnie45MropeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyErnie45MropeDescriptor(infiniopErnie45MropeDescriptor_t desc);
+
+__INFINI_C __export infiniStatus_t infiniopCreateErnie45VisionRopeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45VisionRopeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta);
+
+__INFINI_C __export infiniStatus_t infiniopGetErnie45VisionRopeWorkspaceSize(infiniopErnie45VisionRopeDescriptor_t desc,
+                                                                             size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopErnie45VisionRope(
+    infiniopErnie45VisionRopeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyErnie45VisionRopeDescriptor(infiniopErnie45VisionRopeDescriptor_t desc);
+
+#endif