fix ensure device index bug

cuda/cuda_stream_manager.cpp

@@ -57,6 +57,10 @@ void CudaStreamManager::sync(int idx) {
 }
 
 void CudaStreamManager::setup(const int device) {
+#ifdef MOE_USE_NCCL
+	this->ncclgood = 0;
+#endif
+	this->device = device;
 	checkCudaErrors(cudaSetDevice(device));
 	streams = new cudaStream_t[SMGR_N_STREAMS];
 	handles = new cublasHandle_t[SMGR_N_STREAMS];

cuda/cuda_stream_manager.h

@@ -29,7 +29,7 @@ public:
 #endif
 
 public:
-    CudaStreamManager(int device_): device(device_), ncclgood(0) {
+    CudaStreamManager(int device_): device(device_) {
 		this->setup(device);
     }
 

cuda/moe_comm_kernel.cu

@@ -197,7 +197,7 @@ std::vector<torch::Tensor> moe_cuda_global_gather(
 }
 
 void moe_ensure_nccl(c10d::ProcessGroupNCCL& p, torch::Tensor t) {
-	auto smgr = getCudaStreamManager(0);
+	auto smgr = getCudaStreamManager(t.device().index());
 	smgr->ensure((void*)&p, t.device());
 }