CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-08-04 14:56:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

revamp GPU discovery

Browse source 
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.
This commit is contained in:
Daniel Hiltgen 2026-05-16 15:39:34 -07:00
parent b501e54316
commit facdd253a0
19 changed files with 2274 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

12
cmd/cmd.go
View file

@ -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),
)

48
discover/amd.go
View file

@ -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)

30
discover/amd_test.go
View file

@ -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

49
discover/cuda_compat.go
View file

@ -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
}

63
discover/llama_server.go
View file

@ -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
}

263
discover/native_probe.go Normal file
View file

@ -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)
}

503
discover/native_probe_linux.go Normal file
View file

@ -0,0 +1,503 @@
//go:build linux
package discover
/*
#cgo linux LDFLAGS: -ldl
#include <dlfcn.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
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
}

12
discover/native_probe_linux_nocgo.go Normal file
View file

@ -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")
}

12
discover/native_probe_stub.go Normal file
View file

@ -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")
}

203
discover/native_probe_test.go Normal file
View file

@ -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)
}
})
}
}

470
discover/native_probe_windows.go Normal file
View file

@ -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)
}

129
discover/runner.go
View file

@ -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 ""

144
discover/runner_test.go
View file

@ -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(

2
llm/llama_server.go
View file

@ -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 {

7
llm/rocm_default.go Normal file
View file

@ -0,0 +1,7 @@
//go:build !windows
package llm
func adjustPlatformLibraryPaths(paths, _ []string) []string {
return paths
}

282
llm/rocm_windows.go Normal file
View file

@ -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,
)
}

5
logutil/logutil.go
View file

@ -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
},

19
logutil/logutil_test.go Normal file
View file

@ -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())
}
}

68
ml/device.go
View file

@ -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 {