From facdd253a078531876ac133e342de6e2df75cf42 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Sat, 16 May 2026 15:39:34 -0700 Subject: [PATCH] revamp GPU discovery Scanning the output of llama-server is turning out to be too error prone across llama.cpp updates, so this switches to a thin dynamic library load against the bundled GGML libraries so more details can be gathered from the API. --- cmd/cmd.go | 12 + discover/amd.go | 48 +++ discover/amd_test.go | 30 ++ discover/cuda_compat.go | 49 +-- discover/llama_server.go | 63 +++- discover/native_probe.go | 263 ++++++++++++++ discover/native_probe_linux.go | 503 +++++++++++++++++++++++++++ discover/native_probe_linux_nocgo.go | 12 + discover/native_probe_stub.go | 12 + discover/native_probe_test.go | 203 +++++++++++ discover/native_probe_windows.go | 470 +++++++++++++++++++++++++ discover/runner.go | 129 +++++++ discover/runner_test.go | 144 ++++++++ llm/llama_server.go | 2 +- llm/rocm_default.go | 7 + llm/rocm_windows.go | 282 +++++++++++++++ logutil/logutil.go | 5 +- logutil/logutil_test.go | 19 + ml/device.go | 68 +++- 19 files changed, 2274 insertions(+), 47 deletions(-) create mode 100644 discover/native_probe.go create mode 100644 discover/native_probe_linux.go create mode 100644 discover/native_probe_linux_nocgo.go create mode 100644 discover/native_probe_stub.go create mode 100644 discover/native_probe_test.go create mode 100644 discover/native_probe_windows.go create mode 100644 llm/rocm_default.go create mode 100644 llm/rocm_windows.go create mode 100644 logutil/logutil_test.go diff --git a/cmd/cmd.go b/cmd/cmd.go index dc3188000..79f035fa2 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -42,6 +42,7 @@ import ( "github.com/ollama/ollama/cmd/config" "github.com/ollama/ollama/cmd/launch" "github.com/ollama/ollama/cmd/tui" + "github.com/ollama/ollama/discover" "github.com/ollama/ollama/envconfig" "github.com/ollama/ollama/format" "github.com/ollama/ollama/internal/modelref" @@ -2503,6 +2504,16 @@ func NewCLI() *cobra.Command { _ = runner.Execute(args[1:]) }) + var gpuDiscoverLibDirs []string + gpuDiscoverCmd := &cobra.Command{ + Use: "gpu-discover", + Hidden: true, + RunE: func(cmd *cobra.Command, _ []string) error { + return discover.RunNativeProbeCommand(cmd.Context(), gpuDiscoverLibDirs, os.Stdout) + }, + } + gpuDiscoverCmd.Flags().StringArrayVar(&gpuDiscoverLibDirs, "lib-dir", nil, "Ollama runtime library directory") + envVars := envconfig.AsMap() envs := []envconfig.EnvVar{envVars["OLLAMA_HOST"]} @@ -2569,6 +2580,7 @@ func NewCLI() *cobra.Command { copyCmd, deleteCmd, runnerCmd, + gpuDiscoverCmd, launch.LaunchCmd(checkServerHeartbeat, runInteractiveTUI), ) diff --git a/discover/amd.go b/discover/amd.go index 967805258..77f4eb4dc 100644 --- a/discover/amd.go +++ b/discover/amd.go @@ -65,6 +65,52 @@ func parseGFXTarget(gfx string) (int, int) { return int(major), int(minor) } +// HSA_OVERRIDE_GFX_VERSION changes the effective HIP/rocBLAS target even +// though KFD/sysfs still reports the physical ASIC. +func hsaOverrideGFXTarget() string { + return rocmGFXTargetOverride(os.Getenv("HSA_OVERRIDE_GFX_VERSION")) +} + +func rocmGFXTargetOverride(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + if strings.HasPrefix(value, "gfx") { + if major, minor := parseGFXTarget(value); major != 0 || minor != 0 { + return value + } + return "" + } + + parts := strings.Split(value, ".") + if len(parts) != 3 { + return "" + } + + var digits [3]uint64 + for i, part := range parts { + digit, err := strconv.ParseUint(part, 10, 8) + if err != nil || digit > 0xf { + return "" + } + digits[i] = digit + } + + return "gfx" + + strconv.FormatUint(digits[0], 10) + + strconv.FormatUint(digits[1], 16) + + strconv.FormatUint(digits[2], 16) +} + +func setROCmGFXTarget(device *ml.DeviceInfo, gfx string) { + if gfx == "" || device.Library != "ROCm" { + return + } + device.GFXTarget = gfx + device.ComputeMajor, device.ComputeMinor = parseGFXTarget(gfx) +} + // rocblasGFXTargets scans the rocblas library directory for supported gfx targets // by looking for TensileLibrary_lazy_gfxNNNN.dat files. func rocblasGFXTargets(libDirs []string) map[string]bool { @@ -338,12 +384,14 @@ func filterUnsupportedROCmDevices(devices []ml.DeviceInfo, libDirs []string) []m return devices } + override := hsaOverrideGFXTarget() var filtered []ml.DeviceInfo for _, dev := range devices { if dev.Library != "ROCm" { filtered = append(filtered, dev) continue } + setROCmGFXTarget(&dev, override) gfx := dev.GFXTarget if gfx == "" { filtered = append(filtered, dev) diff --git a/discover/amd_test.go b/discover/amd_test.go index e2db8caca..cbdbaa8dc 100644 --- a/discover/amd_test.go +++ b/discover/amd_test.go @@ -112,6 +112,36 @@ func TestApplyLinuxROCmRefinement(t *testing.T) { } } +func TestFilterUnsupportedROCmDevicesRespectsHSAOverride(t *testing.T) { + t.Setenv("HSA_OVERRIDE_GFX_VERSION", "10.3.0") + + libDir := t.TempDir() + rocblasDir := filepath.Join(libDir, "rocblas", "library") + if err := os.MkdirAll(rocblasDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(rocblasDir, "TensileLibrary_lazy_gfx1030.dat"), nil, 0o644); err != nil { + t.Fatal(err) + } + + devices := filterUnsupportedROCmDevices([]ml.DeviceInfo{{ + DeviceID: ml.DeviceID{ID: "0", Library: "ROCm"}, + Name: "ROCm0", + GFXTarget: "gfx1031", + ComputeMajor: 0x10, + ComputeMinor: 0x31, + }}, []string{libDir}) + if len(devices) != 1 { + t.Fatalf("got %d devices, want 1", len(devices)) + } + if got := devices[0].GFXTarget; got != "gfx1030" { + t.Fatalf("GFXTarget = %q, want gfx1030", got) + } + if got := devices[0].Compute(); got != "gfx1030" { + t.Fatalf("Compute() = %q, want gfx1030", got) + } +} + type fakeROCmNode struct { node int renderMinor int diff --git a/discover/cuda_compat.go b/discover/cuda_compat.go index 8855d8c1e..2dd16f7c8 100644 --- a/discover/cuda_compat.go +++ b/discover/cuda_compat.go @@ -3,21 +3,11 @@ package discover import ( "context" "log/slog" - "os/exec" - "strconv" - "strings" - "sync" - "time" "github.com/ollama/ollama/ml" ) -var ( - cudaDriverVersionMu sync.Mutex - cudaDriverMajorVersion *int -) - -func filterOldCUDADriver(ctx context.Context, devices []ml.DeviceInfo) []ml.DeviceInfo { +func filterOldCUDADriver(_ context.Context, devices []ml.DeviceInfo) []ml.DeviceInfo { oldCUDA := func(dev ml.DeviceInfo) bool { return dev.Library == "CUDA" && dev.ComputeMajor > 0 && dev.ComputeMajor < 7 } @@ -33,9 +23,9 @@ func filterOldCUDADriver(ctx context.Context, devices []ml.DeviceInfo) []ml.Devi return devices } - driver, err := nvidiaDriverMajorVersion(ctx) - if err != nil { - slog.Warn("could not run nvidia-smi to verify CUDA driver compatibility for an older NVIDIA GPU", "error", err) + driver := nvidiaDriverMajorFromDevices(devices) + if driver == 0 { + slog.Warn("could not verify NVIDIA driver compatibility for an older NVIDIA GPU") return devices } if driver >= 570 { @@ -54,30 +44,11 @@ func filterOldCUDADriver(ctx context.Context, devices []ml.DeviceInfo) []ml.Devi return filtered } -func nvidiaDriverMajorVersion(ctx context.Context) (int, error) { - cudaDriverVersionMu.Lock() - defer cudaDriverVersionMu.Unlock() - - if cudaDriverMajorVersion != nil { - return *cudaDriverMajorVersion, nil +func nvidiaDriverMajorFromDevices(devices []ml.DeviceInfo) int { + for _, dev := range devices { + if dev.Library == "CUDA" && dev.NVIDIADriverMajor > 0 { + return dev.NVIDIADriverMajor + } } - - ctx, cancel := context.WithTimeout(ctx, 3*time.Second) - defer cancel() - - output, err := exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits").Output() - if err != nil { - return 0, err - } - - line := strings.TrimSpace(strings.Split(string(output), "\n")[0]) - major, _, _ := strings.Cut(line, ".") - driver, err := strconv.Atoi(major) - if err != nil { - slog.Warn("could not parse nvidia-smi driver version for an older NVIDIA GPU", "version", line, "error", err) - return 0, err - } - - cudaDriverMajorVersion = &driver - return driver, nil + return 0 } diff --git a/discover/llama_server.go b/discover/llama_server.go index a17e203d4..cafb5ff9c 100644 --- a/discover/llama_server.go +++ b/discover/llama_server.go @@ -148,8 +148,14 @@ func llamaServerDiscoverDevices(ctx context.Context, libDirs []string, extraEnvs return nil, status, fmt.Errorf("llama-server --list-devices failed: %w", err) } - combined := string(listOutput) + "\n" + strings.Join(stderrLines, "\n") - return parseLlamaServerDevices(combined, libDirs), status, nil + nativeDevices, nativeStderr, nativeErr := discoverNativeDevices(ctx, llamaServer, libDirs, extraEnvs) + _, _ = status.Write([]byte(nativeStderr)) + if nativeErr != nil { + logNativeProbeFailure(nativeErr, nativeStderr, libDirs) + } + + combined := string(listOutput) + "\n" + strings.Join(stderrLines, "\n") + "\n" + nativeStderr + return parseLlamaServerDevicesWithNative(combined, libDirs, nativeDevices), status, nil } func llamaServerDiscoveryOutput(ctx context.Context) io.Writer { @@ -191,11 +197,24 @@ var ( // It extracts device info, ROCm gfx targets, CUDA compute capabilities, and // CUDA compiled architecture lists. func parseLlamaServerDevices(output string, libDirs []string) []ml.DeviceInfo { + return parseLlamaServerDevicesWithNative(output, libDirs, nil) +} + +func parseLlamaServerDevicesWithNative(output string, libDirs []string, nativeDevices []nativeProbeDevice) []ml.DeviceInfo { // Extract per-device metadata from stderr gfxByIndex := parseROCmGFXTargets(output) + rocmGFXOverride := hsaOverrideGFXTarget() integratedByIndex := parseVulkanUMA(output) ccByIndex := make(map[int]cudaComputeCapability) var cudaArchs []string // compiled architectures for this variant + nativeByIndex := nativeProbeByLibraryIndex(nativeDevices) + for idx, dev := range nativeByIndex["ROCm"] { + if rocmGFXOverride != "" { + gfxByIndex[idx] = rocmGFXOverride + } else if dev.GFXTarget != "" { + gfxByIndex[idx] = dev.GFXTarget + } + } scanner := bufio.NewScanner(strings.NewReader(output)) for scanner.Scan() { @@ -214,6 +233,18 @@ func parseLlamaServerDevices(output string, libDirs []string) []ml.DeviceInfo { cudaArchs = strings.Split(matches[1], ",") } } + if cudaDevices := nativeByIndex["CUDA"]; len(cudaDevices) > 0 { + for idx, dev := range cudaDevices { + if dev.ComputeMajor <= 0 { + continue + } + ccByIndex[idx] = cudaComputeCapability{ + major: dev.ComputeMajor, + minor: dev.ComputeMinor, + arch: fmt.Sprintf("%d%d0", dev.ComputeMajor, dev.ComputeMinor), + } + } + } // Validate CUDA devices against compiled architectures cudaArchSet := make(map[string]bool, len(cudaArchs)) @@ -272,6 +303,7 @@ func parseLlamaServerDevices(output string, libDirs []string) []ml.DeviceInfo { } computeMajor, computeMinor := computeVersion(library, deviceIndex, gfxByIndex, ccByIndex) + nativeDevice, hasNativeDevice := nativeByIndex[library][deviceIndex] dev := ml.DeviceInfo{ DeviceID: ml.DeviceID{ ID: strconv.Itoa(deviceIndex), @@ -287,7 +319,30 @@ func parseLlamaServerDevices(output string, libDirs []string) []ml.DeviceInfo { GFXTarget: gfxByIndex[deviceIndex], Integrated: isIntegratedLlamaServerDevice(library, deviceIndex, integratedByIndex), } - if library == "CUDA" && hasCUDARuntime { + if hasNativeDevice { + if nativeDevice.DeviceID != "" { + dev.PCIID = nativeDevice.DeviceID + } + if nativeDevice.IntegratedKnown { + dev.Integrated = nativeDevice.Integrated + } else { + dev.Integrated = dev.Integrated || nativeDevice.Integrated + } + if dev.ComputeMajor == 0 && nativeDevice.ComputeMajor > 0 { + dev.ComputeMajor = nativeDevice.ComputeMajor + dev.ComputeMinor = nativeDevice.ComputeMinor + } + if nativeDevice.CUDADriverMajor > 0 { + dev.DriverMajor = nativeDevice.CUDADriverMajor + dev.DriverMinor = nativeDevice.CUDADriverMinor + } + if nativeDevice.NVIDIADriverMajor > 0 { + dev.NVIDIADriverMajor = nativeDevice.NVIDIADriverMajor + } + setROCmGFXTarget(&dev, nativeDevice.GFXTarget) + } + setROCmGFXTarget(&dev, rocmGFXOverride) + if library == "CUDA" && dev.DriverMajor == 0 && hasCUDARuntime { dev.DriverMajor = cudaRuntimeMajor dev.DriverMinor = cudaRuntimeMinor } @@ -384,7 +439,7 @@ func inferLibrary(name, description string) string { } func isIntegratedLlamaServerDevice(library string, deviceIndex int, integratedByIndex map[int]bool) bool { - if integratedByIndex[deviceIndex] { + if library == "Vulkan" && integratedByIndex[deviceIndex] { return true } diff --git a/discover/native_probe.go b/discover/native_probe.go new file mode 100644 index 000000000..015a2abc0 --- /dev/null +++ b/discover/native_probe.go @@ -0,0 +1,263 @@ +package discover + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "os" + "os/exec" + "runtime" + "strings" + "time" + + "github.com/ollama/ollama/llm" + "github.com/ollama/ollama/ml" +) + +// Native GPU discovery runs in a short-lived Ollama subprocess so loading GGML +// and driver libraries cannot crash the main server process. The subprocess +// keeps stdout reserved for JSON and lets GGML's default logger write to +// stderr; the parent captures that stderr for trace/debug diagnostics. +const nativeProbeTimeout = 15 * time.Second + +type nativeProbeDevice struct { + Library string `json:"library"` + Index int `json:"index"` + // IndexMatchesBackend means Index is in the same visible-device order that + // llama-server reports, so it is safe to correlate when PCI ID is missing. + IndexMatchesBackend bool `json:"index_matches_backend,omitempty"` + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + DeviceID string `json:"device_id,omitempty"` + Integrated bool `json:"integrated,omitempty"` + IntegratedKnown bool `json:"integrated_known"` + TotalMemory uint64 `json:"total_memory,omitempty"` + FreeMemory uint64 `json:"free_memory,omitempty"` + ComputeMajor int `json:"compute_major,omitempty"` + ComputeMinor int `json:"compute_minor,omitempty"` + CUDADriverMajor int `json:"cuda_driver_major,omitempty"` + CUDADriverMinor int `json:"cuda_driver_minor,omitempty"` + NVIDIADriverMajor int `json:"nvidia_driver_major,omitempty"` + GFXTarget string `json:"gfx_target,omitempty"` +} + +type nativeProbeResult struct { + Devices []nativeProbeDevice `json:"devices"` +} + +type ggmlBackendDevCaps struct { + Async uint8 + HostBuffer uint8 + BufferFromHostPtr uint8 + Events uint8 +} + +type ggmlBackendDevProps struct { + Name uintptr + Description uintptr + MemoryFree uintptr + MemoryTotal uintptr + Type int32 + _ [4]byte + DeviceID uintptr + Caps ggmlBackendDevCaps + _ [4]byte +} + +func discoverNativeDevices(ctx context.Context, llamaServer string, libDirs []string, extraEnvs map[string]string) ([]nativeProbeDevice, string, error) { + if runtime.GOOS != "linux" && runtime.GOOS != "windows" { + return nil, "", nil + } + + exe, err := os.Executable() + if err != nil { + return nil, "", err + } + + ctx, cancel := context.WithTimeout(ctx, nativeProbeTimeout) + defer cancel() + + args := []string{"gpu-discover"} + for _, dir := range libDirs { + args = append(args, "--lib-dir", dir) + } + cmd := exec.CommandContext(ctx, exe, args...) + cmd.WaitDelay = llamaServerDiscoveryWaitDelay + llm.SetupLlamaServerCommandEnv(cmd, llamaServer, libDirs, extraEnvs) + + var stderr bytes.Buffer + cmd.Stderr = &stderr + stdout, err := cmd.Output() + if err != nil { + if ctx.Err() != nil { + return nil, stderr.String(), ctx.Err() + } + return nil, stderr.String(), err + } + + var result nativeProbeResult + if err := json.Unmarshal(stdout, &result); err != nil { + return nil, stderr.String(), err + } + + return result.Devices, stderr.String(), nil +} + +func RunNativeProbeCommand(ctx context.Context, libDirs []string, out io.Writer) error { + if len(libDirs) == 0 { + libDirs = []string{ml.LibOllamaPath} + } + + devices, err := runNativeProbe(ctx, libDirs) + if err != nil { + return err + } + + return json.NewEncoder(out).Encode(nativeProbeResult{Devices: devices}) +} + +func runNativeProbe(ctx context.Context, libDirs []string) ([]nativeProbeDevice, error) { + return runPlatformNativeProbe(ctx, libDirs) +} + +func mergeNativeProbeDevices(base, supplement []nativeProbeDevice) []nativeProbeDevice { + if len(base) == 0 { + var out []nativeProbeDevice + for _, extra := range supplement { + if extra.IndexMatchesBackend { + out = append(out, extra) + } + } + return out + } + + out := append([]nativeProbeDevice(nil), base...) + for _, extra := range supplement { + idx := -1 + for i := range out { + if sameNativeProbeDevice(out[i], extra) { + idx = i + break + } + } + if idx < 0 { + if !extra.IndexMatchesBackend || nativeProbeLibraryIndexExists(out, extra) { + continue + } + out = append(out, extra) + continue + } + mergeNativeProbeDevice(&out[idx], extra) + } + return out +} + +func sameNativeProbeDevice(a, b nativeProbeDevice) bool { + if !strings.EqualFold(a.Library, b.Library) { + return false + } + if a.DeviceID != "" && b.DeviceID != "" { + return strings.EqualFold(a.DeviceID, b.DeviceID) + } + if !a.IndexMatchesBackend || !b.IndexMatchesBackend { + return false + } + return a.Index == b.Index +} + +func mergeNativeProbeDevice(dst *nativeProbeDevice, src nativeProbeDevice) { + dst.IndexMatchesBackend = dst.IndexMatchesBackend || src.IndexMatchesBackend + if dst.Name == "" { + dst.Name = src.Name + } + if dst.Description == "" { + dst.Description = src.Description + } + if dst.DeviceID == "" { + dst.DeviceID = src.DeviceID + } + if src.IntegratedKnown { + dst.Integrated = src.Integrated + dst.IntegratedKnown = true + } else if !dst.IntegratedKnown && src.Integrated { + dst.Integrated = true + } + if dst.TotalMemory == 0 { + dst.TotalMemory = src.TotalMemory + } + if dst.FreeMemory == 0 { + dst.FreeMemory = src.FreeMemory + } + if dst.ComputeMajor == 0 && src.ComputeMajor != 0 { + dst.ComputeMajor = src.ComputeMajor + dst.ComputeMinor = src.ComputeMinor + } + if dst.CUDADriverMajor == 0 && src.CUDADriverMajor != 0 { + dst.CUDADriverMajor = src.CUDADriverMajor + dst.CUDADriverMinor = src.CUDADriverMinor + } + if dst.NVIDIADriverMajor == 0 && src.NVIDIADriverMajor != 0 { + dst.NVIDIADriverMajor = src.NVIDIADriverMajor + } + if dst.GFXTarget == "" { + dst.GFXTarget = src.GFXTarget + } +} + +func nativeProbeLibraryIndexExists(devices []nativeProbeDevice, target nativeProbeDevice) bool { + if !target.IndexMatchesBackend { + return false + } + for _, dev := range devices { + if strings.EqualFold(dev.Library, target.Library) && dev.Index == target.Index { + return true + } + } + return false +} + +func nativeProbeByLibraryIndex(devices []nativeProbeDevice) map[string]map[int]nativeProbeDevice { + out := map[string]map[int]nativeProbeDevice{} + for _, dev := range devices { + if !dev.IndexMatchesBackend { + continue + } + lib := normalizeNativeProbeLibrary(dev.Library) + if lib == "" { + continue + } + if _, ok := out[lib]; !ok { + out[lib] = map[int]nativeProbeDevice{} + } + out[lib][dev.Index] = dev + } + return out +} + +func normalizeNativeProbeLibrary(library string) string { + switch strings.ToLower(library) { + case "cuda": + return "CUDA" + case "hip", "rocm": + return "ROCm" + case "vulkan": + return "Vulkan" + case "metal": + return "Metal" + default: + return library + } +} + +func logNativeProbeFailure(err error, stderr string, libDirs []string) { + if err == nil { + return + } + if stderr != "" { + slog.Debug("native GPU discovery failed", "error", err, "stderr", stderr, "libDirs", libDirs) + return + } + slog.Debug("native GPU discovery failed", "error", err, "libDirs", libDirs) +} diff --git a/discover/native_probe_linux.go b/discover/native_probe_linux.go new file mode 100644 index 000000000..587f389ac --- /dev/null +++ b/discover/native_probe_linux.go @@ -0,0 +1,503 @@ +//go:build linux + +package discover + +/* +#cgo linux LDFLAGS: -ldl + +#include +#include +#include +#include + +static void * ollama_dlopen(const char * path, int global) { + return dlopen(path, RTLD_NOW | (global ? RTLD_GLOBAL : RTLD_LOCAL)); +} + +static void * ollama_dlsym(void * handle, const char * name) { + return dlsym(handle, name); +} + +static const char * ollama_dlerror(void) { + const char * err = dlerror(); + return err ? err : ""; +} + +typedef void * (*ollama_ggml_backend_load_fn)(const char *); +typedef size_t (*ollama_ggml_backend_reg_dev_count_fn)(void *); +typedef void * (*ollama_ggml_backend_reg_dev_get_fn)(void *, size_t); +typedef const char * (*ollama_ggml_backend_reg_name_fn)(void *); +typedef void (*ollama_ggml_backend_dev_get_props_fn)(void *, void *); + +static void * ollama_call_ggml_backend_load(void * fn, const char * path) { + return ((ollama_ggml_backend_load_fn) fn)(path); +} + +static size_t ollama_call_ggml_backend_reg_dev_count(void * fn, void * reg) { + return ((ollama_ggml_backend_reg_dev_count_fn) fn)(reg); +} + +static void * ollama_call_ggml_backend_reg_dev_get(void * fn, void * reg, size_t index) { + return ((ollama_ggml_backend_reg_dev_get_fn) fn)(reg, index); +} + +static const char * ollama_call_ggml_backend_reg_name(void * fn, void * reg) { + return ((ollama_ggml_backend_reg_name_fn) fn)(reg); +} + +static void ollama_call_ggml_backend_dev_get_props(void * fn, void * dev, void * props) { + ((ollama_ggml_backend_dev_get_props_fn) fn)(dev, props); +} + +typedef int (*ollama_cu_init_fn)(unsigned int); +typedef int (*ollama_cu_driver_get_version_fn)(int *); +typedef int (*ollama_cu_device_get_count_fn)(int *); +typedef int (*ollama_cu_device_get_fn)(int *, int); +typedef int (*ollama_cu_device_get_attribute_fn)(int *, int, int); +typedef int (*ollama_cu_device_get_name_fn)(char *, int, int); +typedef int (*ollama_cu_device_total_mem_fn)(size_t *, int); +typedef int (*ollama_cu_device_get_pci_bus_id_fn)(char *, int, int); + +static int ollama_call_cu_init(void * fn) { + return ((ollama_cu_init_fn) fn)(0); +} + +static int ollama_call_cu_driver_get_version(void * fn, int * version) { + return ((ollama_cu_driver_get_version_fn) fn)(version); +} + +static int ollama_call_cu_device_get_count(void * fn, int * count) { + return ((ollama_cu_device_get_count_fn) fn)(count); +} + +static int ollama_call_cu_device_get(void * fn, int * device, int index) { + return ((ollama_cu_device_get_fn) fn)(device, index); +} + +static int ollama_call_cu_device_get_attribute(void * fn, int * value, int attr, int device) { + return ((ollama_cu_device_get_attribute_fn) fn)(value, attr, device); +} + +static int ollama_call_cu_device_get_name(void * fn, char * name, int len, int device) { + return ((ollama_cu_device_get_name_fn) fn)(name, len, device); +} + +static int ollama_call_cu_device_total_mem(void * fn, size_t * total, int device) { + return ((ollama_cu_device_total_mem_fn) fn)(total, device); +} + +static int ollama_call_cu_device_get_pci_bus_id(void * fn, char * pci, int len, int device) { + return ((ollama_cu_device_get_pci_bus_id_fn) fn)(pci, len, device); +} + +typedef int (*ollama_nvml_init_fn)(void); +typedef int (*ollama_nvml_shutdown_fn)(void); +typedef int (*ollama_nvml_system_get_driver_version_fn)(char *, unsigned int); + +static int ollama_call_nvml_init(void * fn) { + return ((ollama_nvml_init_fn) fn)(); +} + +static int ollama_call_nvml_shutdown(void * fn) { + return ((ollama_nvml_shutdown_fn) fn)(); +} + +static int ollama_call_nvml_system_get_driver_version(void * fn, char * version, unsigned int len) { + return ((ollama_nvml_system_get_driver_version_fn) fn)(version, len); +} + +*/ +import "C" + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "unsafe" +) + +const ( + cuSuccess = 0 + cuDeviceAttributeComputeCapabilityMajor = 75 + cuDeviceAttributeComputeCapabilityMinor = 76 + cuDeviceAttributeIntegrated = 18 +) + +type dlHandle struct { + ptr unsafe.Pointer +} + +func runPlatformNativeProbe(ctx context.Context, libDirs []string) ([]nativeProbeDevice, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + ggmlDevices, ggmlErr := probeGGMLDevicesLinux(libDirs) + var cudaDevices []nativeProbeDevice + var cudaErr error + if nativeProbeHasCUDA(libDirs) { + cudaDevices, cudaErr = probeCUDADriverLinux() + } + var rocmDevices []nativeProbeDevice + var rocmErr error + if nativeProbeHasROCm(libDirs) { + rocmDevices, rocmErr = probeROCmSysfsLinux() + } + + devices := mergeNativeProbeDevices(mergeNativeProbeDevices(ggmlDevices, cudaDevices), rocmDevices) + if len(devices) > 0 { + return devices, nil + } + + if ggmlErr != nil { + return nil, ggmlErr + } + if rocmErr != nil { + return nil, rocmErr + } + return nil, cudaErr +} + +func probeGGMLDevicesLinux(libDirs []string) ([]nativeProbeDevice, error) { + if len(libDirs) == 0 { + return nil, errors.New("no library directories provided") + } + + baseDir := libDirs[0] + if baseDir == "" { + return nil, errors.New("empty GGML library directory") + } + + base, err := dlopen(ggmlLibraryFile(baseDir, "ggml-base"), true) + if err != nil { + return nil, err + } + + ggml, err := dlopen(ggmlLibraryFile(baseDir, "ggml"), true) + if err != nil { + return nil, err + } + + backendLoad, err := dlsym(ggml, "ggml_backend_load") + if err != nil { + return nil, err + } + regDevCount, err := dlsym(base, "ggml_backend_reg_dev_count") + if err != nil { + return nil, err + } + regDevGet, err := dlsym(base, "ggml_backend_reg_dev_get") + if err != nil { + return nil, err + } + regName, err := dlsym(base, "ggml_backend_reg_name") + if err != nil { + return nil, err + } + devGetProps, err := dlsym(base, "ggml_backend_dev_get_props") + if err != nil { + return nil, err + } + + var devices []nativeProbeDevice + for _, backendPath := range nativeProbeBackendFiles(libDirs) { + reg := callGGMLBackendLoad(backendLoad, backendPath) + if reg == nil { + continue + } + + library := ggmlProbeLibraryName(callGGMLRegName(regName, reg)) + count := int(callGGMLRegDevCount(regDevCount, reg)) + for i := 0; i < count; i++ { + dev := callGGMLRegDevGet(regDevGet, reg, i) + if dev == nil { + continue + } + props := callGGMLDeviceProps(devGetProps, dev) + if props.MemoryTotal == 0 { + continue + } + devices = append(devices, nativeProbeDevice{ + Library: library, + Index: i, + IndexMatchesBackend: true, + Name: cString(props.Name), + Description: cString(props.Description), + DeviceID: cString(props.DeviceID), + Integrated: ggmlDeviceTypeIntegrated(props.Type), + IntegratedKnown: props.Type == ggmlBackendDeviceTypeGPU || + props.Type == ggmlBackendDeviceTypeIGPU, + TotalMemory: uint64(props.MemoryTotal), + FreeMemory: uint64(props.MemoryFree), + }) + } + } + + return devices, nil +} + +func probeCUDADriverLinux() ([]nativeProbeDevice, error) { + cuda, err := dlopenFirst([]string{"libcuda.so.1", "libcuda.so"}, false) + if err != nil { + return nil, err + } + + cuInit, err := dlsym(cuda, "cuInit") + if err != nil { + return nil, err + } + cuDriverGetVersion, err := dlsym(cuda, "cuDriverGetVersion") + if err != nil { + return nil, err + } + cuDeviceGetCount, err := dlsym(cuda, "cuDeviceGetCount") + if err != nil { + return nil, err + } + cuDeviceGet, err := dlsym(cuda, "cuDeviceGet") + if err != nil { + return nil, err + } + cuDeviceGetAttribute, err := dlsym(cuda, "cuDeviceGetAttribute") + if err != nil { + return nil, err + } + cuDeviceGetName, err := dlsym(cuda, "cuDeviceGetName") + if err != nil { + return nil, err + } + cuDeviceTotalMem, err := dlsymAny(cuda, "cuDeviceTotalMem_v2", "cuDeviceTotalMem") + if err != nil { + return nil, err + } + cuDeviceGetPCIBusID, _ := dlsym(cuda, "cuDeviceGetPCIBusId") + + if ret := C.ollama_call_cu_init(cuInit); ret != cuSuccess { + return nil, fmt.Errorf("cuInit failed: %d", int(ret)) + } + + var driverVersion C.int + driverMajor, driverMinor := 0, 0 + if ret := C.ollama_call_cu_driver_get_version(cuDriverGetVersion, &driverVersion); ret == cuSuccess { + version := int(driverVersion) + driverMajor = version / 1000 + driverMinor = (version - driverMajor*1000) / 10 + } + + nvidiaDriverMajor := 0 + if driver, err := probeNVIDIADriverMajorLinux(); err == nil { + nvidiaDriverMajor = driver + } + + var count C.int + if ret := C.ollama_call_cu_device_get_count(cuDeviceGetCount, &count); ret != cuSuccess { + return nil, fmt.Errorf("cuDeviceGetCount failed: %d", int(ret)) + } + + devices := make([]nativeProbeDevice, 0, int(count)) + for i := 0; i < int(count); i++ { + var device C.int + if ret := C.ollama_call_cu_device_get(cuDeviceGet, &device, C.int(i)); ret != cuSuccess { + continue + } + + major := cudaDeviceAttribute(cuDeviceGetAttribute, cuDeviceAttributeComputeCapabilityMajor, device) + minor := cudaDeviceAttribute(cuDeviceGetAttribute, cuDeviceAttributeComputeCapabilityMinor, device) + integrated := cudaDeviceAttribute(cuDeviceGetAttribute, cuDeviceAttributeIntegrated, device) == 1 + + var name [128]C.char + _ = C.ollama_call_cu_device_get_name(cuDeviceGetName, &name[0], C.int(len(name)), device) + + var total C.size_t + _ = C.ollama_call_cu_device_total_mem(cuDeviceTotalMem, &total, device) + + pci := "" + if cuDeviceGetPCIBusID != nil { + var pciBuf [32]C.char + if ret := C.ollama_call_cu_device_get_pci_bus_id(cuDeviceGetPCIBusID, &pciBuf[0], C.int(len(pciBuf)), device); ret == cuSuccess { + pci = strings.ToLower(C.GoString(&pciBuf[0])) + } + } + + devices = append(devices, nativeProbeDevice{ + Library: "CUDA", + Index: i, + IndexMatchesBackend: true, + Description: C.GoString(&name[0]), + DeviceID: pci, + Integrated: integrated, + IntegratedKnown: true, + TotalMemory: uint64(total), + ComputeMajor: major, + ComputeMinor: minor, + CUDADriverMajor: driverMajor, + CUDADriverMinor: driverMinor, + NVIDIADriverMajor: nvidiaDriverMajor, + }) + } + + return devices, nil +} + +func probeROCmSysfsLinux() ([]nativeProbeDevice, error) { + sysfsDevices, err := readROCmLinuxSysfsDevices("/sys") + if err != nil { + return nil, err + } + + override := hsaOverrideGFXTarget() + // Sysfs stays in physical KFD order; ROCm visibility envs can reindex the + // backend device list, so filtered sysfs data must merge by PCI ID only. + backendIndex := !rocmVisibleDevicesEnvSet() + devices := make([]nativeProbeDevice, 0, len(sysfsDevices)) + for i, sysfsDevice := range sysfsDevices { + gfxTarget := sysfsDevice.gfxTarget + if override != "" { + gfxTarget = override + } + devices = append(devices, nativeProbeDevice{ + Library: "ROCm", + Index: i, + IndexMatchesBackend: backendIndex, + DeviceID: sysfsDevice.pciID, + Integrated: sysfsDevice.integrated, + IntegratedKnown: sysfsDevice.known, + GFXTarget: gfxTarget, + }) + } + return devices, nil +} + +func rocmVisibleDevicesEnvSet() bool { + for _, name := range []string{"HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL"} { + if os.Getenv(name) != "" { + return true + } + } + return false +} + +func probeNVIDIADriverMajorLinux() (int, error) { + nvml, err := dlopenFirst([]string{"libnvidia-ml.so.1", "libnvidia-ml.so"}, false) + if err != nil { + return 0, err + } + initFn, err := dlsym(nvml, "nvmlInit_v2") + if err != nil { + return 0, err + } + shutdownFn, err := dlsym(nvml, "nvmlShutdown") + if err != nil { + return 0, err + } + driverFn, err := dlsym(nvml, "nvmlSystemGetDriverVersion") + if err != nil { + return 0, err + } + if ret := C.ollama_call_nvml_init(initFn); ret != 0 { + return 0, fmt.Errorf("nvmlInit_v2 failed: %d", int(ret)) + } + defer C.ollama_call_nvml_shutdown(shutdownFn) + + var version [80]C.char + if ret := C.ollama_call_nvml_system_get_driver_version(driverFn, &version[0], C.uint(len(version))); ret != 0 { + return 0, fmt.Errorf("nvmlSystemGetDriverVersion failed: %d", int(ret)) + } + return parseNVIDIADriverMajor(C.GoString(&version[0])) +} + +func cudaDeviceAttribute(fn unsafe.Pointer, attr int, device C.int) int { + var value C.int + if ret := C.ollama_call_cu_device_get_attribute(fn, &value, C.int(attr), device); ret != cuSuccess { + return 0 + } + return int(value) +} + +func dlopenFirst(names []string, global bool) (dlHandle, error) { + var errs []string + for _, name := range names { + handle, err := dlopen(name, global) + if err == nil { + return handle, nil + } + errs = append(errs, err.Error()) + } + return dlHandle{}, errors.New(strings.Join(errs, "; ")) +} + +func dlopen(path string, global bool) (dlHandle, error) { + cpath := C.CString(path) + defer C.free(unsafe.Pointer(cpath)) + + handle := C.ollama_dlopen(cpath, boolToCInt(global)) + if handle == nil { + return dlHandle{}, fmt.Errorf("dlopen %s: %s", path, C.GoString(C.ollama_dlerror())) + } + return dlHandle{ptr: handle}, nil +} + +func dlsym(handle dlHandle, name string) (unsafe.Pointer, error) { + cname := C.CString(name) + defer C.free(unsafe.Pointer(cname)) + + sym := C.ollama_dlsym(handle.ptr, cname) + if sym == nil { + return nil, fmt.Errorf("dlsym %s: %s", name, C.GoString(C.ollama_dlerror())) + } + return sym, nil +} + +func dlsymAny(handle dlHandle, names ...string) (unsafe.Pointer, error) { + var errs []string + for _, name := range names { + sym, err := dlsym(handle, name) + if err == nil { + return sym, nil + } + errs = append(errs, err.Error()) + } + return nil, errors.New(strings.Join(errs, "; ")) +} + +func callGGMLBackendLoad(fn unsafe.Pointer, path string) unsafe.Pointer { + cpath := C.CString(path) + defer C.free(unsafe.Pointer(cpath)) + return C.ollama_call_ggml_backend_load(fn, cpath) +} + +func callGGMLRegDevCount(fn unsafe.Pointer, reg unsafe.Pointer) uintptr { + return uintptr(C.ollama_call_ggml_backend_reg_dev_count(fn, reg)) +} + +func callGGMLRegDevGet(fn unsafe.Pointer, reg unsafe.Pointer, index int) unsafe.Pointer { + return C.ollama_call_ggml_backend_reg_dev_get(fn, reg, C.size_t(index)) +} + +func callGGMLRegName(fn unsafe.Pointer, reg unsafe.Pointer) string { + return C.GoString(C.ollama_call_ggml_backend_reg_name(fn, reg)) +} + +func callGGMLDeviceProps(fn unsafe.Pointer, dev unsafe.Pointer) ggmlBackendDevProps { + var props ggmlBackendDevProps + C.ollama_call_ggml_backend_dev_get_props(fn, dev, unsafe.Pointer(&props)) + return props +} + +func cString(ptr uintptr) string { + if ptr == 0 { + return "" + } + return C.GoString((*C.char)(unsafe.Pointer(ptr))) +} + +func boolToCInt(v bool) C.int { + if v { + return 1 + } + return 0 +} diff --git a/discover/native_probe_linux_nocgo.go b/discover/native_probe_linux_nocgo.go new file mode 100644 index 000000000..caf620c83 --- /dev/null +++ b/discover/native_probe_linux_nocgo.go @@ -0,0 +1,12 @@ +//go:build linux && !cgo + +package discover + +import ( + "context" + "errors" +) + +func runPlatformNativeProbe(context.Context, []string) ([]nativeProbeDevice, error) { + return nil, errors.New("native GPU discovery requires cgo on Linux") +} diff --git a/discover/native_probe_stub.go b/discover/native_probe_stub.go new file mode 100644 index 000000000..b7baae577 --- /dev/null +++ b/discover/native_probe_stub.go @@ -0,0 +1,12 @@ +//go:build !linux && !windows + +package discover + +import ( + "context" + "errors" +) + +func runPlatformNativeProbe(context.Context, []string) ([]nativeProbeDevice, error) { + return nil, errors.New("native GPU discovery is not implemented on this platform") +} diff --git a/discover/native_probe_test.go b/discover/native_probe_test.go new file mode 100644 index 000000000..582b3782b --- /dev/null +++ b/discover/native_probe_test.go @@ -0,0 +1,203 @@ +package discover + +import ( + "testing" + "unsafe" + + "github.com/ollama/ollama/ml" +) + +func TestGGMLBackendDevPropsLayout(t *testing.T) { + if unsafe.Sizeof(uintptr(0)) != 8 { + t.Skip("GGML probe layout assertions are for 64-bit builds") + } + + var props ggmlBackendDevProps + if got, want := unsafe.Sizeof(props), uintptr(56); got != want { + t.Fatalf("ggmlBackendDevProps size = %d, want %d", got, want) + } + checks := []struct { + name string + got uintptr + want uintptr + }{ + {"Name", unsafe.Offsetof(props.Name), 0}, + {"Description", unsafe.Offsetof(props.Description), 8}, + {"MemoryFree", unsafe.Offsetof(props.MemoryFree), 16}, + {"MemoryTotal", unsafe.Offsetof(props.MemoryTotal), 24}, + {"Type", unsafe.Offsetof(props.Type), 32}, + {"DeviceID", unsafe.Offsetof(props.DeviceID), 40}, + {"Caps", unsafe.Offsetof(props.Caps), 48}, + } + for _, tt := range checks { + t.Run(tt.name, func(t *testing.T) { + if tt.got != tt.want { + t.Fatalf("offset = %d, want %d", tt.got, tt.want) + } + }) + } + if got, want := unsafe.Sizeof(ggmlBackendDevCaps{}), uintptr(4); got != want { + t.Fatalf("ggmlBackendDevCaps size = %d, want %d", got, want) + } +} + +func TestParseLlamaServerDevicesUsesNativeCUDAComputeCapability(t *testing.T) { + output := `system_info: n_threads = 4 | CUDA : ARCHS = 750,800 | +Available devices: + CUDA0: NVIDIA GeForce GTX 1060 6GB (6063 MiB, 5900 MiB free) +` + devices := parseLlamaServerDevicesWithNative(output, []string{"/lib/ollama", "/lib/ollama/cuda_v13"}, []nativeProbeDevice{{ + Library: "CUDA", + Index: 0, + IndexMatchesBackend: true, + DeviceID: "0000:01:00.0", + ComputeMajor: 6, + ComputeMinor: 1, + CUDADriverMajor: 13, + NVIDIADriverMajor: 570, + }}) + if len(devices) != 0 { + t.Fatalf("got %d devices, want unsupported CUDA device filtered", len(devices)) + } + + output = `system_info: n_threads = 4 | CUDA : ARCHS = 610,750,800 | +Available devices: + CUDA0: NVIDIA GeForce GTX 1060 6GB (6063 MiB, 5900 MiB free) +` + devices = parseLlamaServerDevicesWithNative(output, []string{"/lib/ollama", "/lib/ollama/cuda_v12"}, []nativeProbeDevice{{ + Library: "CUDA", + Index: 0, + IndexMatchesBackend: true, + DeviceID: "0000:01:00.0", + ComputeMajor: 6, + ComputeMinor: 1, + CUDADriverMajor: 12, + NVIDIADriverMajor: 570, + }}) + if len(devices) != 1 { + t.Fatalf("got %d devices, want 1", len(devices)) + } + got := devices[0] + if got.Compute() != "6.1" { + t.Fatalf("compute = %q, want 6.1", got.Compute()) + } + if got.PCIID != "0000:01:00.0" { + t.Fatalf("PCIID = %q, want 0000:01:00.0", got.PCIID) + } + if got.Driver() != "12.0" { + t.Fatalf("driver = %q, want 12.0", got.Driver()) + } + if got.NVIDIADriverMajor != 570 { + t.Fatalf("NVIDIADriverMajor = %d, want 570", got.NVIDIADriverMajor) + } +} + +func TestParseLlamaServerDevicesUsesNativeROCmMetadata(t *testing.T) { + output := `ggml_vulkan: 0 = AMD Radeon RX 7600 | uma: 1 | fp16: 1 | +Available devices: + ROCm0: AMD Radeon RX 7600 (8176 MiB, 7900 MiB free) +` + devices := parseLlamaServerDevicesWithNative(output, []string{"/lib/ollama", "/lib/ollama/rocm"}, []nativeProbeDevice{{ + Library: "ROCm", + Index: 0, + IndexMatchesBackend: true, + DeviceID: "0000:03:00.0", + GFXTarget: "gfx1102", + Integrated: false, + IntegratedKnown: true, + }}) + if len(devices) != 1 { + t.Fatalf("got %d devices, want 1", len(devices)) + } + got := devices[0] + if got.PCIID != "0000:03:00.0" { + t.Fatalf("PCIID = %q, want 0000:03:00.0", got.PCIID) + } + if got.GFXTarget != "gfx1102" { + t.Fatalf("GFXTarget = %q, want gfx1102", got.GFXTarget) + } + if got.Compute() != "gfx1102" { + t.Fatalf("compute = %q, want gfx1102", got.Compute()) + } + if got.Integrated { + t.Fatalf("Integrated = true, want false") + } +} + +func TestNVIDIADriverMajorFromDevices(t *testing.T) { + devices := []ml.DeviceInfo{ + {DeviceID: ml.DeviceID{Library: "CUDA"}, NVIDIADriverMajor: 565}, + } + if got := nvidiaDriverMajorFromDevices(devices); got != 565 { + t.Fatalf("driver = %d, want 565", got) + } +} + +func TestMergeNativeProbeDevicesAvoidsUnreliableIndexMatch(t *testing.T) { + tests := []struct { + name string + base []nativeProbeDevice + supplement []nativeProbeDevice + wantLen int + wantPCI string + wantKnown bool + wantIGPU bool + }{ + { + name: "filtered sysfs cannot overwrite a different backend device by index", + base: []nativeProbeDevice{{ + Library: "ROCm", + Index: 0, + IndexMatchesBackend: true, + DeviceID: "0000:03:00.0", + }}, + supplement: []nativeProbeDevice{{ + Library: "ROCm", + Index: 0, + DeviceID: "0000:04:00.0", + Integrated: true, + IntegratedKnown: true, + }}, + wantLen: 1, + wantPCI: "0000:03:00.0", + }, + { + name: "filtered sysfs can still merge by PCI ID", + base: []nativeProbeDevice{{ + Library: "ROCm", + Index: 0, + IndexMatchesBackend: true, + DeviceID: "0000:04:00.0", + }}, + supplement: []nativeProbeDevice{{ + Library: "ROCm", + Index: 1, + DeviceID: "0000:04:00.0", + Integrated: true, + IntegratedKnown: true, + }}, + wantLen: 1, + wantPCI: "0000:04:00.0", + wantKnown: true, + wantIGPU: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := mergeNativeProbeDevices(tt.base, tt.supplement) + if len(got) != tt.wantLen { + t.Fatalf("got %d devices, want %d: %#v", len(got), tt.wantLen, got) + } + if got[0].DeviceID != tt.wantPCI { + t.Fatalf("DeviceID = %q, want %q", got[0].DeviceID, tt.wantPCI) + } + if got[0].IntegratedKnown != tt.wantKnown { + t.Fatalf("IntegratedKnown = %v, want %v", got[0].IntegratedKnown, tt.wantKnown) + } + if got[0].Integrated != tt.wantIGPU { + t.Fatalf("Integrated = %v, want %v", got[0].Integrated, tt.wantIGPU) + } + }) + } +} diff --git a/discover/native_probe_windows.go b/discover/native_probe_windows.go new file mode 100644 index 000000000..62c97e4ee --- /dev/null +++ b/discover/native_probe_windows.go @@ -0,0 +1,470 @@ +//go:build windows + +package discover + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "unsafe" + + "github.com/ollama/ollama/llm" + "golang.org/x/sys/windows" +) + +const ( + cuSuccessWindows = 0 + cuDeviceAttributeComputeCapabilityMajorW = 75 + cuDeviceAttributeComputeCapabilityMinorW = 76 + cuDeviceAttributeIntegratedW = 18 + hipSuccessWindows = 0 + hipDeviceAttributeIntegratedWindows = 16 +) + +func runPlatformNativeProbe(ctx context.Context, libDirs []string) ([]nativeProbeDevice, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + ggmlDevices, ggmlErr := probeGGMLDevicesWindows(libDirs) + var cudaDevices []nativeProbeDevice + var cudaErr error + if nativeProbeHasCUDA(libDirs) { + cudaDevices, cudaErr = probeCUDADriverWindows() + } + var rocmDevices []nativeProbeDevice + var rocmErr error + if nativeProbeHasROCm(libDirs) { + rocmDevices, rocmErr = probeHIPRuntimeWindows(libDirs) + } + + devices := mergeNativeProbeDevices(mergeNativeProbeDevices(ggmlDevices, cudaDevices), rocmDevices) + if len(devices) > 0 { + return devices, nil + } + + if ggmlErr != nil { + return nil, ggmlErr + } + if rocmErr != nil { + return nil, rocmErr + } + return nil, cudaErr +} + +func probeGGMLDevicesWindows(libDirs []string) ([]nativeProbeDevice, error) { + if len(libDirs) == 0 || libDirs[0] == "" { + return nil, errors.New("empty GGML library directory") + } + + base, err := loadDLLFromPath(ggmlLibraryFile(libDirs[0], "ggml-base")) + if err != nil { + return nil, err + } + ggml, err := loadDLLFromPath(ggmlLibraryFile(libDirs[0], "ggml")) + if err != nil { + return nil, err + } + + backendLoad, err := findProc(ggml, "ggml_backend_load") + if err != nil { + return nil, err + } + regDevCount, err := findProc(base, "ggml_backend_reg_dev_count") + if err != nil { + return nil, err + } + regDevGet, err := findProc(base, "ggml_backend_reg_dev_get") + if err != nil { + return nil, err + } + regName, err := findProc(base, "ggml_backend_reg_name") + if err != nil { + return nil, err + } + devGetProps, err := findProc(base, "ggml_backend_dev_get_props") + if err != nil { + return nil, err + } + + var devices []nativeProbeDevice + for _, backendPath := range nativeProbeBackendFiles(libDirs) { + cpath, err := windows.BytePtrFromString(backendPath) + if err != nil { + return nil, err + } + reg, _, _ := backendLoad.Call(uintptr(unsafe.Pointer(cpath))) + if reg == 0 { + continue + } + + regNamePtr, _, _ := regName.Call(reg) + library := ggmlProbeLibraryName(windowsCString(regNamePtr)) + count, _, _ := regDevCount.Call(reg) + for i := uintptr(0); i < count; i++ { + dev, _, _ := regDevGet.Call(reg, i) + if dev == 0 { + continue + } + var props ggmlBackendDevProps + devGetProps.Call(dev, uintptr(unsafe.Pointer(&props))) + if props.MemoryTotal == 0 { + continue + } + devices = append(devices, nativeProbeDevice{ + Library: library, + Index: int(i), + IndexMatchesBackend: true, + Name: windowsCString(props.Name), + Description: windowsCString(props.Description), + DeviceID: windowsCString(props.DeviceID), + Integrated: ggmlDeviceTypeIntegrated(props.Type), + IntegratedKnown: props.Type == ggmlBackendDeviceTypeGPU || + props.Type == ggmlBackendDeviceTypeIGPU, + TotalMemory: uint64(props.MemoryTotal), + FreeMemory: uint64(props.MemoryFree), + }) + } + } + + return devices, nil +} + +func probeCUDADriverWindows() ([]nativeProbeDevice, error) { + cuda, err := loadDLLFromSystem32("nvcuda.dll") + if err != nil { + return nil, err + } + cuInit, err := findProc(cuda, "cuInit") + if err != nil { + return nil, err + } + cuDriverGetVersion, err := findProc(cuda, "cuDriverGetVersion") + if err != nil { + return nil, err + } + cuDeviceGetCount, err := findProc(cuda, "cuDeviceGetCount") + if err != nil { + return nil, err + } + cuDeviceGet, err := findProc(cuda, "cuDeviceGet") + if err != nil { + return nil, err + } + cuDeviceGetAttribute, err := findProc(cuda, "cuDeviceGetAttribute") + if err != nil { + return nil, err + } + cuDeviceGetName, err := findProc(cuda, "cuDeviceGetName") + if err != nil { + return nil, err + } + cuDeviceTotalMem, err := procAny(cuda, "cuDeviceTotalMem_v2", "cuDeviceTotalMem") + if err != nil { + return nil, err + } + cuDeviceGetPCIBusID, _ := findProc(cuda, "cuDeviceGetPCIBusId") + + if ret, _, _ := cuInit.Call(0); ret != cuSuccessWindows { + return nil, fmt.Errorf("cuInit failed: %d", ret) + } + + driverMajor, driverMinor := 0, 0 + var driverVersion int32 + if ret, _, _ := cuDriverGetVersion.Call(uintptr(unsafe.Pointer(&driverVersion))); ret == cuSuccessWindows { + version := int(driverVersion) + driverMajor = version / 1000 + driverMinor = (version - driverMajor*1000) / 10 + } + + nvidiaDriverMajor := 0 + if driver, err := probeNVIDIADriverMajorWindows(); err == nil { + nvidiaDriverMajor = driver + } + + var count int32 + if ret, _, _ := cuDeviceGetCount.Call(uintptr(unsafe.Pointer(&count))); ret != cuSuccessWindows { + return nil, fmt.Errorf("cuDeviceGetCount failed: %d", ret) + } + + devices := make([]nativeProbeDevice, 0, int(count)) + for i := 0; i < int(count); i++ { + var device int32 + if ret, _, _ := cuDeviceGet.Call(uintptr(unsafe.Pointer(&device)), uintptr(i)); ret != cuSuccessWindows { + continue + } + + major := cudaDeviceAttributeWindows(cuDeviceGetAttribute, cuDeviceAttributeComputeCapabilityMajorW, device) + minor := cudaDeviceAttributeWindows(cuDeviceGetAttribute, cuDeviceAttributeComputeCapabilityMinorW, device) + integrated := cudaDeviceAttributeWindows(cuDeviceGetAttribute, cuDeviceAttributeIntegratedW, device) == 1 + + name := make([]byte, 128) + cuDeviceGetName.Call(uintptr(unsafe.Pointer(&name[0])), uintptr(len(name)), uintptr(device)) + + var total uintptr + cuDeviceTotalMem.Call(uintptr(unsafe.Pointer(&total)), uintptr(device)) + + pci := "" + if cuDeviceGetPCIBusID != nil { + pciBuf := make([]byte, 32) + if ret, _, _ := cuDeviceGetPCIBusID.Call(uintptr(unsafe.Pointer(&pciBuf[0])), uintptr(len(pciBuf)), uintptr(device)); ret == cuSuccessWindows { + pci = strings.ToLower(byteCString(pciBuf)) + } + } + + devices = append(devices, nativeProbeDevice{ + Library: "CUDA", + Index: i, + IndexMatchesBackend: true, + Description: byteCString(name), + DeviceID: pci, + Integrated: integrated, + IntegratedKnown: true, + TotalMemory: uint64(total), + ComputeMajor: major, + ComputeMinor: minor, + CUDADriverMajor: driverMajor, + CUDADriverMinor: driverMinor, + NVIDIADriverMajor: nvidiaDriverMajor, + }) + } + + return devices, nil +} + +func probeHIPRuntimeWindows(libDirs []string) ([]nativeProbeDevice, error) { + hipPath, err := llm.WindowsROCmRuntimeDLLPath(libDirs) + if err != nil { + return nil, err + } + hip, err := loadDLLFromPath(hipPath) + if err != nil { + return nil, err + } + hipGetDeviceCount, err := findProc(hip, "hipGetDeviceCount") + if err != nil { + return nil, err + } + hipDeviceGetName, err := findProc(hip, "hipDeviceGetName") + if err != nil { + return nil, err + } + hipDeviceTotalMem, err := findProc(hip, "hipDeviceTotalMem") + if err != nil { + return nil, err + } + hipDeviceGetPCIBusID, _ := findProc(hip, "hipDeviceGetPCIBusId") + hipDeviceGetAttribute, _ := findProc(hip, "hipDeviceGetAttribute") + + var count int32 + if ret, _, _ := hipGetDeviceCount.Call(uintptr(unsafe.Pointer(&count))); ret != hipSuccessWindows { + return nil, fmt.Errorf("hipGetDeviceCount failed: %d", ret) + } + + devices := make([]nativeProbeDevice, 0, int(count)) + for i := 0; i < int(count); i++ { + name := make([]byte, 128) + hipDeviceGetName.Call(uintptr(unsafe.Pointer(&name[0])), uintptr(len(name)), uintptr(i)) + + var total uintptr + hipDeviceTotalMem.Call(uintptr(unsafe.Pointer(&total)), uintptr(i)) + + pci := "" + if hipDeviceGetPCIBusID != nil { + pciBuf := make([]byte, 32) + if ret, _, _ := hipDeviceGetPCIBusID.Call(uintptr(unsafe.Pointer(&pciBuf[0])), uintptr(len(pciBuf)), uintptr(i)); ret == hipSuccessWindows { + pci = strings.ToLower(byteCString(pciBuf)) + } + } + + integrated, integratedKnown := false, false + if hipDeviceGetAttribute != nil { + integrated = hipDeviceAttributeWindows(hipDeviceGetAttribute, hipDeviceAttributeIntegratedWindows, int32(i)) == 1 + integratedKnown = true + } + + devices = append(devices, nativeProbeDevice{ + Library: "ROCm", + Index: i, + IndexMatchesBackend: true, + Description: byteCString(name), + DeviceID: pci, + Integrated: integrated, + IntegratedKnown: integratedKnown, + TotalMemory: uint64(total), + }) + } + + return devices, nil +} + +func probeNVIDIADriverMajorWindows() (int, error) { + nvml, err := loadDLLFromSystem32("nvml.dll") + if err != nil { + nvml, err = loadDLLFromDirs([]string{"nvml.dll"}, nvidiaNVMLDirsWindows()) + } + if err != nil { + return 0, err + } + initFn, err := findProc(nvml, "nvmlInit_v2") + if err != nil { + return 0, err + } + shutdownFn, err := findProc(nvml, "nvmlShutdown") + if err != nil { + return 0, err + } + driverFn, err := findProc(nvml, "nvmlSystemGetDriverVersion") + if err != nil { + return 0, err + } + if ret, _, _ := initFn.Call(); ret != 0 { + return 0, fmt.Errorf("nvmlInit_v2 failed: %d", ret) + } + defer shutdownFn.Call() + + version := make([]byte, 80) + if ret, _, _ := driverFn.Call(uintptr(unsafe.Pointer(&version[0])), uintptr(len(version))); ret != 0 { + return 0, fmt.Errorf("nvmlSystemGetDriverVersion failed: %d", ret) + } + return parseNVIDIADriverMajor(byteCString(version)) +} + +func cudaDeviceAttributeWindows(fn *windows.Proc, attr int, device int32) int { + var value int32 + if ret, _, _ := fn.Call(uintptr(unsafe.Pointer(&value)), uintptr(attr), uintptr(device)); ret != cuSuccessWindows { + return 0 + } + return int(value) +} + +func hipDeviceAttributeWindows(fn *windows.Proc, attr int, device int32) int { + var value int32 + if ret, _, _ := fn.Call(uintptr(unsafe.Pointer(&value)), uintptr(attr), uintptr(device)); ret != hipSuccessWindows { + return 0 + } + return int(value) +} + +func findProc(dll *windows.DLL, name string) (*windows.Proc, error) { + return dll.FindProc(name) +} + +// Use LoadLibraryEx so GPU discovery does not honor the current directory or PATH for DLL resolution. +func loadDLLFromSystem32(name string) (*windows.DLL, error) { + return loadDLLWithFlags(name, windows.LOAD_LIBRARY_SEARCH_SYSTEM32) +} + +func loadDLLFromPath(path string) (*windows.DLL, error) { + absPath, err := filepath.Abs(path) + if err != nil { + return nil, err + } + return loadDLLWithFlags(absPath, windows.LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR|windows.LOAD_LIBRARY_SEARCH_DEFAULT_DIRS) +} + +func loadDLLWithFlags(name string, flags uintptr) (*windows.DLL, error) { + handle, err := windows.LoadLibraryEx(name, 0, flags) + if err != nil { + return nil, fmt.Errorf("failed to load %s: %w", name, err) + } + return &windows.DLL{Name: name, Handle: handle}, nil +} + +func loadDLLFromDirs(names, dirs []string) (*windows.DLL, error) { + var errs []string + for _, name := range names { + for _, dir := range dirs { + path := filepath.Join(dir, name) + if _, err := os.Stat(path); err != nil { + continue + } + dll, err := loadDLLFromPath(path) + if err == nil { + return dll, nil + } + errs = append(errs, err.Error()) + } + } + if len(errs) == 0 { + return nil, fmt.Errorf("no matching DLL found: %s", strings.Join(names, ", ")) + } + return nil, errors.New(strings.Join(errs, "; ")) +} + +func nvidiaNVMLDirsWindows() []string { + var dirs []string + for _, root := range windowsProgramFilesDirs() { + dirs = append(dirs, filepath.Join(root, "NVIDIA Corporation", "NVSMI")) + } + return uniqueAbsDirs(dirs) +} + +func windowsProgramFilesDirs() []string { + return uniqueAbsDirs([]string{ + os.Getenv("ProgramW6432"), + os.Getenv("ProgramFiles"), + }) +} + +func uniqueAbsDirs(dirs []string) []string { + seen := map[string]bool{} + var out []string + for _, dir := range dirs { + if dir == "" { + continue + } + absDir, err := filepath.Abs(dir) + if err != nil { + continue + } + absDir = filepath.Clean(absDir) + key := strings.ToLower(absDir) + if seen[key] { + continue + } + seen[key] = true + out = append(out, absDir) + } + return out +} + +func procAny(dll *windows.DLL, names ...string) (*windows.Proc, error) { + var errs []string + for _, name := range names { + proc, err := dll.FindProc(name) + if err == nil { + return proc, nil + } + errs = append(errs, err.Error()) + } + return nil, errors.New(strings.Join(errs, "; ")) +} + +func windowsCString(ptr uintptr) string { + if ptr == 0 { + return "" + } + var bytes []byte + for p := ptr; ; p++ { + b := *(*byte)(unsafe.Pointer(p)) + if b == 0 { + return string(bytes) + } + bytes = append(bytes, b) + } +} + +func byteCString(data []byte) string { + for i, b := range data { + if b == 0 { + return string(data[:i]) + } + } + return string(data) +} diff --git a/discover/runner.go b/discover/runner.go index 0984965b3..aa65b2f86 100644 --- a/discover/runner.go +++ b/discover/runner.go @@ -186,6 +186,7 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml. devices[i].FilterID = devices[i].ID devices[i].ID = strconv.Itoa(postFilteredID[devices[i].Library]) } + remapFilterIDForUserVisibleDevices(&devices[i]) postFilteredID[devices[i].Library]++ } } @@ -405,6 +406,8 @@ func filterOverlapByLibrary(supported map[string]map[string]map[string]int, need } func bootstrapDevicesWithMetalRetry(firstAttemptCtx, retryParentCtx context.Context, timeout time.Duration, ollamaLibDirs []string, extraEnvs map[string]string) []ml.DeviceInfo { + extraEnvs = normalizeDiscoveryEnv(ollamaLibDirs, extraEnvs) + runDiscovery := func(ctx context.Context, extraEnvs map[string]string) ([]ml.DeviceInfo, *llm.StatusWriter, error) { start := time.Now() defer func() { @@ -440,6 +443,34 @@ func bootstrapDevicesWithMetalRetry(firstAttemptCtx, retryParentCtx context.Cont return devices } +func normalizeDiscoveryEnv(ollamaLibDirs []string, extraEnvs map[string]string) map[string]string { + return normalizeDiscoveryEnvForGOOS(runtime.GOOS, ollamaLibDirs, extraEnvs) +} + +func normalizeDiscoveryEnvForGOOS(goos string, ollamaLibDirs []string, extraEnvs map[string]string) map[string]string { + if goos != "linux" || len(ollamaLibDirs) == 0 || filepath.Base(ollamaLibDirs[len(ollamaLibDirs)-1]) != "rocm" { + return extraEnvs + } + + if extraEnvs["ROCR_VISIBLE_DEVICES"] != "" || envconfig.RocrVisibleDevices() != "" { + return extraEnvs + } + + source, tokens := rocmNumericVisibleDeviceSource(extraEnvs) + if len(tokens) == 0 { + return extraEnvs + } + + env := make(map[string]string, len(extraEnvs)+1) + for k, v := range extraEnvs { + env[k] = v + } + env["ROCR_VISIBLE_DEVICES"] = strings.Join(tokens, ",") + env[source] = visibleDeviceOrdinals(len(tokens)) + slog.Debug("normalizing AMD visible devices for ROCm discovery", "from_env", source, "ROCR_VISIBLE_DEVICES", env["ROCR_VISIBLE_DEVICES"], "visible_ordinals", env[source]) + return env +} + type bootstrapDevicesResult struct { devices []ml.DeviceInfo status *llm.StatusWriter @@ -471,6 +502,104 @@ func runBootstrapDevicesWithStatusWatchdog( } } +func remapFilterIDForUserVisibleDevices(device *ml.DeviceInfo) { + tokens := visibleDeviceFilterTokens(runtime.GOOS, device.Library) + if len(tokens) == 0 { + return + } + + id := device.FilterID + if id == "" { + id = device.ID + } + index, err := strconv.Atoi(id) + if err != nil || index < 0 || index >= len(tokens) { + return + } + + device.FilterID = tokens[index] +} + +func visibleDeviceFilterTokens(goos, library string) []string { + switch library { + case "CUDA": + return splitVisibleDeviceList(envconfig.CudaVisibleDevices()) + case "ROCm": + if goos == "linux" { + if tokens := splitVisibleDeviceList(envconfig.RocrVisibleDevices()); len(tokens) > 0 { + return tokens + } + if _, tokens := rocmNumericVisibleDeviceSource(nil); len(tokens) > 0 { + return tokens + } + return nil + } + for _, value := range []string{envconfig.HipVisibleDevices(), envconfig.GpuDeviceOrdinal(), envconfig.CudaVisibleDevices()} { + if tokens := splitNumericVisibleDeviceList(value); len(tokens) > 0 { + return tokens + } + } + case "Vulkan": + return splitVisibleDeviceList(envconfig.VkVisibleDevices()) + } + + return nil +} + +func rocmNumericVisibleDeviceSource(extraEnvs map[string]string) (string, []string) { + for _, name := range []string{"HIP_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL", "CUDA_VISIBLE_DEVICES"} { + value := extraEnvs[name] + if value == "" { + switch name { + case "HIP_VISIBLE_DEVICES": + value = envconfig.HipVisibleDevices() + case "GPU_DEVICE_ORDINAL": + value = envconfig.GpuDeviceOrdinal() + case "CUDA_VISIBLE_DEVICES": + value = envconfig.CudaVisibleDevices() + } + } + if tokens := splitNumericVisibleDeviceList(value); len(tokens) > 0 { + return name, tokens + } + } + return "", nil +} + +func splitVisibleDeviceList(value string) []string { + fields := strings.Split(value, ",") + tokens := make([]string, 0, len(fields)) + for _, field := range fields { + field = strings.TrimSpace(field) + if field != "" { + tokens = append(tokens, field) + } + } + return tokens +} + +func splitNumericVisibleDeviceList(value string) []string { + tokens := splitVisibleDeviceList(value) + if len(tokens) == 0 { + return nil + } + for _, token := range tokens { + index, err := strconv.Atoi(token) + if err != nil || index < 0 { + return nil + } + } + return tokens +} + +func visibleDeviceOrdinals(count int) string { + ordinals := make([]string, count) + for i := range ordinals { + ordinals[i] = strconv.Itoa(i) + } + return strings.Join(ordinals, ",") +} + func lastDiscoveryStatusError(status *llm.StatusWriter) string { if status == nil { return "" diff --git a/discover/runner_test.go b/discover/runner_test.go index 7768aa88f..7c675ca45 100644 --- a/discover/runner_test.go +++ b/discover/runner_test.go @@ -131,6 +131,150 @@ func TestRecordPersistentRunnerEnv(t *testing.T) { } } +func TestRemapFilterIDForUserVisibleDevices(t *testing.T) { + tests := []struct { + name string + env map[string]string + device ml.DeviceInfo + wantID string + wantFilter string + }{ + { + name: "cuda numeric parent filter", + env: map[string]string{"CUDA_VISIBLE_DEVICES": "1"}, + device: ml.DeviceInfo{ + DeviceID: ml.DeviceID{Library: "CUDA", ID: "0"}, + FilterID: "0", + }, + wantID: "0", + wantFilter: "1", + }, + { + name: "cuda uuid parent filter", + env: map[string]string{"CUDA_VISIBLE_DEVICES": "GPU-f3a94ab8-b31d-61ff-9fbb-ce91ac1cdd95"}, + device: ml.DeviceInfo{ + DeviceID: ml.DeviceID{Library: "CUDA", ID: "0"}, + FilterID: "0", + }, + wantID: "0", + wantFilter: "GPU-f3a94ab8-b31d-61ff-9fbb-ce91ac1cdd95", + }, + { + name: "rocm hip parent filter", + env: map[string]string{"HIP_VISIBLE_DEVICES": "2,0"}, + device: ml.DeviceInfo{ + DeviceID: ml.DeviceID{Library: "ROCm", ID: "1"}, + FilterID: "1", + }, + wantID: "1", + wantFilter: "0", + }, + { + name: "vulkan parent filter", + env: map[string]string{"GGML_VK_VISIBLE_DEVICES": "1"}, + device: ml.DeviceInfo{ + DeviceID: ml.DeviceID{Library: "Vulkan", ID: "0"}, + FilterID: "0", + }, + wantID: "0", + wantFilter: "1", + }, + { + name: "no parent filter keeps internal filter id", + device: ml.DeviceInfo{ + DeviceID: ml.DeviceID{Library: "CUDA", ID: "0"}, + FilterID: "3", + }, + wantID: "0", + wantFilter: "3", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for key, value := range tt.env { + t.Setenv(key, value) + } + + remapFilterIDForUserVisibleDevices(&tt.device) + + if tt.device.ID != tt.wantID { + t.Fatalf("ID = %q, want %q", tt.device.ID, tt.wantID) + } + if tt.device.FilterID != tt.wantFilter { + t.Fatalf("FilterID = %q, want %q", tt.device.FilterID, tt.wantFilter) + } + }) + } +} + +func TestNormalizeROCmDiscoveryEnv(t *testing.T) { + tests := []struct { + name string + env map[string]string + extra map[string]string + wantROCR string + wantSource string + wantOrdinal string + wantSame bool + }{ + { + name: "hip becomes rocr", + env: map[string]string{"HIP_VISIBLE_DEVICES": "2"}, + wantROCR: "2", + wantSource: "HIP_VISIBLE_DEVICES", + wantOrdinal: "0", + }, + { + name: "gpu ordinal becomes rocr", + env: map[string]string{"GPU_DEVICE_ORDINAL": "3"}, + wantROCR: "3", + wantSource: "GPU_DEVICE_ORDINAL", + wantOrdinal: "0", + }, + { + name: "cuda numeric becomes rocr", + env: map[string]string{"CUDA_VISIBLE_DEVICES": "2,0"}, + wantROCR: "2,0", + wantSource: "CUDA_VISIBLE_DEVICES", + wantOrdinal: "0,1", + }, + { + name: "rocr wins", + env: map[string]string{"ROCR_VISIBLE_DEVICES": "1", "HIP_VISIBLE_DEVICES": "2"}, + wantSame: true, + }, + { + name: "cuda uuid does not become rocr", + env: map[string]string{"CUDA_VISIBLE_DEVICES": "GPU-f3a94ab8-b31d-61ff-9fbb-ce91ac1cdd95"}, + wantSame: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for key, value := range tt.env { + t.Setenv(key, value) + } + + got := normalizeDiscoveryEnvForGOOS("linux", []string{"/lib/ollama", "/lib/ollama/rocm"}, tt.extra) + + if tt.wantSame { + if got != nil && got["ROCR_VISIBLE_DEVICES"] != "" { + t.Fatalf("ROCR_VISIBLE_DEVICES = %q, want unset", got["ROCR_VISIBLE_DEVICES"]) + } + return + } + if got["ROCR_VISIBLE_DEVICES"] != tt.wantROCR { + t.Fatalf("ROCR_VISIBLE_DEVICES = %q, want %q", got["ROCR_VISIBLE_DEVICES"], tt.wantROCR) + } + if got[tt.wantSource] != tt.wantOrdinal { + t.Fatalf("%s = %q, want %q", tt.wantSource, got[tt.wantSource], tt.wantOrdinal) + } + }) + } +} + func TestBootstrapDevicesWithStatusWatchdogReturnsResult(t *testing.T) { want := []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA", ID: "0"}}} devices, _, err := runBootstrapDevicesWithStatusWatchdog( diff --git a/llm/llama_server.go b/llm/llama_server.go index 0394d1e9f..a45244aa9 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -457,7 +457,7 @@ func llamaServerLibraryPaths(exe string, gpuLibs []string, envUpdates map[string addPath(dir) } } - return libraryPaths + return adjustPlatformLibraryPaths(libraryPaths, gpuLibs) } func findLlamaServerGPUBackend(dir string) string { diff --git a/llm/rocm_default.go b/llm/rocm_default.go new file mode 100644 index 000000000..50850726f --- /dev/null +++ b/llm/rocm_default.go @@ -0,0 +1,7 @@ +//go:build !windows + +package llm + +func adjustPlatformLibraryPaths(paths, _ []string) []string { + return paths +} diff --git a/llm/rocm_windows.go b/llm/rocm_windows.go new file mode 100644 index 000000000..4858d9ee0 --- /dev/null +++ b/llm/rocm_windows.go @@ -0,0 +1,282 @@ +//go:build windows + +package llm + +import ( + "errors" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + "unsafe" + + "golang.org/x/sys/windows" +) + +var windowsROCmRuntimeDLLNames = []string{"amdhip64_7.dll", "amdhip64_6.dll", "amdhip64.dll"} +var windowsROCmCompanionDLLNames = []string{"amd_comgr.dll", "amd_comgr_2.dll"} + +type windowsROCmRuntimeDLL struct { + path string + source string + name string + bundledPath string + bundledVer windowsFileVersion + bundledVerOK bool + systemPath string + systemVer windowsFileVersion + systemVerOK bool + systemDir string + bundledDir string +} + +type windowsFileVersion struct { + ms uint32 + ls uint32 +} + +func WindowsROCmRuntimeDLLPath(libDirs []string) (string, error) { + choice, err := windowsROCmRuntimeDLLChoice(libDirs) + if err != nil { + return "", err + } + return choice.path, nil +} + +func adjustPlatformLibraryPaths(paths, gpuLibs []string) []string { + rocmDir := firstWindowsROCmLibDir(gpuLibs) + if rocmDir == "" { + return paths + } + + choice, err := windowsROCmRuntimeDLLChoice(gpuLibs) + if err != nil { + slog.Debug("windows ROCm runtime selection unavailable", "error", err) + return paths + } + + slog.Debug("selected windows ROCm runtime", + "runtime_source", choice.source, + "name", choice.name, + "path", choice.path, + "bundled", choice.bundledPath, + "bundled_version", choice.bundledVer.String(), + "system", choice.systemPath, + "system_version", choice.systemVer.String(), + ) + + if choice.source != "system" || choice.systemDir == "" { + return paths + } + + before := choice.bundledDir + if before == "" { + before = rocmDir + } + + return insertPathBefore(paths, choice.systemDir, before) +} + +// AMD's Windows driver also ships ROCm runtime DLLs in System32. Within the +// same DLL name/major, use the newer driver or bundled copy; never search the +// ROCm SDK installation paths. +func windowsROCmRuntimeDLLChoice(libDirs []string) (windowsROCmRuntimeDLL, error) { + systemDir, err := windows.GetSystemDirectory() + if err != nil { + return windowsROCmRuntimeDLL{}, err + } + systemDir = filepath.Clean(systemDir) + + for _, name := range windowsROCmRuntimeDLLNames { + bundledPath := firstExistingFile(libDirs, name) + if bundledPath == "" { + continue + } + + choice := windowsROCmRuntimeDLL{ + path: bundledPath, + source: "bundled", + name: name, + bundledPath: bundledPath, + bundledDir: filepath.Dir(bundledPath), + systemPath: filepath.Join(systemDir, name), + systemDir: systemDir, + } + choice.bundledVer, choice.bundledVerOK = readWindowsFileVersion(choice.bundledPath) + if fileExists(choice.systemPath) { + choice.systemVer, choice.systemVerOK = readWindowsFileVersion(choice.systemPath) + if choice.systemVerOK && choice.bundledVerOK && choice.systemVer.Compare(choice.bundledVer) > 0 { + choice.path = choice.systemPath + choice.source = "system" + } + } + if choice.source == "system" && !systemROCmCompanionDLLsCompatible(libDirs, choice.systemDir) { + choice.path = choice.bundledPath + choice.source = "bundled" + } + return choice, nil + } + + for _, name := range windowsROCmRuntimeDLLNames { + path := filepath.Join(systemDir, name) + if !fileExists(path) { + continue + } + choice := windowsROCmRuntimeDLL{ + path: path, + source: "system", + name: name, + systemPath: path, + systemDir: systemDir, + } + choice.systemVer, choice.systemVerOK = readWindowsFileVersion(path) + return choice, nil + } + + return windowsROCmRuntimeDLL{}, errors.New("no amdhip64 runtime DLL found") +} + +func systemROCmCompanionDLLsCompatible(libDirs []string, systemDir string) bool { + for _, name := range windowsROCmCompanionDLLNames { + bundledPath := firstExistingFile(libDirs, name) + if bundledPath == "" { + continue + } + + systemPath := filepath.Join(systemDir, name) + if !fileExists(systemPath) { + continue + } + + bundledVer, bundledOK := readWindowsFileVersion(bundledPath) + systemVer, systemOK := readWindowsFileVersion(systemPath) + if !bundledOK || !systemOK { + slog.Debug("keeping bundled ROCm runtime because companion DLL version is unavailable", + "name", name, "bundled", bundledPath, "system", systemPath) + return false + } + if systemVer.Compare(bundledVer) < 0 { + slog.Debug("keeping bundled ROCm runtime because system companion DLL is older", + "name", name, + "bundled", bundledPath, + "bundled_version", bundledVer.String(), + "system", systemPath, + "system_version", systemVer.String(), + ) + return false + } + } + return true +} + +func firstWindowsROCmLibDir(libDirs []string) string { + for _, dir := range libDirs { + if dir == "" { + continue + } + base := strings.ToLower(filepath.Base(dir)) + if strings.Contains(base, "rocm") || strings.Contains(base, "hip") { + return filepath.Clean(dir) + } + if fileExists(filepath.Join(dir, "ggml-hip.dll")) || fileExists(filepath.Join(dir, "libggml-hip.dll")) { + return filepath.Clean(dir) + } + } + return "" +} + +func firstExistingFile(dirs []string, name string) string { + for _, dir := range dirs { + if dir == "" { + continue + } + path := filepath.Join(dir, name) + if fileExists(path) { + return filepath.Clean(path) + } + } + return "" +} + +func insertPathBefore(paths []string, insert, before string) []string { + insert = filepath.Clean(insert) + before = filepath.Clean(before) + insertKey := strings.ToLower(insert) + beforeKey := strings.ToLower(before) + + out := make([]string, 0, len(paths)+1) + inserted := false + for _, path := range paths { + clean := filepath.Clean(path) + key := strings.ToLower(clean) + if key == insertKey { + continue + } + if !inserted && key == beforeKey { + out = append(out, insert) + inserted = true + } + out = append(out, path) + } + if !inserted { + out = append(out, insert) + } + return out +} + +func fileExists(path string) bool { + info, err := os.Stat(path) + return err == nil && !info.IsDir() +} + +func readWindowsFileVersion(path string) (windowsFileVersion, bool) { + var zero windows.Handle + infoSize, err := windows.GetFileVersionInfoSize(path, &zero) + if err != nil || infoSize == 0 { + return windowsFileVersion{}, false + } + + versionInfo := make([]byte, infoSize) + if err := windows.GetFileVersionInfo(path, 0, infoSize, unsafe.Pointer(&versionInfo[0])); err != nil { + return windowsFileVersion{}, false + } + + var fixedInfo *windows.VS_FIXEDFILEINFO + var fixedInfoLen uint32 + if err := windows.VerQueryValue(unsafe.Pointer(&versionInfo[0]), `\`, unsafe.Pointer(&fixedInfo), &fixedInfoLen); err != nil { + return windowsFileVersion{}, false + } + if fixedInfo == nil { + return windowsFileVersion{}, false + } + return windowsFileVersion{ms: fixedInfo.FileVersionMS, ls: fixedInfo.FileVersionLS}, true +} + +func (v windowsFileVersion) Compare(other windowsFileVersion) int { + if v.ms < other.ms { + return -1 + } + if v.ms > other.ms { + return 1 + } + if v.ls < other.ls { + return -1 + } + if v.ls > other.ls { + return 1 + } + return 0 +} + +func (v windowsFileVersion) String() string { + if v.ms == 0 && v.ls == 0 { + return "" + } + return fmt.Sprintf("%d.%d.%d.%d", + (v.ms>>16)&0xffff, + v.ms&0xffff, + (v.ls>>16)&0xffff, + v.ls&0xffff, + ) +} diff --git a/logutil/logutil.go b/logutil/logutil.go index 00daf6a6e..e93df39fe 100644 --- a/logutil/logutil.go +++ b/logutil/logutil.go @@ -23,8 +23,9 @@ func NewLogger(w io.Writer, level slog.Level) *slog.Logger { attr.Value = slog.StringValue("TRACE") } case slog.SourceKey: - source := attr.Value.Any().(*slog.Source) - source.File = filepath.Base(source.File) + if source, ok := attr.Value.Any().(*slog.Source); ok { + source.File = filepath.Base(source.File) + } } return attr }, diff --git a/logutil/logutil_test.go b/logutil/logutil_test.go new file mode 100644 index 000000000..42e26a3ff --- /dev/null +++ b/logutil/logutil_test.go @@ -0,0 +1,19 @@ +package logutil + +import ( + "bytes" + "log/slog" + "strings" + "testing" +) + +func TestNewLoggerAllowsSourceAttr(t *testing.T) { + var buf bytes.Buffer + logger := NewLogger(&buf, slog.LevelDebug) + + logger.Debug("message", "source", "runtime") + + if !strings.Contains(buf.String(), `source=runtime`) { + t.Fatalf("expected user source attr in log, got %q", buf.String()) + } +} diff --git a/ml/device.go b/ml/device.go index 823df2c67..1c760ae2b 100644 --- a/ml/device.go +++ b/ml/device.go @@ -10,6 +10,7 @@ import ( "log/slog" "math" "net/http" + "os" "runtime" "slices" "sort" @@ -312,6 +313,10 @@ type DeviceInfo struct { DriverMajor int `json:"driver_major,omitempty"` DriverMinor int `json:"driver_minor,omitempty"` + // NVIDIADriverMajor is the NVIDIA kernel driver branch. CUDA driver APIs + // expose a separate CUDA compatibility version, so keep this distinct. + NVIDIADriverMajor int `json:"-"` + // GFXTarget is the AMD GPU gfx target string (e.g. "gfx1100") for ROCm // device validation. Empty on non-AMD devices. GFXTarget string `json:"gfx_target,omitempty"` @@ -595,12 +600,15 @@ func (d DeviceInfo) PreferredLibrary(other DeviceInfo) bool { func (d DeviceInfo) updateVisibleDevicesEnv(env map[string]string, mustFilter bool) { var envVar string + var rocmOrdinalEnv string switch d.Library { case "ROCm": // ROCm must be filtered as it can crash the runner on unsupported devices envVar = "ROCR_VISIBLE_DEVICES" if runtime.GOOS != "linux" { - envVar = "HIP_VISIBLE_DEVICES" + envVar = rocmNonLinuxVisibleDevicesEnv() + } else { + rocmOrdinalEnv = rocmLinuxOrdinalVisibleDevicesEnv() } case "CUDA": if !mustFilter { @@ -618,6 +626,7 @@ func (d DeviceInfo) updateVisibleDevicesEnv(env map[string]string, mustFilter bo return } v, existing := env[envVar] + childOrdinal := visibleDeviceCount(v) if existing { v = v + "," } @@ -627,6 +636,63 @@ func (d DeviceInfo) updateVisibleDevicesEnv(env map[string]string, mustFilter bo v = v + d.ID } env[envVar] = v + + if rocmOrdinalEnv != "" { + v, existing = env[rocmOrdinalEnv] + if existing { + v = v + "," + } + v = v + strconv.Itoa(childOrdinal) + env[rocmOrdinalEnv] = v + } +} + +func visibleDeviceCount(value string) int { + count := 0 + for _, field := range strings.Split(value, ",") { + if strings.TrimSpace(field) != "" { + count++ + } + } + return count +} + +func rocmLinuxOrdinalVisibleDevicesEnv() string { + if runtime.GOOS != "linux" || os.Getenv("ROCR_VISIBLE_DEVICES") != "" { + return "" + } + for _, name := range []string{"HIP_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL", "CUDA_VISIBLE_DEVICES"} { + if numericVisibleDeviceList(os.Getenv(name)) { + return name + } + } + return "" +} + +func rocmNonLinuxVisibleDevicesEnv() string { + for _, name := range []string{"HIP_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL", "CUDA_VISIBLE_DEVICES"} { + if numericVisibleDeviceList(os.Getenv(name)) { + return name + } + } + return "HIP_VISIBLE_DEVICES" +} + +func numericVisibleDeviceList(value string) bool { + fields := strings.Split(value, ",") + found := false + for _, field := range fields { + field = strings.TrimSpace(field) + if field == "" { + continue + } + index, err := strconv.Atoi(field) + if err != nil || index < 0 { + return false + } + found = true + } + return found } type BaseRunner interface {