ml

package
v0.15.2 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jan 26, 2026 License: MIT Imports: 20 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

View Source 
var LibOllamaPath string = func() string {
	exe, err := os.Executable()
	if err != nil {
		return ""
	}

	if eval, err := filepath.EvalSymlinks(exe); err == nil {
		exe = eval
	}

	var libPath string
	switch runtime.GOOS {
	case "windows":
		libPath = filepath.Join(filepath.Dir(exe), "lib", "ollama")
	case "linux":
		libPath = filepath.Join(filepath.Dir(exe), "..", "lib", "ollama")
	case "darwin":
		libPath = filepath.Dir(exe)
	}

	cwd, err := os.Getwd()
	if err != nil {
		return ""
	}

	paths := []string{
		libPath,

		filepath.Join(filepath.Dir(exe), "build", "lib", "ollama"),
		filepath.Join(cwd, "build", "lib", "ollama"),
	}

	for _, p := range paths {
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}

	return filepath.Dir(exe)
}()

LibPath is a path to lookup dynamic libraries in development it's usually 'build/lib/ollama' in distribution builds it's 'lib/ollama' on Windows '../lib/ollama' on Linux and the executable's directory on macOS note: distribution builds, additional GPU-specific libraries are found in subdirectories of the returned path, such as 'cuda_v12', 'rocm', etc.

Functions

func ByLibrary 

func ByLibrary(l []DeviceInfo) [][]DeviceInfo

func ByPerformance 

func ByPerformance(l []DeviceInfo) [][]DeviceInfo

ByPerformance groups devices by similar speed

func FlashAttentionSupported 

func FlashAttentionSupported(l []DeviceInfo) bool

For each GPU, check if it does NOT support flash attention

func GetVisibleDevicesEnv 

func GetVisibleDevicesEnv(l []DeviceInfo, mustFilter bool) map[string]string

Given the list of GPUs this instantiation is targeted for, figure out the visible devices environment variables Set mustFilter true to enable filtering of CUDA devices

func LibraryPaths 

func LibraryPaths(l []DeviceInfo) []string

func RegisterBackend 

func RegisterBackend(name string, f func(string, BackendParams) (Backend, error))

func WithRoPEBase 

func WithRoPEBase(base float32) func(*RoPEOptions)

func WithRoPEFreqs 

func WithRoPEFreqs(freqs Tensor) func(*RoPEOptions)

Types

type Backend 

type Backend interface {
	Config() fs.Config
	Get(name string) Tensor
	NewContext() Context
}

func NewBackend 

func NewBackend(modelPath string, params BackendParams) (Backend, error)

type BackendCacheConfig 

type BackendCacheConfig interface {
	CacheConfig() CacheConfig
}

BackendCacheConfig should be implemented by backends that need special output from the cache to meet specific requirements. It is frequently implemented in conjunction with ScaledDotProductAttention.

type BackendMemory 

type BackendMemory struct {
	// InputWeights are always located on the CPU and cannot be moved
	InputWeights uint64

	// CPU model components are located in system memory. This does not
	// include unified memory allocated through the GPU.
	CPU DeviceMemory

	// GPU model components are located on one or more GPUs.
	GPUs []DeviceMemory
}

BackendMemory provides the amount of memory required to load the model per device based on the BackendParams. In some cases, not all required allocations will be known at this point. However, the size of the most recent allocation is guaranteed to be provided so that if it failed, the caller can accommodate that to make forward progress.

func (BackendMemory) Log 

func (m BackendMemory) Log(level slog.Level)

Log prints a high level summary of the memory

func (BackendMemory) LogValue 

func (m BackendMemory) LogValue() slog.Value

type BackendParams 

type BackendParams struct {
	// AllocMemory causes the backend to allocate memory for the model. If
	// false, this is only being used for discovering the required amount of
	// memory and cannot load the model for running.
	AllocMemory bool

	// NumThreads sets the number of threads to use if running on the CPU
	NumThreads int

	// GPULayers is the set of layers to offload to GPUs
	GPULayers GPULayersList

	// FlashAttention indicates that we should use a fused flash attention kernel
	FlashAttention bool
}

BackendParams controls how the backend loads and executes models

type BaseRunner 

type BaseRunner interface {
	// GetPort returns the localhost port number the runner is running on
	GetPort() int

	// HasExited indicates if the runner is no longer running.  This can be used during
	// bootstrap to detect if a given filtered device is incompatible and triggered an assert
	HasExited() bool
}

type ByFreeMemory 

type ByFreeMemory []DeviceInfo

Sort by Free Space. iGPUs are reported first, thus Reverse() yields the largest discrete GPU first

func (ByFreeMemory) Len 

func (a ByFreeMemory) Len() int

func (ByFreeMemory) Less 

func (a ByFreeMemory) Less(i, j int) bool

func (ByFreeMemory) Swap 

func (a ByFreeMemory) Swap(i, j int)

type CacheConfig 

type CacheConfig struct {
	// CachePadding specifies the multiple for the number of tokens of cache history
	// that will be returned from cache Get for k, v and mask. The capacity of the
	// cache itself will also be increased to a multiple of this size if needed.
	CachePadding int

	// PermutedV performs Permute(ctx, 1, 2, 0, 3) on v tensors stored via Put
	// and return the permuted version via Get. This uses the cache copy operation
	// to avoid a Contiguous call on the permuted tensor.
	PermutedV bool

	// MaskDType specifies the data type for generating the mask. If unset it will
	// default to DTypeF32.
	MaskDType DType

	// MaskBatchPadding specifies the multiple for the batch size dimension in the mask.
	// Any position that does not correspond to an actual token will be filled with -Inf.
	MaskBatchPadding int
}

CacheConfig controls optimizations (mostly backend-specific) that may transform the output the cache to work better with specific kernels.

type Context 

type Context interface {
	Empty(dtype DType, shape ...int) Tensor
	Zeros(dtype DType, shape ...int) Tensor
	// FromBytes(dtype DType, s []byte, shape ...int) Tensor
	FromFloats(s []float32, shape ...int) Tensor
	FromInts(s []int32, shape ...int) Tensor
	RandomNormal(shape []int, dtype DType, loc, scale float32, key Tensor) Tensor

	// Arange creates a 1D tensor with values within an interval (start, stop] increased by step.
	Arange(start, stop, step float32, dtype DType) Tensor

	Forward(...Tensor) Context

	Compute(...Tensor)

	// MaxGraphNodes() int
	Close()

	// Input returns a context appropriate for creating tensors that are
	// inputs to the model (which includes things like output locations)
	Input() Context

	// Layer returns a context appropriate for creating intermediate tensors
	Layer(int) Context

	// Load a tensor from "filename" safetensors file, and compare with the input tensor
	// Returns error if the shape is inconsistent, or similarity measures are below 99%
	CompareWith(filename string, tensors map[string]Tensor, abortOnError bool) error
}

type DType 

type DType int
const (
	DTypeBool DType = iota
	DTypeUint8
	DTypeUint16
	DTypeUint32
	DTypeUint64
	DTypeInt8
	DTypeInt16
	DTypeInt32
	DTypeInt64
	DTypeFloat16
	DTypeFloat32
	DTypeFloat64
	DTypeBfloat16
	DTypeComplex64
)

type DeviceComparison 

type DeviceComparison int
const (
	UniqueDevice      DeviceComparison = iota
	SameBackendDevice                  // The device is the same, and the library/backend is the same
	DuplicateDevice                    // The same physical device but different library/backend (overlapping device)
)

type DeviceID 

type DeviceID struct {
	// ID is an identifier for the device for matching with system
	// management libraries.  The ID is only unique for other devices
	// using the same Library.
	// This ID represents a "post filtered" view of the enumerated devices
	// if the ID is numeric
	ID string `json:"id"`

	// Library identifies which library is used for the device (e.g. CUDA, ROCm, etc.)
	Library string `json:"backend,omitempty"`
}

Minimal unique device identification

type DeviceInfo 

type DeviceInfo struct {
	DeviceID

	// Name is the name of the device as labeled by the backend. It
	// may not be persistent across instances of the runner.
	Name string `json:"name"`

	// Description is the longer user-friendly identification of the device
	Description string `json:"description"`

	// FilterID is populated with the unfiltered device ID if a numeric ID is used
	// so the device can be included.
	FilterID string `json:"filter_id,omitempty"`

	// Integrated is set true for integrated GPUs, false for Discrete GPUs
	Integrated bool `json:"integration,omitempty"`

	// PCIID is the bus, device and domain ID of the device for deduplication
	// when discovered by multiple backends
	PCIID string `json:"pci_id,omitempty"`

	// TotalMemory is the total amount of memory the device can use for loading models
	TotalMemory uint64 `json:"total_memory"`

	// FreeMemory is the amount of memory currently available on the device for loading models
	FreeMemory uint64 `json:"free_memory,omitempty"`

	// ComputeMajor is the major version of capabilities of the device
	// if unsupported by the backend, -1 will be returned
	ComputeMajor int

	// ComputeMinor is the minor version of capabilities of the device
	// if unsupported by the backend, -1 will be returned
	ComputeMinor int

	// Driver Information
	DriverMajor int `json:"driver_major,omitempty"`
	DriverMinor int `json:"driver_minor,omitempty"`

	// Where backends were loaded from
	LibraryPath []string
}

func GetDevicesFromRunner 

func GetDevicesFromRunner(ctx context.Context, runner BaseRunner) ([]DeviceInfo, error)

func (DeviceInfo) AddInitValidation 

func (d DeviceInfo) AddInitValidation(env map[string]string)

Set the init validation environment variable

func (DeviceInfo) Compare 

func (a DeviceInfo) Compare(b DeviceInfo) DeviceComparison

func (DeviceInfo) Compute 

func (d DeviceInfo) Compute() string

func (DeviceInfo) Driver 

func (d DeviceInfo) Driver() string

func (DeviceInfo) IsBetter 

func (a DeviceInfo) IsBetter(b DeviceInfo) bool

For a SameBackendDevice, return true if b is better than a e.g. newer GPU library version

func (DeviceInfo) MinimumMemory 

func (d DeviceInfo) MinimumMemory() uint64

MinimumMemory reports the amount of memory that should be set aside on the device for overhead (e.g. VRAM consumed by context structures independent of model allocations)

func (DeviceInfo) NeedsInitValidation 

func (d DeviceInfo) NeedsInitValidation() bool

NeedsInitValidation returns true if the device in question has the potential to crash at inference time and requires deeper validation before we include it in the supported devices list.

func (DeviceInfo) PreferredLibrary 

func (d DeviceInfo) PreferredLibrary(other DeviceInfo) bool

PreferredLibrary returns true if this library is preferred over the other input library Used to filter out Vulkan in favor of CUDA or ROCm

type DeviceMemory 

type DeviceMemory struct {
	DeviceID

	// Name is the name of the device as labeled by the backend. It
	// may not be persistent across instances of the runner.
	Name string

	// Weights is the per-layer memory needed for the model weights.
	Weights []uint64

	// Cache is the per-layer memory needed for the KV cache.
	Cache []uint64

	// Graph is the size of the compute graph. It is not per-layer.
	Graph uint64
}

DeviceMemory provides a breakdown of the memory needed per device, such as a CPU or GPU.

func (DeviceMemory) LogValue 

func (m DeviceMemory) LogValue() slog.Value

func (DeviceMemory) Size 

func (m DeviceMemory) Size() uint64

Size returns the total size of the memory required by this device

type ErrNoMem 

type ErrNoMem struct {
	BackendMemory
}

ErrNoMem is returned when panicing due to insufficient memory. It includes the attempted memory allocation.

func (ErrNoMem) Error 

func (e ErrNoMem) Error() string

type FilteredRunnerDiscovery 

type FilteredRunnerDiscovery interface {
	RunnerDiscovery

	// GetActiveDeviceIDs returns the filtered set of devices actively in
	// use by this runner for running models.  If the runner is a bootstrap runner, no devices
	// will be active yet so no device IDs are returned.
	// This routine will not query the underlying device and will return immediately
	GetActiveDeviceIDs() []DeviceID
}

type GPULayers 

type GPULayers struct {
	DeviceID

	// Layers is a set of layer indicies to load
	Layers []int
}

GPULayers is a set of layers to be allocated on a single GPU

func (GPULayers) FirstLayer 

func (g GPULayers) FirstLayer() int

FirstLayer returns the smallest layer index scheduled on this GPU, or MaxInt when empty.

func (GPULayers) String 

func (g GPULayers) String() string

type GPULayersList 

type GPULayersList []GPULayers

GPULayersList is a set of layer allocations across multiple GPUs

func (GPULayersList) Hash 

func (l GPULayersList) Hash() uint64

Hash is an identifier of this layer assignment

func (GPULayersList) Len 

func (l GPULayersList) Len() int

func (GPULayersList) Less 

func (l GPULayersList) Less(i, j int) bool

Sort by the ordering of the layers offloaded

func (GPULayersList) String 

func (l GPULayersList) String() string

func (GPULayersList) Sum 

func (l GPULayersList) Sum() int

Sum is the total number of layers assigned across all GPUs

func (GPULayersList) Swap 

func (l GPULayersList) Swap(i, j int)

type RoPEOptions 

type RoPEOptions struct {
	Base  *float32
	Freqs Tensor
}

type RunnerDiscovery 

type RunnerDiscovery interface {
	BaseRunner

	// GetDeviceInfos will perform a query of the underlying device libraries
	// for device identification and free VRAM information
	// During bootstrap scenarios, this routine may take seconds to complete
	GetDeviceInfos(ctx context.Context) []DeviceInfo
}

type SamplingMode 

type SamplingMode int
const (
	SamplingModeNearest SamplingMode = iota
	SamplingModeBilinear
)

type SystemInfo 

type SystemInfo struct {
	// ThreadCount is the optimal number of threads to use for inference
	ThreadCount int `json:"threads,omitempty"`

	// TotalMemory is the total amount of system memory
	TotalMemory uint64 `json:"total_memory,omitempty"`

	// FreeMemory is the amount of memory currently available on the system for loading models
	FreeMemory uint64 `json:"free_memory,omitempty"`

	// FreeSwap is the amount of system swap space reported as available
	FreeSwap uint64 `json:"free_swap,omitempty"`
}

type Tensor 

type Tensor interface {
	ToString() string
	RoPE(ctx Context, dims int, traditional bool, scale float32, offset int, options ...func(*RoPEOptions)) Tensor
	ScaledDotProductAttention(ctx Context, keys, values Tensor, scale float64, maskMode string, mask Tensor, sinks Tensor) Tensor
	TakeAxes(ctx Context, indicies Tensor, axes int) Tensor

	Dim(n int) int
	Stride(n int) int

	Shape() []int
	DType() DType

	// Bytes() []byte
	Floats() []float32
	Ints() []int32

	Add(ctx Context, t2 Tensor) Tensor
	Sub(ctx Context, t2 Tensor) Tensor

	Max(ctx Context, axes []int, keepDims bool) Tensor
	Min(ctx Context, axes []int, keepDims bool) Tensor

	Matmul(ctx Context, a2 Tensor) Tensor

	Softmax(ctx Context) Tensor
	L2Norm(ctx Context, eps float32) Tensor
	LayerNorm(ctx Context, weight, bias Tensor, eps float32) Tensor
	RMSNorm(ctx Context, weight Tensor, eps float32) Tensor
	Scale(ctx Context, s float64) Tensor

	AvgPool2D(ctx Context, k, s int, p float32) Tensor
	Conv2D(ctx Context, weight Tensor, stride0, stride1, padding0, padding1, dilation0, dilation1, groups int) Tensor
	Conv3D(ctx Context, weight Tensor, stride0, stride1, stride2, padding0, padding1, padding2, dilation0, dilation1, dilation2, groups int) Tensor

	// Sin(ctx Context) Tensor
	// Cos(ctx Context) Tensor
	// Tanh(ctx Context) Tensor
	GELU(ctx Context, up ...Tensor) Tensor

	Reshape(ctx Context, shape ...int) Tensor
	AsStrided(ctx Context, shape, strides []int, offset int) Tensor
	Transpose(ctx Context, shape ...int) Tensor
	Contiguous(ctx Context, allowColMajor bool) Tensor

	Scatter(ctx Context, indicies []Tensor, updates Tensor, axes []int) Tensor

	Copy(ctx Context, t2 Tensor) Tensor
}

Source Files

View all Source files

Directories

Path Synopsis
nn
pooling
rope

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.