[Metax] Assign data stream to CUDA (#2173)

metax666 · web-flow · commit a5c6ea46050e · 2025-11-17T11:29:51.000+08:00
diff --git a/.github/workflows/_Metax-X86.yaml b/.github/workflows/_Metax-X86.yaml
@@ -78,7 +78,7 @@ jobs:
 
         run: |
           cd backends/metax_gpu/tests
-          bash run_test.sh -j 8
+          bash run_test.sh -j 32
 
       - name: push whl
         env:
diff --git a/.github/workflows/_Metax_work_private.yaml b/.github/workflows/_Metax_work_private.yaml
@@ -79,7 +79,7 @@ jobs:
 
         run: |
           cd backends/metax_gpu/tests
-          bash run_test.sh -j 8
+          bash run_test.sh -j 32
 
       - name: push whl
         env:
diff --git a/backends/metax_gpu/runtime/runtime.cc b/backends/metax_gpu/runtime/runtime.cc
@@ -579,7 +579,8 @@ C_Status AsyncMemCpyH2D(const C_Device device,
     return C_ERROR;
   }
 
-  cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyHostToDevice);
+  cudaErr = cudaMemcpyAsync(
+      dst, src, size, cudaMemcpyHostToDevice, (cudaStream_t)stream);
   if (cudaErr != cudaSuccess) {
     return C_ERROR;
   }
@@ -605,7 +606,8 @@ C_Status AsyncMemCpyD2H(const C_Device device,
     return C_ERROR;
   }
 
-  cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToHost);
+  cudaErr = cudaMemcpyAsync(
+      dst, src, size, cudaMemcpyDeviceToHost, (cudaStream_t)stream);
   if (cudaErr != cudaSuccess) {
     return C_ERROR;
   }
@@ -633,7 +635,8 @@ C_Status AsyncMemCpyD2D(const C_Device device,
     return C_ERROR;
   }
 
-  cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToDevice);
+  cudaErr = cudaMemcpyAsync(
+      dst, src, size, cudaMemcpyDeviceToDevice, (cudaStream_t)stream);
   if (cudaErr != cudaSuccess) {
     return C_ERROR;
   }

Original file line number	Diff line number	Diff line change
`@@ -579,7 +579,8 @@ C_Status AsyncMemCpyH2D(const C_Device device,`
`579`	`579`	`return C_ERROR;`
`580`	`580`	`}`
`581`	`581`
`582`		`- cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyHostToDevice);`
	`582`	`+ cudaErr = cudaMemcpyAsync(`
	`583`	`+ dst, src, size, cudaMemcpyHostToDevice, (cudaStream_t)stream);`
`583`	`584`	`if (cudaErr != cudaSuccess) {`
`584`	`585`	`return C_ERROR;`
`585`	`586`	`}`
`@@ -605,7 +606,8 @@ C_Status AsyncMemCpyD2H(const C_Device device,`
`605`	`606`	`return C_ERROR;`
`606`	`607`	`}`
`607`	`608`
`608`		`- cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToHost);`
	`609`	`+ cudaErr = cudaMemcpyAsync(`
	`610`	`+ dst, src, size, cudaMemcpyDeviceToHost, (cudaStream_t)stream);`
`609`	`611`	`if (cudaErr != cudaSuccess) {`
`610`	`612`	`return C_ERROR;`
`611`	`613`	`}`
`@@ -633,7 +635,8 @@ C_Status AsyncMemCpyD2D(const C_Device device,`
`633`	`635`	`return C_ERROR;`
`634`	`636`	`}`
`635`	`637`
`636`		`- cudaErr = cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToDevice);`
	`638`	`+ cudaErr = cudaMemcpyAsync(`
	`639`	`+ dst, src, size, cudaMemcpyDeviceToDevice, (cudaStream_t)stream);`
`637`	`640`	`if (cudaErr != cudaSuccess) {`
`638`	`641`	`return C_ERROR;`
`639`	`642`	`}`