Selecting the Right GPU

If the DGX Station comes with a dedicated GPU in addition to the GB300, applications must select the appropriate GPU for each workload: AI workloads and graphics workloads.

The following code snippet shows how an application can identify each device.

#include <cuda_runtime.h>
#include <nvml.h>
#include <cstdio>

int main(void)
{
    int device_count = 0, graphics_device = -1, compute_device = -1;

    if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) {
        printf("No CUDA-capable GPU detected.\n");
        return 0;
    }

    nvmlInit_v2();
    for (int dev = 0; dev < device_count; ++dev) {
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        printf("Device %d: %s (cc %d.%d)\n",
            dev, prop.name, prop.major, prop.minor);

        // Match to NVML by PCI bus id -- CUDA and NVML order devices
        // differently, so the bus id is the stable common key.
        char busid[32];
        snprintf(busid, sizeof(busid), "%08x:%02x:%02x.0",
                prop.pciDomainID, prop.pciBusID, prop.pciDeviceID);

        nvmlDevice_t ndev;
        nvmlDeviceGetHandleByPciBusId_v2(busid, &ndev);

        nvmlEnableState_t mode;
        nvmlDeviceGetDisplayMode(ndev, &mode);
        if (mode == NVML_FEATURE_ENABLED)
            graphics_device = dev;
        else
            compute_device = dev;
    }
    nvmlShutdown();

    printf("Graphics Device: %d\n", graphics_device);
    printf("Compute Device: %d\n", compute_device);

    return 0;
}