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

Code before:
	}

	if envconfig.Debug() {
		params = append(params, "--verbose")
	}

	if opts.MainGPU > 0 {
		params = append(params, "--main-gpu", strconv.Itoa(opts.MainGPU))
	}

	if len(adapters) > 0 {
		// TODO: applying multiple adapters is not supported by the llama.cpp server yet
		params = append(params, "--lora", adapters[0])
	}

	if len(projectors) > 0 {
		// TODO: applying multiple projectors is not supported by the llama.cpp server yet
		params = append(params, "--mmproj", projectors[0])
	}

	defaultThreads := systemInfo.GetOptimalThreadCount()
	if opts.NumThread > 0 {
		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
	} else if defaultThreads > 0 {
		params = append(params, "--threads", strconv.Itoa(defaultThreads))
	}

	// isEmbeddingModel checks for common GGML attributes that help distinguish most embedding models from normal models.
	isEmbeddingModel := false
	if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
		isEmbeddingModel = true
	}

	// Validates and applies KV cache parameters
	setCacheTypeParam := func(paramName, cacheType string) {
		if cacheType == "" {
			return
		}

		validCacheTypes := []string{"f32", "f16", "q8_0", "q5_1", "q5_0", "iq4_nl", "q4_1", "q4_0"}
		if !slices.Contains(validCacheTypes, cacheType) {
			slog.Warn("invalid cache type, ignoring", "param", paramName, "type", cacheType)
			return
		}

		// For embedding models, only allow f16 and f32
		if isEmbeddingModel && cacheType != "f16" && cacheType != "f32" {
			slog.Warn("only f16 and f32 cache types are supported for embedding models, ignoring",
				"param", paramName, "type", cacheType)
			return
		}

Review comment (performance):
The checks here encompass what we need which is awesome. Without going overboard, I'm wondering if there's a way for us to pull out the checks 1/ flash attention compatibility and 2/ kv cache types so we can write some unit tests around that?

Code after:

	if opts.MainGPU > 0 {
		params = append(params, "--main-gpu", strconv.Itoa(opts.MainGPU))
	}

	if len(adapters) > 0 {
		for _, adapter := range adapters {
			params = append(params, "--lora", adapter)
		}
	}

	if len(projectors) > 0 {
		// TODO: applying multiple projectors is not supported by the llama.cpp server yet
		params = append(params, "--mmproj", projectors[0])
	}

	defaultThreads := systemInfo.GetOptimalThreadCount()
	if opts.NumThread > 0 {
		params = append(params, "--threads", strconv.Itoa(opts.NumThread))
	} else if defaultThreads > 0 {
		params = append(params, "--threads", strconv.Itoa(defaultThreads))
	}

	fa := envconfig.FlashAttention()
	if fa && !gpus.FlashAttentionSupported() {
		slog.Warn("flash attention enabled but not supported by gpu")
		fa = false
	}

	if fa && !ggml.SupportsFlashAttention() {
		slog.Warn("flash attention enabled but not supported by model")
		fa = false
	}

	kvct := envconfig.KvCacheType()

	if fa {
		slog.Info("enabling flash attention")
		params = append(params, "--flash-attn")

		// Flash Attention also supports kv cache quantization
		// Enable if the requested and kv cache type is supported by the model
		if kvct != "" && ggml.SupportsKVCacheType(kvct) {
			params = append(params, "--kv-cache-type", kvct)
		} else {
			slog.Warn("kv cache type not supported by model", "type", kvct)
		}
	} else if kvct != "" && kvct != "f16" {
		slog.Warn("quantized kv cache requested but flash attention disabled", "type", kvct)
	}