fix: fix

Author: MMadhatter

rtp_llm/models_py/bindings/common/RtpProcessGroup.cc

@@ -38,7 +38,8 @@ RtpProcessGroup::RtpProcessGroup(RtpProcessGroupType type) {
 }
 
 void RtpProcessGroup::broadcast(std::vector<torch::Tensor>& input, int rootRank) {
-    std::vector<BufferPtr> buffers;
+    ScopedCUDAStreamContext stream_ctx(device_);
+    std::vector<BufferPtr>  buffers;
     for (auto& tensor : input) {
         buffers.push_back(torchTensor2Buffer(tensor));
     }
@@ -67,32 +68,36 @@ ReduceOp getReduceOp(c10d::ReduceOp reduce_op) {
 
 std::vector<torch::Tensor> RtpProcessGroup::all_reduce(std::vector<torch::Tensor>& input) {
     RTP_LLM_CHECK_WITH_INFO(input.size() == 1, "AllReduce input size must be 1 , but got %d", input.size());
-    auto     tensor      = input[0];
-    auto     dest_tensor = torch::empty_like(tensor);
-    ReduceOp reduce_op   = ReduceOp::Sum;
+    ScopedCUDAStreamContext stream_ctx(device_);
+    auto                    tensor      = input[0];
+    auto                    dest_tensor = torch::empty_like(tensor);
+    ReduceOp                reduce_op   = ReduceOp::Sum;
     device_->allReduce({torchTensor2Buffer(tensor), reduce_op, false, mode_, torchTensor2Buffer(dest_tensor)});
     check_cuda_error();
     return {dest_tensor};
 }
 
 void RtpProcessGroup::send(std::vector<torch::Tensor>& input, int dst_rank) {
     RTP_LLM_CHECK_WITH_INFO(input.size() == 1, "Send input size must be 1 , but got %d", input.size());
-    BatchSendRecvParams params;
+    ScopedCUDAStreamContext stream_ctx(device_);
+    BatchSendRecvParams     params;
     params.p2p_params.push_back({SendRecvType::kSend, torchTensor2Buffer(input[0]), dst_rank});
     device_->batchSendRecv(params, mode_);
     check_cuda_error();
 }
 
 void RtpProcessGroup::recv(std::vector<torch::Tensor>& input, int src_rank) {
     RTP_LLM_CHECK_WITH_INFO(input.size() == 1, "Send input size must be 1 , but got %d", input.size());
-    BatchSendRecvParams params;
+    ScopedCUDAStreamContext stream_ctx(device_);
+    BatchSendRecvParams     params;
     params.p2p_params.push_back({SendRecvType::kRecv, torchTensor2Buffer(input[0]), src_rank});
     device_->batchSendRecv(params, mode_);
     check_cuda_error();
 }
 
 std::vector<torch::Tensor> RtpProcessGroup::all_gather(std::vector<torch::Tensor>& input) {
     RTP_LLM_CHECK_WITH_INFO(input.size() == 1, "AllGather input size must be 1 , but got %d", input.size());
+    ScopedCUDAStreamContext stream_ctx(device_);
     auto output = torch::empty({input[0].size(0), input[0].size(1) * world_size_}, input[0].options());
     device_->allGather({{torchTensor2Buffer(output)}, mode_, {torchTensor2Buffer(input[0])}, false});
     check_cuda_error();

rtp_llm/models_py/bindings/common/RtpProcessGroup.h

@@ -2,6 +2,7 @@
 
 #include <torch/csrc/distributed/c10d/Types.hpp>
 #include <torch/csrc/distributed/c10d/Backend.hpp>
+#include <ATen/cuda/CUDAContext.h>
 #include <vector>
 #include "rtp_llm/cpp/core/Types.h"                       // for ParallelMode
 #include "rtp_llm/models_py/bindings/common/Torch_ext.h"  // for DefaultDeviceType
@@ -15,6 +16,28 @@ enum class RtpProcessGroupType {
     CP_GROUP        = 3,
 };
 
+class ScopedCUDAStreamContext {
+public:
+    explicit ScopedCUDAStreamContext(DefaultDeviceType* device): device_(device) {
+        original_stream_ = device_->getStream();
+        current_stream_  = at::cuda::getCurrentCUDAStream(at::cuda::current_device()).stream();
+        device_->setStream(current_stream_);
+    }
+
+    ~ScopedCUDAStreamContext() {
+        device_->setStream(original_stream_);
+    }
+    ScopedCUDAStreamContext(const ScopedCUDAStreamContext&)            = delete;
+    ScopedCUDAStreamContext& operator=(const ScopedCUDAStreamContext&) = delete;
+    ScopedCUDAStreamContext(ScopedCUDAStreamContext&&)                 = delete;
+    ScopedCUDAStreamContext& operator=(ScopedCUDAStreamContext&&)      = delete;
+
+private:
+    DefaultDeviceType* device_;
+    cudaStream_t       original_stream_;
+    cudaStream_t       current_stream_;
+};
+
 class RtpProcessGroup {
 public:
     RtpProcessGroup(RtpProcessGroupType type);