"lm_1b/data_utils.py" did not exist on "761c59c461bd15355c3bd805bc329120840ee487"
cuda_common.go 1.66 KB
Newer Older
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1
2
//go:build linux || windows

3
package discover
Daniel Hiltgen's avatar
Daniel Hiltgen committed
4
5

import (
Daniel Hiltgen's avatar
Daniel Hiltgen committed
6
	"fmt"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
7
	"log/slog"
8
9
10
11
	"os"
	"regexp"
	"runtime"
	"strconv"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
13
14
	"strings"
)

15
16
17
18
// Jetson devices have JETSON_JETPACK="x.y.z" factory set to the Jetpack version installed.
// Included to drive logic for reducing Ollama-allocated overhead on L4T/Jetson devices.
var CudaTegra string = os.Getenv("JETSON_JETPACK")

Daniel Hiltgen's avatar
Daniel Hiltgen committed
19
func cudaVariant(gpuInfo CudaGPUInfo) string {
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
	if runtime.GOARCH == "arm64" && runtime.GOOS == "linux" {
		if CudaTegra != "" {
			ver := strings.Split(CudaTegra, ".")
			if len(ver) > 0 {
				return "jetpack" + ver[0]
			}
		} else if data, err := os.ReadFile("/etc/nv_tegra_release"); err == nil {
			r := regexp.MustCompile(` R(\d+) `)
			m := r.FindSubmatch(data)
			if len(m) != 2 {
				slog.Info("Unexpected format for /etc/nv_tegra_release.  Set JETSON_JETPACK to select version")
			} else {
				if l4t, err := strconv.Atoi(string(m[1])); err == nil {
					// Note: mapping from L4t -> JP is inconsistent (can't just subtract 30)
					// https://developer.nvidia.com/embedded/jetpack-archive
					switch l4t {
					case 35:
						return "jetpack5"
					case 36:
						return "jetpack6"
					default:
						slog.Info("unsupported L4T version", "nv_tegra_release", string(data))
					}
				}
			}
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
46
		return "sbsa"
47
48
	}

49
50
	// driver 12.0 has problems with the cuda v12 library, so run v11 on those older drivers
	if gpuInfo.DriverMajor < 12 || (gpuInfo.DriverMajor == 12 && gpuInfo.DriverMinor == 0) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
51
52
		// The detected driver is older than Feb 2023
		slog.Warn("old CUDA driver detected - please upgrade to a newer driver", "version", fmt.Sprintf("%d.%d", gpuInfo.DriverMajor, gpuInfo.DriverMinor))
53
54
55
56
		return "v11"
	}
	return "v12"
}