Skip to content

Commit cee6d1e

Browse files
committed
Revert "simplify code through directly using the comm stream without converting it to a CUDAStream and back to a phi::Stream. This change is made in the deep_ep.cpp file, specifically in the get_comm_stream function. The modified line now directly retrieves the comm stream as a cudaStream_t and constructs a phi::Stream from it, eliminating unnecessary conversions."
This reverts commit f6f5d0d.
1 parent fc00d22 commit cee6d1e

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

csrc/deep_ep.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1900,7 +1900,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
19001900
.def("get_comm_stream",
19011901
[](deep_ep::Buffer &self) {
19021902
int device_id = self.get_local_device_id();
1903-
cudaStream_t comm_stream = self.get_comm_stream().stream();
1903+
cudaStream_t comm_stream = at::cuda::CUDAStream(self.get_comm_stream()).stream();
19041904
auto s = phi::Stream(reinterpret_cast<phi::StreamId>(comm_stream));
19051905
#if defined(PADDLE_WITH_CUDA)
19061906
return phi::CUDAStream(phi::GPUPlace(device_id), s);

0 commit comments

Comments
 (0)