[Code scan] Preserve CUDA execution context in edge_index CUDA path

[njzjz]

This issue was found by a Codex global scan of the repository at commit 19f9265.

The CUDA edge-index implementation switches the current device and does not restore it:

deepmd-gnn/op/edge_index_cuda.cu

Lines 132 to 134 in 19f9265

	if (device_index >= 0) {
	check_cuda(cudaSetDevice(device_index), "cudaSetDevice");
	}

It also launches kernels and Thrust scans on the implicit default stream rather than PyTorch's current stream:

deepmd-gnn/op/edge_index_cuda.cu

Lines 147 to 153 in 19f9265

	mark_edges_kernel<<<blocks, threads>>>(nlist, atype, mm, flags.ptr,
	total_slots, nloc, nnei, nall, nmm);
	check_cuda(cudaGetLastError(), "mark_edges_kernel launch");
	thrust::inclusive_scan(thrust::device, thrust::device_pointer_cast(flags.ptr),
	thrust::device_pointer_cast(flags.ptr + total_slots),
	thrust::device_pointer_cast(prefix.ptr));

deepmd-gnn/op/edge_index_cuda.cu

Lines 161 to 165 in 19f9265

	if (edge_count > 0) {
	scatter_edges_kernel<<<blocks, threads>>>(nlist, flags.ptr, prefix.ptr,
	edge_index, total_slots, nloc,
	nnei, nall);
	check_cuda(cudaGetLastError(), "scatter_edges_kernel launch");

In multi-GPU or custom-stream PyTorch programs, this can leave the calling thread on a different CUDA device and can introduce unexpected synchronization/order behavior.

Suggested fix: add an RAII device guard around cudaSetDevice, and use the current PyTorch CUDA stream for kernel launches and thrust::cuda::par.on(stream) if that is compatible with the stable ABI/linking constraints. If stable ABI constraints prevent using PyTorch stream APIs, document the default-stream behavior explicitly.