diff --git a/llm/llama_server.go b/llm/llama_server.go index 86f3d3c5..540aa564 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -77,6 +77,14 @@ const ( openEndedGenerationContextMultiplier = 10 ) +const ( + llamaArgFitTargetEnv = "LLAMA_ARG_FIT_TARGET" + bytesPerMiB = 1 << 20 + + // mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights. + mmprojOffloadHeadroom = 1 << 30 +) + // DefaultEmbeddingNumBatchForContext caps the embedding batch default to the // active context length before it is passed to llama-server. func DefaultEmbeddingNumBatchForContext(numCtx int) int { @@ -420,7 +428,7 @@ func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec. cmd.Stderr = out } cmd.SysProcAttr = LlamaServerSysProcAttr - SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvs) + SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvsForStart()) slog.Info("starting llama-server", "cmd", cmd) slog.Debug("subprocess", "", filteredEnv(cmd.Env)) @@ -621,11 +629,6 @@ func appendMainGPUArgs(params []string, opts api.Options) []string { return append(params, "--split-mode", "none", "--main-gpu", strconv.Itoa(*opts.MainGPU)) } -const ( - // mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights. - mmprojOffloadHeadroom = 1 << 30 -) - func appendMMProjArgs(params []string, launch llamaServerLaunchConfig) []string { if len(launch.projectors) == 0 { return params @@ -658,9 +661,6 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay requiredMemory := mmprojMemory + mmprojOffloadHeadroom for _, gpu := range gpus { - if gpu.Integrated && gpu.Library != "Metal" { - return true, "shared-memory-gpu" - } memory := gpu.FreeMemory if memory == 0 || (gpu.TotalMemory > 0 && gpu.TotalMemory < memory) { memory = gpu.TotalMemory @@ -673,6 +673,47 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay return false, "" } +func (launch llamaServerLaunchConfig) extraEnvsForStart() map[string]string { + pad, ok := launch.mmprojFitTargetMiB() + if !ok { + return launch.extraEnvs + } + + if existing, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok { + existingTarget, err := strconv.ParseUint(existing, 10, 64) + if err != nil { + slog.Warn("invalid llama-server fit target", "env", llamaArgFitTargetEnv, "value", existing, "error", err) + return launch.extraEnvs + } + + envs := cloneStringMap(launch.extraEnvs) + envs[llamaArgFitTargetEnv] = strconv.FormatUint(existingTarget+pad, 10) + return envs + } + + if _, ok := os.LookupEnv(llamaArgFitTargetEnv); ok { + // Preserve an inherited user override. SetupLlamaServerCommandEnv + // will pass it through unless extraEnvs overrides it. + return launch.extraEnvs + } + + envs := cloneStringMap(launch.extraEnvs) + envs[llamaArgFitTargetEnv] = strconv.FormatUint(pad, 10) + return envs +} + +func (launch llamaServerLaunchConfig) mmprojFitTargetMiB() (uint64, bool) { + if len(launch.projectors) == 0 || launch.mmprojMemory == 0 { + return 0, false + } + if disable, _ := launch.mmprojOffloadDisabled(); disable { + return 0, false + } + + requiredMemory := launch.mmprojMemory + mmprojOffloadHeadroom + return (requiredMemory + bytesPerMiB - 1) / bytesPerMiB, true +} + // mmprojMemoryRequirement is a stopgap until fit accounts for mmproj memory directly. func mmprojMemoryRequirement(modelPath string, f *ggml.GGML, projectors []string) (uint64, error) { if len(projectors) == 0 { diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index 70e66302..d39b5b08 100644 --- a/llm/llama_server_test.go +++ b/llm/llama_server_test.go @@ -2121,13 +2121,22 @@ func TestAppendMMProjArgs(t *testing.T) { want: []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"}, }, { - name: "integrated rocm gpu disables projector offload", + name: "integrated rocm gpu keeps projector offload when projector fits", projectors: []string{"model.gguf"}, opts: defaultOpts, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "ROCm"}, Integrated: true, FreeMemory: 32 << 30}}, mmprojMemory: 933 << 20, modelLayers: 81, - want: []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"}, + want: []string{"base", "--mmproj", "model.gguf"}, + }, + { + name: "integrated cuda gpu keeps projector offload", + projectors: []string{"model.gguf"}, + opts: defaultOpts, + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}}, + mmprojMemory: 933 << 20, + modelLayers: 81, + want: []string{"base", "--mmproj", "model.gguf"}, }, { name: "integrated metal gpu keeps projector offload", @@ -2195,6 +2204,75 @@ func TestAppendMMProjArgs(t *testing.T) { } } +func TestMMProjFitTargetExtraEnvs(t *testing.T) { + t.Setenv(llamaArgFitTargetEnv, "") + _ = os.Unsetenv(llamaArgFitTargetEnv) + + const ( + projectorMemoryMiB = uint64(933) + projectorPadMiB = projectorMemoryMiB + mmprojOffloadHeadroom/bytesPerMiB + ) + + fitTargetValue := func(mib uint64) string { + return fmt.Sprint(mib) + } + + assertFitTarget := func(t *testing.T, got map[string]string, wantMiB uint64) { + t.Helper() + if got[llamaArgFitTargetEnv] != fitTargetValue(wantMiB) { + t.Fatalf("fit target = %q, want %d", got[llamaArgFitTargetEnv], wantMiB) + } + } + + newLaunch := func(extraEnvs map[string]string) llamaServerLaunchConfig { + return llamaServerLaunchConfig{ + projectors: []string{"model.gguf"}, + mmprojMemory: projectorMemoryMiB * bytesPerMiB, + opts: api.DefaultOptions(), + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}}, + modelLayers: 81, + extraEnvs: extraEnvs, + } + } + + t.Run("sets projector pad when no fit target exists", func(t *testing.T) { + launch := newLaunch(map[string]string{"KEEP": "1"}) + + got := launch.extraEnvsForStart() + assertFitTarget(t, got, projectorPadMiB) + if _, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok { + t.Fatal("extraEnvsForStart mutated launch.extraEnvs") + } + }) + + for _, tt := range []struct { + name string + launchFitTargetMiB uint64 + }{ + {name: "adds projector pad to existing launch fit target", launchFitTargetMiB: 2048}, + {name: "adds projector pad to smaller launch fit target", launchFitTargetMiB: 512}, + } { + t.Run(tt.name, func(t *testing.T) { + launch := newLaunch(map[string]string{ + llamaArgFitTargetEnv: fitTargetValue(tt.launchFitTargetMiB), + }) + + got := launch.extraEnvsForStart() + assertFitTarget(t, got, tt.launchFitTargetMiB+projectorPadMiB) + }) + } + + t.Run("preserves inherited user fit target", func(t *testing.T) { + t.Setenv(llamaArgFitTargetEnv, fitTargetValue(256)) + launch := newLaunch(map[string]string{}) + + got := launch.extraEnvsForStart() + if _, ok := got[llamaArgFitTargetEnv]; ok { + t.Fatalf("user env should not be overridden, got extra env %q", got[llamaArgFitTargetEnv]) + } + }) +} + func TestMMProjMemoryRequirement(t *testing.T) { if got, err := mmprojMemoryRequirement("model.gguf", nil, nil); err != nil || got != 0 { t.Fatalf("no projector memory = %d, %v; want 0, nil", got, err)