[PP] Refactor PP to async mode (#11852)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: bluecoffee8 <jasperli2002@gmail.com>
Co-authored-by: zhangxiaolei123456 <zhangxiaolei.666@bytedance.com>
Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com>

python/sglang/srt/managers/tp_worker.py

@@ -376,6 +376,7 @@ class TpModelWorker(BaseTpWorker):
         self,
         model_worker_batch: ModelWorkerBatch,
         forward_batch: Optional[ForwardBatch] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
         is_verify: bool = False,
         skip_attn_backend_init=False,
     ) -> GenerationBatchResult:
@@ -391,14 +392,6 @@ class TpModelWorker(BaseTpWorker):
             # FIXME(lsyin): unify the interface of forward_batch
             assert forward_batch is not None
 
-        pp_proxy_tensors = None
-        if not self.pp_group.is_first_rank:
-            pp_proxy_tensors = PPProxyTensors(
-                self.pp_group.recv_tensor_dict(
-                    all_gather_group=self.get_attention_tp_group()
-                )
-            )
-
         if self.pp_group.is_last_rank:
             if self.is_dllm():
                 return self._forward_batch_generation_dllm(forward_batch)