How to use cudaSetDevice to select devices while using MIG on A100?

Accelerated Computing CUDA CUDA Programming and Performance
1

Hi,

I am using A100 with Multi-Instance GPU (MIG). As described in the official document, it is straightforward to use environment variable CUDA_VISIBLE_DEVICES to select devices to run a single process. I am wondering if I could use CUDA runtime API cudaSetDevice or similar APIs to select devices for a thread in a multi-threading inference with CUDA? Thanks.

$ nvidia-smi -L
GPU 0: A100-SXM4-40GB (UUID: GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0)
  MIG 3g.20gb Device 0: (UUID: MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/1/0)
  MIG 3g.20gb Device 1: (UUID: MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/2/0)


$ CUDA_VISIBLE_DEVICES=MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/1/0 ./BlackScholes &
$ CUDA_VISIBLE_DEVICES=MIG-GPU-e86cb44c-6756-fd30-cd4a-1e6da3caf9b0/2/0 ./BlackScholes &
2

Currently, no. See here

With CUDA 11, only enumeration of a single MIG instance is supported.

CUDA applications treat a CI and its parent GI as a single CUDA device. CUDA is limited to use a single CI and will pick the first one available if several of them are visible. To summarize, there are two constraints:

CUDA can only enumerate a single compute instance

CUDA will not enumerate non-MIG GPU if any compute instance is enumerated on any other GPU

Note that these constraints may be relaxed in future NVIDIA driver releases for MIG.

Perhaps you may be interested in Triton inference server.