Repository: ollama/ollama
File: llm/memory.go (Go)

Code before:
	overhead := envconfig.GpuOverhead()
	availableList := make([]string, len(gpus))
	for i, gpu := range gpus {
		availableList[i] = format.HumanBytes2(gpu.FreeMemory)
	}
	slog.Debug("evaluating", "library", gpus[0].Library, "gpu_count", len(gpus), "available", availableList)

	for _, projector := range projectors {
		weight, graph := projectorMemoryRequirements(projector)
		projectorWeights += weight
		projectorGraph += graph

		// multimodal models require at least 2048 context
		opts.NumCtx = max(opts.NumCtx, 2048)
	}
	if projectorWeights == 0 && projectorGraph == 0 {
		projectorWeights, projectorGraph = f.VisionGraphSize()
	}

	layers := f.Tensors().GroupLayers()
	// add one layer worth of memory as a buffer
	// use layer with maximum size to avoid vram overflow
	maxKey, maxLayer := MapMaxFunc(layers, func(a, b ggml.Layer) int {
		return cmp.Compare(a.Size(), b.Size())
	})
	layerSize = maxLayer.Size()

	var kvct string
	if envconfig.FlashAttention() &&
		discover.GetGPUInfo().FlashAttentionSupported() &&
		f.SupportsFlashAttention() {
		requested := strings.ToLower(envconfig.KvCacheType())
		if requested != "" && f.SupportsKVCacheType(requested) {
			kvct = requested
		}
	}

	kv, graphPartialOffload, graphFullOffload := f.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)), numParallel, kvct)

	if len(kv) > 0 {
		i, err := strconv.Atoi(strings.TrimPrefix(maxKey, "blk."))
		if err != nil {
			slog.Error(fmt.Sprintf("Invalid block key: %v", err))
		}
		layerSize += kv[i]
	}

	var kvTotal uint64
	for _, kvLayer := range kv {
		kvTotal += kvLayer
	}

Review comment (performance):
I'm not sure that the new comment here is really correct. Does this part actually make a difference in your tests or is it reversing the allocation order that fixes the problem?

As the original comment mentions, this is just a buffer of some extra space, not something that we technically need.

Note that the layer selected is not necessarily the largest when taking the KV cache into account. This is because for large context sizes, the KV cache will dominate and the layer with the largest weights is not necessarily the one with the largest KV cache.

To be honest, I think this is getting overly complicated (especially with all of the generics, etc. - we don't need to write a standard library function here). Unless this is actually critical and not just acting as an extra buffer, I would just leave this part as it is and use layer 0 (including for the KV cache).

Code after:
		// multimodal models require at least 2048 context
		opts.NumCtx = max(opts.NumCtx, 2048)
	}
	if projectorWeights == 0 && projectorGraph == 0 {
		projectorWeights, projectorGraph = f.VisionGraphSize()
	}

	layers := f.Tensors().GroupLayers()
	// add one layer (chosing the max layer) worth of memory as a buffer
	layerSize = slices.MaxFunc(slices.Collect(maps.Values(layers)), func(a, b ggml.Layer) int {
		return cmp.Compare(a.Size(), b.Size())
	}).Size()

	var kvct string
	if envconfig.FlashAttention() &&
		discover.GetGPUInfo().FlashAttentionSupported() &&
		f.SupportsFlashAttention() {
		requested := strings.ToLower(envconfig.KvCacheType())
		if requested != "" && f.SupportsKVCacheType(requested) {
			kvct = requested
		}
	}

	kv, graphPartialOffload, graphFullOffload := f.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)), numParallel, kvct)

	if len(kv) > 0 {
		layerSize += kv[0]
	}

	var kvTotal uint64
	for _, kvLayer := range kv {
		kvTotal += kvLayer
	}

	if graphPartialOffload == 0 {
		graphPartialOffload = f.KV().GQA() * kvTotal / 6
	}
	if graphFullOffload == 0 {
		graphFullOffload = graphPartialOffload
	}

	// on metal there's no partial offload overhead
	if gpus[0].Library == "metal" {
		graphPartialOffload = graphFullOffload
	} else if len(gpus) > 1 {
		// multigpu should always use the partial graph size
		graphFullOffload = graphPartialOffload
	}

	if layer, ok := layers["output_norm"]; ok {
		memoryLayerOutput += layer.Size()