feat: add flash_attn support

ollama · May 4, 2024 · 3024414 · 3024414
1 parent 3fca6e7
commit 3024414
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/api/types.go b/api/types.go
@@ -151,6 +151,7 @@ type Runner struct {
 	UseMMap   bool `json:"use_mmap,omitempty"`
 	UseMLock  bool `json:"use_mlock,omitempty"`
 	NumThread int  `json:"num_thread,omitempty"`
+	FlashAttn bool `json:"flash_attn,omitempty"`
 
 	// Unused: RopeFrequencyBase is ignored. Instead the value in the model will be used
 	RopeFrequencyBase float32 `json:"rope_frequency_base,omitempty"`
@@ -428,6 +429,7 @@ func DefaultOptions() Options {
 			UseMLock:  false,
 			UseMMap:   true,
 			UseNUMA:   false,
+			FlashAttn: false, // for CPU only compatibility
 		},
 	}
 }

diff --git a/llm/ext_server/server.cpp b/llm/ext_server/server.cpp
@@ -2508,7 +2508,8 @@ static void server_params_parse(int argc, char **argv, server_params &sparams,
         {
             params.use_mmap = false;
         }
-        else if (arg == "--numa") {
+        else if (arg == "--numa")
+        {
             if (++i >= argc) {
                 invalid_param = true;
                 break;
@@ -2540,15 +2541,17 @@ static void server_params_parse(int argc, char **argv, server_params &sparams,
                 break;
             }
             params.n_parallel = std::stoi(argv[i]);
-        } else if (arg == "-n" || arg == "--n-predict")
+        }
+        else if (arg == "-n" || arg == "--n-predict")
         {
             if (++i >= argc)
             {
                 invalid_param = true;
                 break;
             }
             params.n_predict = std::stoi(argv[i]);
-        } else if (arg == "-spf" || arg == "--system-prompt-file")
+        }
+        else if (arg == "-spf" || arg == "--system-prompt-file")
         {
             if (++i >= argc)
             {