add custom op declaration

Author: DrRyanHuang

custom_ops/gpu_ops/cpp_extensions.cc

@@ -530,7 +530,7 @@ paddle::Tensor FusedHadamardQuantFp8Func(
 int64_t init_custom_all_reduce(const std::vector<int64_t>& fake_ipc_ptrs,
                       paddle::Tensor& rank_data, int64_t rank, bool full_nvlink);
 
-void all_reduce(int64_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+void all_reduce(paddle::Tensor& inp, paddle::Tensor& out, int64_t _fa,
                 int64_t reg_buffer, int64_t reg_buffer_sz_bytes);
 
 void dispose(int64_t _fa);

custom_ops/gpu_ops/custom_all_reduce/all_reduce.cu

@@ -49,7 +49,7 @@ fptr_t init_custom_all_reduce(const std::vector<fptr_t>& fake_ipc_ptrs,
  * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
  * copied into _reg_buffer.
  */
-void all_reduce(fptr_t _fa, paddle::Tensor& inp, paddle::Tensor& out,
+void all_reduce(paddle::Tensor& inp, paddle::Tensor& out, fptr_t _fa,
                 fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
   auto fa = reinterpret_cast<paddle::CustomAllreduce*>(_fa);
   auto stream = inp.stream();
@@ -163,3 +163,12 @@ fptr_t open_mem_handle(paddle::Tensor& mem_handle) {
 void free_shared_buffer(fptr_t buffer) {
   CUDACHECK(cudaFree(reinterpret_cast<void*>(buffer)));
 }
+
+
+PD_BUILD_STATIC_OP(all_reduce)
+    .Inputs({"inp",
+             "out"})
+    .Outputs({"new_out"})
+    .Attrs({"_fa: int64_t", "_reg_buffer: int64_t", "reg_buffer_sz_bytes: int64_t"})
+    .SetInplaceMap({{"out", "new_out"}})
+    .SetKernelFn(PD_KERNEL(all_reduce));

fastdeploy/distributed/communication.py

@@ -42,22 +42,17 @@ def use_custom_allreduce(custom_all_reduce_max_bytes: int = 8192 * 1024):
     _TP_AR = CustomAllreduce(model_parallel_group, custom_all_reduce_max_bytes)
 
 
-try:
-
-    @paddle.jit.marker.unified
-    def tensor_model_parallel_all_reduce(
-        input_: paddle.Tensor,
-    ) -> paddle.Tensor:
-        """All-reduce the input tensor across model parallel group."""
-        global _TP_AR
-        if _TP_AR is not None and _TP_AR.should_custom_ar(input_):
-            _TP_AR.custom_all_reduce(input_)
-        elif paddle.in_dynamic_mode():
-            hcg = fleet.get_hybrid_communicate_group()
-            mp_group = hcg.get_model_parallel_group()
-            dist.all_reduce(input_, group=mp_group)
-        else:
-            dist.all_reduce(input_)
-
-except:
-    tensor_model_parallel_all_reduce = None
+@paddle.jit.marker.unified
+def tensor_model_parallel_all_reduce(
+    input_: paddle.Tensor,
+) -> paddle.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    global _TP_AR
+    if _TP_AR is not None and _TP_AR.should_custom_ar(input_):
+        _TP_AR.custom_all_reduce(input_)
+    elif paddle.in_dynamic_mode():
+        hcg = fleet.get_hybrid_communicate_group()
+        mp_group = hcg.get_model_parallel_group()
+        dist.all_reduce(input_, group=mp_group)
+    else:
+        dist.all_reduce(input_)

fastdeploy/distributed/custom_all_reduce/custom_all_reduce.py

@@ -158,9 +158,9 @@ def all_reduce(
         if out is None:
             out = paddle.empty_like(inp)
         if registered:
-            all_reduce(self._ptr, inp, out, 0, 0)
+            all_reduce(inp, out, self._ptr, 0, 0)
         else:
-            all_reduce(self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size)
+            all_reduce(inp, out, self._ptr, self.buffer_ptrs[self.rank], self.max_size)
         return out
 
     def start_capture(self):

fastdeploy/model_executor/layers/attention/mla_attention_backend.py

@@ -89,6 +89,9 @@ class MLAAttentionMetadata(AttentionMetadata):
     kv_signal_metadata: Optional[paddle.Tensor] = None
     kv_signal_data_list: List[Optional[paddle.Tensor]] = field(default_factory=list)
 
+    max_enc_len_this_time: Optional[paddle.Tensor] = None
+    max_dec_len_this_time: Optional[paddle.Tensor] = None
+
 
 class MLAAttentionBackend(AttentionBackend):
     """