update tests

qwen3-coder: fix tool definition type rendering

.gitignore

@@ -6,6 +6,7 @@
 dist
 build
 .cache
+.gocache
 *.exe
 .idea
 test_data

CMakeLists.txt

@@ -98,14 +98,17 @@ check_language(HIP)
 if(CMAKE_HIP_COMPILER)
     set(HIP_PLATFORM "amd")
 
-    find_package(hip REQUIRED)
     if(NOT AMDGPU_TARGETS)
+        find_package(hip REQUIRED)
         list(FILTER AMDGPU_TARGETS INCLUDE REGEX "^gfx(900|94[012]|101[02]|1030|110[012]|120[01])$")
-    elseif(WIN32 AND WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX)
+    endif()
+
+    if(WIN32 AND WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX)
         list(FILTER AMDGPU_TARGETS EXCLUDE REGEX ${WINDOWS_AMDGPU_TARGETS_EXCLUDE_REGEX})
     endif()
 
     if(AMDGPU_TARGETS)
+        find_package(hip REQUIRED)
         add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ml/backend/ggml/ggml/src/ggml-hip)
 
         if (WIN32)

api/client.go

@@ -45,6 +45,12 @@ func checkError(resp *http.Response, body []byte) error {
 		return nil
 	}
 
+	if resp.StatusCode == http.StatusUnauthorized {
+		authError := AuthorizationError{StatusCode: resp.StatusCode}
+		json.Unmarshal(body, &authError)
+		return authError
+	}
+
 	apiError := StatusError{StatusCode: resp.StatusCode}
 
 	err := json.Unmarshal(body, &apiError)
@@ -214,7 +220,8 @@ func (c *Client) stream(ctx context.Context, method, path string, data any, fn f
 	scanner.Buffer(scanBuf, maxBufferSize)
 	for scanner.Scan() {
 		var errorResponse struct {
-			Error string `json:"error,omitempty"`
+			Error     string `json:"error,omitempty"`
+			SigninURL string `json:"signin_url,omitempty"`
 		}
 
 		bts := scanner.Bytes()
@@ -223,14 +230,10 @@ func (c *Client) stream(ctx context.Context, method, path string, data any, fn f
 		}
 
 		if response.StatusCode == http.StatusUnauthorized {
-			pubKey, pkErr := auth.GetPublicKey()
-			if pkErr != nil {
-				return pkErr
-			}
 			return AuthorizationError{
 				StatusCode: response.StatusCode,
 				Status:     response.Status,
-				PublicKey:  pubKey,
+				SigninURL:  errorResponse.SigninURL,
 			}
 		} else if response.StatusCode >= http.StatusBadRequest {
 			return StatusError{
@@ -411,6 +414,24 @@ func (c *Client) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse,
 	return &resp, nil
 }
 
+// WebSearch performs a web search via the Ollama server.
+func (c *Client) WebSearch(ctx context.Context, req *WebSearchRequest) (*WebSearchResponse, error) {
+	var resp WebSearchResponse
+	if err := c.do(ctx, http.MethodPost, "/api/web_search", req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+// WebFetch fetches the contents of a web page via the Ollama server.
+func (c *Client) WebFetch(ctx context.Context, req *WebFetchRequest) (*WebFetchResponse, error) {
+	var resp WebFetchResponse
+	if err := c.do(ctx, http.MethodPost, "/api/web_fetch", req, &resp); err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
 // Embeddings generates an embedding from a model.
 func (c *Client) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) {
 	var resp EmbeddingResponse
@@ -439,8 +460,13 @@ func (c *Client) Version(ctx context.Context) (string, error) {
 	return version.Version, nil
 }
 
-// Signout will disconnect an ollama instance from ollama.com
-func (c *Client) Signout(ctx context.Context, encodedKey string) error {
+// Signout will signout a client for a local ollama server.
+func (c *Client) Signout(ctx context.Context) error {
+	return c.do(ctx, http.MethodPost, "/api/signout", nil, nil)
+}
+
+// Disconnect will disconnect an ollama instance from ollama.com.
+func (c *Client) Disconnect(ctx context.Context, encodedKey string) error {
 	return c.do(ctx, http.MethodDelete, fmt.Sprintf("/api/user/keys/%s", encodedKey), nil, nil)
 }
 

api/client_test.go

@@ -262,3 +262,135 @@ func TestClientDo(t *testing.T) {
 		})
 	}
 }
+
+func TestClientWebSearch(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			t.Fatalf("expected POST, got %s", r.Method)
+		}
+		if !strings.HasSuffix(r.URL.Path, "/api/web_search") {
+			t.Fatalf("unexpected path: %s", r.URL.Path)
+		}
+
+		var req WebSearchRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			t.Fatalf("failed to decode request: %v", err)
+		}
+
+		if req.Query != "what is ollama" {
+			t.Fatalf("unexpected query: %s", req.Query)
+		}
+
+		if req.MaxResults != 3 {
+			t.Fatalf("unexpected max_results: %d", req.MaxResults)
+		}
+
+		resp := WebSearchResponse{
+			Results: []WebSearchResult{{
+				Title:   "Ollama",
+				URL:     "https://ollama.com",
+				Content: "Cloud models are now available...",
+			}},
+		}
+		if err := json.NewEncoder(w).Encode(resp); err != nil {
+			t.Fatalf("failed to encode response: %v", err)
+		}
+	}))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Fatalf("parse server URL: %v", err)
+	}
+
+	client := NewClient(u, ts.Client())
+
+	resp, err := client.WebSearch(t.Context(), &WebSearchRequest{Query: "what is ollama", MaxResults: 3})
+	if err != nil {
+		t.Fatalf("WebSearch returned error: %v", err)
+	}
+
+	if len(resp.Results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(resp.Results))
+	}
+
+	if resp.Results[0].Title != "Ollama" {
+		t.Fatalf("unexpected title: %s", resp.Results[0].Title)
+	}
+}
+
+func TestClientWebSearchUnauthorized(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusUnauthorized)
+		_ = json.NewEncoder(w).Encode(map[string]string{
+			"signin_url": "https://ollama.com/connect",
+		})
+	}))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Fatalf("parse server URL: %v", err)
+	}
+
+	client := NewClient(u, ts.Client())
+
+	_, err = client.WebSearch(t.Context(), &WebSearchRequest{Query: "what is ollama"})
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+
+	if _, ok := err.(AuthorizationError); !ok {
+		t.Fatalf("expected AuthorizationError, got %T", err)
+	}
+}
+
+func TestClientWebFetch(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			t.Fatalf("expected POST, got %s", r.Method)
+		}
+		if !strings.HasSuffix(r.URL.Path, "/api/web_fetch") {
+			t.Fatalf("unexpected path: %s", r.URL.Path)
+		}
+
+		var req WebFetchRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			t.Fatalf("failed to decode request: %v", err)
+		}
+
+		if req.URL != "https://ollama.com" {
+			t.Fatalf("unexpected url: %s", req.URL)
+		}
+
+		resp := WebFetchResponse{
+			Title:   "Ollama",
+			Content: "Cloud models are now available...",
+			Links:   []string{"https://ollama.com/models"},
+		}
+		if err := json.NewEncoder(w).Encode(resp); err != nil {
+			t.Fatalf("failed to encode response: %v", err)
+		}
+	}))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Fatalf("parse server URL: %v", err)
+	}
+
+	client := NewClient(u, ts.Client())
+
+	resp, err := client.WebFetch(t.Context(), &WebFetchRequest{URL: "https://ollama.com"})
+	if err != nil {
+		t.Fatalf("WebFetch returned error: %v", err)
+	}
+
+	if resp.Title != "Ollama" {
+		t.Fatalf("unexpected title: %s", resp.Title)
+	}
+
+	if len(resp.Links) != 1 || resp.Links[0] != "https://ollama.com/models" {
+		t.Fatalf("unexpected links: %v", resp.Links)
+	}
+}

api/types.go

@@ -41,7 +41,7 @@ func (e StatusError) Error() string {
 type AuthorizationError struct {
 	StatusCode int
 	Status     string
-	PublicKey  string `json:"public_key"`
+	SigninURL  string `json:"signin_url"`
 }
 
 func (e AuthorizationError) Error() string {
@@ -453,6 +453,40 @@ type EmbeddingResponse struct {
 	Embedding []float64 `json:"embedding"`
 }
 
+// WebSearchRequest is the request passed to [Client.WebSearch].
+type WebSearchRequest struct {
+	// Query is the search query string.
+	Query string `json:"query"`
+
+	// MaxResults is the optional maximum number of results to return (default 5, max 10).
+	MaxResults int `json:"max_results,omitempty"`
+}
+
+// WebSearchResult represents a single web search result.
+type WebSearchResult struct {
+	Title   string `json:"title"`
+	URL     string `json:"url"`
+	Content string `json:"content"`
+}
+
+// WebSearchResponse is the response from [Client.WebSearch].
+type WebSearchResponse struct {
+	Results []WebSearchResult `json:"results"`
+}
+
+// WebFetchRequest is the request passed to [Client.WebFetch].
+type WebFetchRequest struct {
+	// URL is the address of the page to fetch.
+	URL string `json:"url"`
+}
+
+// WebFetchResponse is the response from [Client.WebFetch].
+type WebFetchResponse struct {
+	Title   string   `json:"title"`
+	Content string   `json:"content"`
+	Links   []string `json:"links"`
+}
+
 // CreateRequest is the request passed to [Client.Create].
 type CreateRequest struct {
 	// Model is the model name to create.

auth/auth.go

@@ -18,46 +18,13 @@ import (
 
 const defaultPrivateKey = "id_ed25519"
 
-func keyPath() (string, error) {
-	fileIsReadable := func(fp string) bool {
-		info, err := os.Stat(fp)
-		if err != nil {
-			return false
-		}
-
-		// Check that it's a regular file, not a directory or other file type
-		if !info.Mode().IsRegular() {
-			return false
-		}
-
-		// Try to open it to check readability
-		file, err := os.Open(fp)
-		if err != nil {
-			return false
-		}
-		file.Close()
-		return true
-	}
-
-	systemPath := filepath.Join("/usr/share/ollama/.ollama", defaultPrivateKey)
-	if fileIsReadable(systemPath) {
-		return systemPath, nil
-	}
-
+func GetPublicKey() (string, error) {
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", err
 	}
 
-	return filepath.Join(home, ".ollama", defaultPrivateKey), nil
-}
-
-func GetPublicKey() (string, error) {
-	keyPath, err := keyPath()
-	if err != nil {
-		return "", err
-	}
-
+	keyPath := filepath.Join(home, ".ollama", defaultPrivateKey)
 	privateKeyFile, err := os.ReadFile(keyPath)
 	if err != nil {
 		slog.Info(fmt.Sprintf("Failed to load private key: %v", err))
@@ -84,11 +51,12 @@ func NewNonce(r io.Reader, length int) (string, error) {
 }
 
 func Sign(ctx context.Context, bts []byte) (string, error) {
-	keyPath, err := keyPath()
+	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", err
 	}
 
+	keyPath := filepath.Join(home, ".ollama", defaultPrivateKey)
 	privateKeyFile, err := os.ReadFile(keyPath)
 	if err != nil {
 		slog.Info(fmt.Sprintf("Failed to load private key: %v", err))

cmd/cmd.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"crypto/ed25519"
 	"crypto/rand"
-	"encoding/base64"
 	"encoding/json"
 	"encoding/pem"
 	"errors"
@@ -15,7 +14,6 @@ import (
 	"math"
 	"net"
 	"net/http"
-	"net/url"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -37,7 +35,6 @@ import (
 	"golang.org/x/term"
 
 	"github.com/ollama/ollama/api"
-	"github.com/ollama/ollama/auth"
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/format"
 	"github.com/ollama/ollama/parser"
@@ -50,7 +47,7 @@ import (
 	"github.com/ollama/ollama/version"
 )
 
-const ConnectInstructions = "To sign in, navigate to:\n    https://ollama.com/connect?name=%s&key=%s\n\n"
+const ConnectInstructions = "To sign in, navigate to:\n    %s\n\n"
 
 // ensureThinkingSupport emits a warning if the model does not advertise thinking support
 func ensureThinkingSupport(ctx context.Context, client *api.Client, name string) {
@@ -452,16 +449,10 @@ func RunHandler(cmd *cobra.Command, args []string) error {
 		if err := loadOrUnloadModel(cmd, &opts); err != nil {
 			var sErr api.AuthorizationError
 			if errors.As(err, &sErr) && sErr.StatusCode == http.StatusUnauthorized {
-				pubKey, pkErr := auth.GetPublicKey()
-				if pkErr != nil {
-					return pkErr
-				}
-				// the server and the client both have the same public key
-				if pubKey == sErr.PublicKey {
-					h, _ := os.Hostname()
-					encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
-					fmt.Printf("You need to be signed in to Ollama to run Cloud models.\n\n")
-					fmt.Printf(ConnectInstructions, url.PathEscape(h), encKey)
+				fmt.Printf("You need to be signed in to Ollama to run Cloud models.\n\n")
+
+				if sErr.SigninURL != "" {
+					fmt.Printf(ConnectInstructions, sErr.SigninURL)
 				}
 				return nil
 			}
@@ -493,6 +484,16 @@ func SigninHandler(cmd *cobra.Command, args []string) error {
 
 	user, err := client.Whoami(cmd.Context())
 	if err != nil {
+		var aErr api.AuthorizationError
+		if errors.As(err, &aErr) && aErr.StatusCode == http.StatusUnauthorized {
+			fmt.Println("You need to be signed in to Ollama to run Cloud models.")
+			fmt.Println()
+
+			if aErr.SigninURL != "" {
+				fmt.Printf(ConnectInstructions, aErr.SigninURL)
+			}
+			return nil
+		}
 		return err
 	}
 
@@ -502,34 +503,27 @@ func SigninHandler(cmd *cobra.Command, args []string) error {
 		return nil
 	}
 
-	pubKey, pkErr := auth.GetPublicKey()
-	if pkErr != nil {
-		return pkErr
-	}
-	encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
-
-	h, _ := os.Hostname()
-	fmt.Printf(ConnectInstructions, url.PathEscape(h), encKey)
-
 	return nil
 }
 
 func SignoutHandler(cmd *cobra.Command, args []string) error {
-	pubKey, pkErr := auth.GetPublicKey()
-	if pkErr != nil {
-		return pkErr
-	}
-	encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
-
 	client, err := api.ClientFromEnvironment()
 	if err != nil {
 		return err
 	}
 
-	err = client.Signout(cmd.Context(), encKey)
+	err = client.Signout(cmd.Context())
 	if err != nil {
-		return err
+		var aErr api.AuthorizationError
+		if errors.As(err, &aErr) && aErr.StatusCode == http.StatusUnauthorized {
+			fmt.Println("You are not signed in to ollama.com")
+			fmt.Println()
+			return nil
+		} else {
+			return err
+		}
 	}
+
 	fmt.Println("You have signed out of ollama.com")
 	fmt.Println()
 	return nil
@@ -546,6 +540,25 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	n := model.ParseName(args[0])
+	if strings.HasSuffix(n.Host, ".ollama.ai") || strings.HasSuffix(n.Host, ".ollama.com") {
+		_, err := client.Whoami(cmd.Context())
+		if err != nil {
+			var aErr api.AuthorizationError
+			if errors.As(err, &aErr) && aErr.StatusCode == http.StatusUnauthorized {
+				fmt.Println("You need to be signed in to push models to ollama.com.")
+				fmt.Println()
+
+				if aErr.SigninURL != "" {
+					fmt.Printf(ConnectInstructions, aErr.SigninURL)
+				}
+				return nil
+			}
+
+			return err
+		}
+	}
+
 	p := progress.NewProgress(os.Stderr)
 	defer p.Stop()
 
@@ -582,7 +595,6 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 
 	request := api.PushRequest{Name: args[0], Insecure: insecure}
 
-	n := model.ParseName(args[0])
 	if err := client.Push(cmd.Context(), &request, fn); err != nil {
 		if spinner != nil {
 			spinner.Stop()
@@ -1106,6 +1118,51 @@ type runOptions struct {
 	ShowConnect  bool
 }
 
+func (r runOptions) Copy() runOptions {
+	var messages []api.Message
+	if r.Messages != nil {
+		messages = make([]api.Message, len(r.Messages))
+		copy(messages, r.Messages)
+	}
+
+	var images []api.ImageData
+	if r.Images != nil {
+		images = make([]api.ImageData, len(r.Images))
+		copy(images, r.Images)
+	}
+
+	var opts map[string]any
+	if r.Options != nil {
+		opts = make(map[string]any, len(r.Options))
+		for k, v := range r.Options {
+			opts[k] = v
+		}
+	}
+
+	var think *api.ThinkValue
+	if r.Think != nil {
+		cThink := *r.Think
+		think = &cThink
+	}
+
+	return runOptions{
+		Model:        r.Model,
+		ParentModel:  r.ParentModel,
+		Prompt:       r.Prompt,
+		Messages:     messages,
+		WordWrap:     r.WordWrap,
+		Format:       r.Format,
+		System:       r.System,
+		Images:       images,
+		Options:      opts,
+		MultiModal:   r.MultiModal,
+		KeepAlive:    r.KeepAlive,
+		Think:        think,
+		HideThinking: r.HideThinking,
+		ShowConnect:  r.ShowConnect,
+	}
+}
+
 type displayResponseState struct {
 	lineLength int
 	wordBuffer string

cmd/cmd_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -491,9 +492,35 @@ func TestPushHandler(t *testing.T) {
 						w.(http.Flusher).Flush()
 					}
 				},
+				"/api/me": func(w http.ResponseWriter, r *http.Request) {
+					if r.Method != http.MethodPost {
+						t.Errorf("expected POST request, got %s", r.Method)
+					}
+				},
 			},
 			expectedOutput: "\nYou can find your model at:\n\n\thttps://ollama.com/test-model\n",
 		},
+		{
+			name:      "not signed in push",
+			modelName: "notsignedin-model",
+			serverResponse: map[string]func(w http.ResponseWriter, r *http.Request){
+				"/api/me": func(w http.ResponseWriter, r *http.Request) {
+					if r.Method != http.MethodPost {
+						t.Errorf("expected POST request, got %s", r.Method)
+					}
+					w.Header().Set("Content-Type", "application/json")
+					w.WriteHeader(http.StatusUnauthorized)
+					err := json.NewEncoder(w).Encode(map[string]string{
+						"error":      "unauthorized",
+						"signin_url": "https://somethingsomething",
+					})
+					if err != nil {
+						t.Fatal(err)
+					}
+				},
+			},
+			expectedOutput: "You need to be signed in to push",
+		},
 		{
 			name:      "unauthorized push",
 			modelName: "unauthorized-model",
@@ -508,6 +535,11 @@ func TestPushHandler(t *testing.T) {
 						t.Fatal(err)
 					}
 				},
+				"/api/me": func(w http.ResponseWriter, r *http.Request) {
+					if r.Method != http.MethodPost {
+						t.Errorf("expected POST request, got %s", r.Method)
+					}
+				},
 			},
 			expectedError: "you are not authorized to push to this namespace, create the model under a namespace you own",
 		},
@@ -525,6 +557,9 @@ func TestPushHandler(t *testing.T) {
 			defer mockServer.Close()
 
 			t.Setenv("OLLAMA_HOST", mockServer.URL)
+			tmpDir := t.TempDir()
+			t.Setenv("HOME", tmpDir)
+			t.Setenv("USERPROFILE", tmpDir)
 			initializeKeypair()
 
 			cmd := &cobra.Command{}
@@ -561,7 +596,7 @@ func TestPushHandler(t *testing.T) {
 					t.Errorf("expected no error, got %v", err)
 				}
 				if tt.expectedOutput != "" {
-					if got := string(stdout); got != tt.expectedOutput {
+					if got := string(stdout); !strings.Contains(got, tt.expectedOutput) {
 						t.Errorf("expected output %q, got %q", tt.expectedOutput, got)
 					}
 				}
@@ -919,3 +954,286 @@ func TestNewCreateRequest(t *testing.T) {
 		})
 	}
 }
+
+func TestRunOptions_Copy(t *testing.T) {
+	// Setup test data
+	originalKeepAlive := &api.Duration{Duration: 5 * time.Minute}
+	originalThink := &api.ThinkValue{Value: "test reasoning"}
+
+	original := runOptions{
+		Model:       "test-model",
+		ParentModel: "parent-model",
+		Prompt:      "test prompt",
+		Messages: []api.Message{
+			{Role: "user", Content: "hello"},
+			{Role: "assistant", Content: "hi there"},
+		},
+		WordWrap: true,
+		Format:   "json",
+		System:   "system prompt",
+		Images: []api.ImageData{
+			[]byte("image1"),
+			[]byte("image2"),
+		},
+		Options: map[string]any{
+			"temperature": 0.7,
+			"max_tokens":  1000,
+			"top_p":       0.9,
+		},
+		MultiModal:   true,
+		KeepAlive:    originalKeepAlive,
+		Think:        originalThink,
+		HideThinking: false,
+		ShowConnect:  true,
+	}
+
+	// Test the copy
+	copied := original.Copy()
+
+	// Test 1: Verify the copy is not the same instance
+	if &copied == &original {
+		t.Error("Copy should return a different instance")
+	}
+
+	// Test 2: Verify all fields are copied correctly
+	tests := []struct {
+		name string
+		got  interface{}
+		want interface{}
+	}{
+		{"Model", copied.Model, original.Model},
+		{"ParentModel", copied.ParentModel, original.ParentModel},
+		{"Prompt", copied.Prompt, original.Prompt},
+		{"WordWrap", copied.WordWrap, original.WordWrap},
+		{"Format", copied.Format, original.Format},
+		{"System", copied.System, original.System},
+		{"MultiModal", copied.MultiModal, original.MultiModal},
+		{"HideThinking", copied.HideThinking, original.HideThinking},
+		{"ShowConnect", copied.ShowConnect, original.ShowConnect},
+	}
+
+	for _, tt := range tests {
+		if !reflect.DeepEqual(tt.got, tt.want) {
+			t.Errorf("%s mismatch: got %v, want %v", tt.name, tt.got, tt.want)
+		}
+	}
+
+	// Test 3: Verify Messages slice is deeply copied
+	if len(copied.Messages) != len(original.Messages) {
+		t.Errorf("Messages length mismatch: got %d, want %d", len(copied.Messages), len(original.Messages))
+	}
+
+	if len(copied.Messages) > 0 && &copied.Messages[0] == &original.Messages[0] {
+		t.Error("Messages should be different instances")
+	}
+
+	// Modify original to verify independence
+	if len(original.Messages) > 0 {
+		originalContent := original.Messages[0].Content
+		original.Messages[0].Content = "modified"
+		if len(copied.Messages) > 0 && copied.Messages[0].Content == "modified" {
+			t.Error("Messages should be independent after copy")
+		}
+		// Restore for other tests
+		original.Messages[0].Content = originalContent
+	}
+
+	// Test 4: Verify Images slice is deeply copied
+	if len(copied.Images) != len(original.Images) {
+		t.Errorf("Images length mismatch: got %d, want %d", len(copied.Images), len(original.Images))
+	}
+
+	if len(copied.Images) > 0 && &copied.Images[0] == &original.Images[0] {
+		t.Error("Images should be different instances")
+	}
+
+	// Modify original to verify independence
+	if len(original.Images) > 0 {
+		originalImage := original.Images[0]
+		original.Images[0] = []byte("modified")
+		if len(copied.Images) > 0 && string(copied.Images[0]) == "modified" {
+			t.Error("Images should be independent after copy")
+		}
+		// Restore for other tests
+		original.Images[0] = originalImage
+	}
+
+	// Test 5: Verify Options map is deeply copied
+	if len(copied.Options) != len(original.Options) {
+		t.Errorf("Options length mismatch: got %d, want %d", len(copied.Options), len(original.Options))
+	}
+
+	if len(copied.Options) > 0 && &copied.Options == &original.Options {
+		t.Error("Options map should be different instances")
+	}
+
+	// Modify original to verify independence
+	if len(original.Options) > 0 {
+		originalTemp := original.Options["temperature"]
+		original.Options["temperature"] = 0.9
+		if copied.Options["temperature"] == 0.9 {
+			t.Error("Options should be independent after copy")
+		}
+		// Restore for other tests
+		original.Options["temperature"] = originalTemp
+	}
+
+	// Test 6: Verify KeepAlive pointer is copied (shallow copy)
+	if copied.KeepAlive != original.KeepAlive {
+		t.Error("KeepAlive pointer should be the same (shallow copy)")
+	}
+
+	// Test 7: Verify Think pointer creates a new instance
+	if original.Think != nil && copied.Think == original.Think {
+		t.Error("Think should be a different instance")
+	}
+
+	if original.Think != nil && copied.Think != nil {
+		if !reflect.DeepEqual(copied.Think.Value, original.Think.Value) {
+			t.Errorf("Think.Value mismatch: got %v, want %v", copied.Think.Value, original.Think.Value)
+		}
+	}
+
+	// Test 8: Test with zero values
+	zeroOriginal := runOptions{}
+	zeroCopy := zeroOriginal.Copy()
+
+	if !reflect.DeepEqual(zeroCopy, zeroOriginal) {
+		fmt.Printf("orig: %#v\ncopy: %#v\n", zeroOriginal, zeroCopy)
+		t.Error("Copy of zero value should equal original zero value")
+	}
+}
+
+func TestRunOptions_Copy_EmptySlicesAndMaps(t *testing.T) {
+	// Test with empty slices and maps
+	original := runOptions{
+		Messages: []api.Message{},
+		Images:   []api.ImageData{},
+		Options:  map[string]any{},
+	}
+
+	copied := original.Copy()
+
+	if copied.Messages == nil {
+		t.Error("Empty Messages slice should remain empty, not nil")
+	}
+
+	if copied.Images == nil {
+		t.Error("Empty Images slice should remain empty, not nil")
+	}
+
+	if copied.Options == nil {
+		t.Error("Empty Options map should remain empty, not nil")
+	}
+
+	if len(copied.Messages) != 0 {
+		t.Error("Empty Messages slice should remain empty")
+	}
+
+	if len(copied.Images) != 0 {
+		t.Error("Empty Images slice should remain empty")
+	}
+
+	if len(copied.Options) != 0 {
+		t.Error("Empty Options map should remain empty")
+	}
+}
+
+func TestRunOptions_Copy_NilPointers(t *testing.T) {
+	// Test with nil pointers
+	original := runOptions{
+		KeepAlive: nil,
+		Think:     nil,
+	}
+
+	copied := original.Copy()
+
+	if copied.KeepAlive != nil {
+		t.Error("Nil KeepAlive should remain nil")
+	}
+
+	if copied.Think != nil {
+		t.Error("Nil Think should remain nil")
+	}
+}
+
+func TestRunOptions_Copy_ThinkValueVariants(t *testing.T) {
+	tests := []struct {
+		name  string
+		think *api.ThinkValue
+	}{
+		{"nil Think", nil},
+		{"bool true", &api.ThinkValue{Value: true}},
+		{"bool false", &api.ThinkValue{Value: false}},
+		{"string value", &api.ThinkValue{Value: "reasoning text"}},
+		{"int value", &api.ThinkValue{Value: 42}},
+		{"nil value", &api.ThinkValue{Value: nil}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			original := runOptions{Think: tt.think}
+			copied := original.Copy()
+
+			if tt.think == nil {
+				if copied.Think != nil {
+					t.Error("Nil Think should remain nil")
+				}
+				return
+			}
+
+			if copied.Think == nil {
+				t.Error("Non-nil Think should not become nil")
+				return
+			}
+
+			if copied.Think == original.Think {
+				t.Error("Think should be a different instance")
+			}
+
+			if !reflect.DeepEqual(copied.Think.Value, original.Think.Value) {
+				t.Errorf("Think.Value mismatch: got %v, want %v", copied.Think.Value, original.Think.Value)
+			}
+		})
+	}
+}
+
+func TestRunOptions_Copy_Independence(t *testing.T) {
+	// Test that modifications to original don't affect copy
+	originalThink := &api.ThinkValue{Value: "original"}
+	original := runOptions{
+		Model:    "original-model",
+		Messages: []api.Message{{Role: "user", Content: "original"}},
+		Options:  map[string]any{"key": "value"},
+		Think:    originalThink,
+	}
+
+	copied := original.Copy()
+
+	// Modify original
+	original.Model = "modified-model"
+	if len(original.Messages) > 0 {
+		original.Messages[0].Content = "modified"
+	}
+	original.Options["key"] = "modified"
+	if original.Think != nil {
+		original.Think.Value = "modified"
+	}
+
+	// Verify copy is unchanged
+	if copied.Model == "modified-model" {
+		t.Error("Copy Model should not be affected by original modification")
+	}
+
+	if len(copied.Messages) > 0 && copied.Messages[0].Content == "modified" {
+		t.Error("Copy Messages should not be affected by original modification")
+	}
+
+	if copied.Options["key"] == "modified" {
+		t.Error("Copy Options should not be affected by original modification")
+	}
+
+	if copied.Think != nil && copied.Think.Value == "modified" {
+		t.Error("Copy Think should not be affected by original modification")
+	}
+}

cmd/interactive.go

@@ -195,16 +195,24 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 				fmt.Println("Usage:\n  /load <modelname>")
 				continue
 			}
+			origOpts := opts.Copy()
+
 			opts.Model = args[1]
 			opts.Messages = []api.Message{}
 			fmt.Printf("Loading model '%s'\n", opts.Model)
 			opts.Think, err = inferThinkingOption(nil, &opts, thinkExplicitlySet)
 			if err != nil {
+				if strings.Contains(err.Error(), "not found") {
+					fmt.Printf("Couldn't find model '%s'\n", opts.Model)
+					opts = origOpts.Copy()
+					continue
+				}
 				return err
 			}
 			if err := loadOrUnloadModel(cmd, &opts); err != nil {
 				if strings.Contains(err.Error(), "not found") {
-					fmt.Printf("error: %v\n", err)
+					fmt.Printf("Couldn't find model '%s'\n", opts.Model)
+					opts = origOpts.Copy()
 					continue
 				}
 				if strings.Contains(err.Error(), "does not support thinking") {

docs/cloud.md

@@ -0,0 +1,40 @@
+# Cloud
+
+| Ollama's cloud is currently in preview. For full documentation, see [Ollama's documentation](https://docs.ollama.com/cloud).
+
+## Cloud Models
+
+[Cloud models](https://ollama.com/cloud) are a new kind of model in Ollama that can run without a powerful GPU. Instead, cloud models are automatically offloaded to Ollama's cloud while offering the same capabilities as local models, making it possible to keep using your local tools while running larger models that wouldn’t fit on a personal computer.
+
+Ollama currently supports the following cloud models, with more coming soon:
+
+- `gpt-oss:20b-cloud`
+- `gpt-oss:120b-cloud`
+- `deepseek-v3.1:671b-cloud`
+- `qwen3-coder:480b-cloud`
+
+### Get started
+
+To run a cloud model, open the terminal and run:
+
+```
+ollama run gpt-oss:120b-cloud
+```
+
+To run cloud models with integrations that work with Ollama, first download the cloud model:
+
+```
+ollama pull qwen3-coder:480b-cloud
+```
+
+Then sign in to Ollama:
+
+```
+ollama signin
+```
+
+Finally, access the model using the model name `qwen3-coder:480b-cloud` via Ollama's local API or tooling.
+
+## Cloud API access
+
+Cloud models can also be accessed directly on ollama.com's API. For more information, see the [docs](https://docs.ollama.com/cloud).

docs/turbo.md

@@ -1,107 +0,0 @@
-# Turbo
-
-> ⚠️ Turbo is preview
-
-Ollama’s [Turbo](https://ollama.com/turbo) is a new way to run open-source models with acceleration from datacenter-grade hardware.
-
-Currently, the following models are available in Turbo:
-
-- `gpt-oss:20b`
-- `gpt-oss:120b`
-
-## Get started
-
-### Ollama for macOS & Windows
-
-Download Ollama
-
-- Select a model such as `gpt-oss:20b` or `gpt-oss:120b`
-- Click on **Turbo**. You’ll be prompted to create an account or sign in
-
-### Ollama’s CLI
-
-- [Sign up](https://ollama.com/signup) for an Ollama account
-- Add your Ollama key [to ollama.com](https://ollama.com/settings/keys).
-
-  On macOS and Linux:
-
-  ```shell
-  cat ~/.ollama/id_ed25519.pub
-  ```
-
-  On Windows:
-
-  ```
-  type "%USERPROFILE%\.ollama\id_ed25519.pub"
-  ```
-
-- Then run a model setting `OLLAMA_HOST` to `ollama.com`:
-  ```shell
-  OLLAMA_HOST=ollama.com ollama run gpt-oss:120b
-  ```
-
-### Ollama’s Python library
-
-- Download Ollama's [Python library](https://github.com/ollama/ollama-python)
-- [Sign up](https://ollama.com/signup) for an Ollama account
-- Create an API key by visiting https://ollama.com/settings/keys
-
-```python
-from ollama import Client
-
-client = Client(
-    host="https://ollama.com",
-    headers={'Authorization': '<api key>'}
-)
-
-messages = [
-  {
-    'role': 'user',
-    'content': 'Why is the sky blue?',
-  },
-]
-
-for part in client.chat('gpt-oss:120b', messages=messages, stream=True):
-  print(part['message']['content'], end='', flush=True)
-```
-
-### Ollama’s JavaScript library
-
-- Download Ollama's [JavaScript library](https://github.com/ollama/ollama-js)
-- [Sign up](https://ollama.com/signup) for an Ollama account
-- Create an API key by visiting https://ollama.com/settings/keys
-
-```typescript
-import { Ollama } from 'ollama';
-
-const ollama = new Ollama({
-  host: 'https://ollama.com',
-  headers: {
-	  Authorization: "Bearer <api key>"
-  }
-});
-
-const response = await ollama.chat({
-  model: 'gpt-oss:120b',
-  messages: [{ role: 'user', content: 'Explain quantum computing' }],
-  stream: true
-});
-
-for await (const part of response) {
-    process.stdout.write(part.message.content)
-}
-```
-
-### Community integrations
-
-Turbo mode is also compatible with several community integrations.
-
-#### Open WebUI
-
-- Go to **settings** → **Admin settings** → **Connections**
-- Under **Ollama API,** click **+**
-- For the **URL** put `https://ollama.com`
-- For the **API key,** create an API key on https://ollama.com/settings/keys and add it.
-- Click **Save**
-
-Now, if you navigate to the model selector, Turbo models should be available under **External**.

fs/ggml/ggml.go

@@ -244,6 +244,7 @@ func (kv KV) OllamaEngineRequired() bool {
 		"gemma3n",
 		"mistral3",
 		"qwen3",
+		"qwen3moe",
 		"llama4",
 		"mllama",
 		"qwen25vl",

harmony/harmonyparser.go

@@ -1,6 +1,7 @@
 package harmony
 
 import (
+	"encoding/json"
 	"fmt"
 	"log/slog"
 	"strings"
@@ -265,6 +266,8 @@ type HarmonyMessageHandler struct {
 	state           harmonyMessageState
 	HarmonyParser   *HarmonyParser
 	FunctionNameMap *FunctionNameMap
+	toolAccumulator *HarmonyToolCallAccumulator
+	convertedTools  map[string]struct{}
 }
 
 // NewHarmonyMessageHandler creates a new message handler
@@ -277,6 +280,7 @@ func NewHarmonyMessageHandler() *HarmonyMessageHandler {
 			HeaderEndTag:    "<|message|>",
 		},
 		FunctionNameMap: NewFunctionNameMap(),
+		convertedTools:  make(map[string]struct{}),
 	}
 }
 
@@ -384,8 +388,85 @@ func NewFunctionNameMap() *FunctionNameMap {
 	}
 }
 
+// Init initializes the handler with tools and optional last message
+// Implements the Parser interface
+func (h *HarmonyMessageHandler) Init(tools []api.Tool, lastMessage *api.Message) []api.Tool {
+	// Initialize the harmony parser
+	if h.HarmonyParser == nil {
+		h.HarmonyParser = &HarmonyParser{
+			MessageStartTag: "<|start|>",
+			MessageEndTag:   "<|end|>",
+			HeaderEndTag:    "<|message|>",
+		}
+	}
+
+	// Handle prefill for chat mode
+	if lastMessage != nil {
+		h.HarmonyParser.AddImplicitStartOrPrefill(lastMessage)
+	} else {
+		h.HarmonyParser.AddImplicitStart()
+	}
+
+	// Initialize tool accumulator
+	h.toolAccumulator = h.CreateToolParser()
+
+	// Process tools and return renamed versions
+	if len(tools) == 0 {
+		return tools
+	}
+
+	processedTools := make([]api.Tool, len(tools))
+	copy(processedTools, tools)
+	for i, tool := range processedTools {
+		if tool.Function.Name != "" {
+			processedTools[i].Function.Name = h.FunctionNameMap.ConvertAndAdd(tool.Function.Name)
+			h.convertedTools[tool.Function.Name] = struct{}{}
+		}
+	}
+	return processedTools
+}
+
+// Add implements the Parser interface - processes streamed content and extracts content, thinking, and tool calls
+func (h *HarmonyMessageHandler) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
+	content, thinking, toolContent := h.AddContent(s, h.toolAccumulator)
+	if toolContent != "" {
+		h.toolAccumulator.Add(toolContent)
+	}
+
+	// tool calls always happen one at a time, and always at the end of a message,
+	// so for simplicity we defer parsing them until we know we're done
+	if done {
+		toolName, raw := h.toolAccumulator.Drain()
+		if toolName != nil {
+			name := strings.TrimPrefix(*toolName, "functions.")
+			name = h.FunctionNameMap.OriginalFromConverted(name)
+			var args api.ToolCallFunctionArguments
+			if err := json.Unmarshal([]byte(raw), &args); err != nil {
+				return "", "", nil, fmt.Errorf("error parsing tool call: raw='%s', err=%w", raw, err)
+			}
+			calls = append(calls, api.ToolCall{Function: api.ToolCallFunction{Name: name, Arguments: args}})
+		}
+	}
+
+	return content, thinking, calls, nil
+}
+
+// HasToolSupport implements the Parser interface
+func (h *HarmonyMessageHandler) HasToolSupport() bool {
+	return true
+}
+
+// HasThinkingSupport implements the Parser interface
+func (h *HarmonyMessageHandler) HasThinkingSupport() bool {
+	return true
+}
+
 func (m *FunctionNameMap) ConvertAndAdd(userFunctionName string) string {
 	harmonyFunctionName := m.deriveName(userFunctionName)
+	// built-in functions should not be renamed
+	if userFunctionName == "browser.open" || userFunctionName == "browser.search" || userFunctionName == "browser.find" || userFunctionName == "python" {
+		harmonyFunctionName = userFunctionName
+	}
 	m.userToHarmony[userFunctionName] = harmonyFunctionName
 	m.harmonyToUser[harmonyFunctionName] = userFunctionName
 	return harmonyFunctionName

harmony/harmonyparser_test.go

@@ -513,6 +513,7 @@ func TestFunctionConvertAndAdd(t *testing.T) {
 		{name: "dupes from different user-specified names", in: []string{"get weather", "get_weather", "get-weather"}, want: []string{"get_weather", "get_weather_2", "get_weather_3"}},
 		{name: "non dupes after dupes", in: []string{"get weather", "get_weather", "get-weather", "something-different"}, want: []string{"get_weather", "get_weather_2", "get_weather_3", "something_different"}},
 		{name: "multiple sets of dupes", in: []string{"a", "a", "b", "a", "a", "b", "a"}, want: []string{"a", "a_2", "b", "a_3", "a_4", "b_2", "a_5"}},
+		{name: "built-in functions should not be renamed", in: []string{"browser.open", "python", "not.a.built-in.function", "browser.not_a_real_built_in"}, want: []string{"browser.open", "python", "not_a_built_in_function", "browser_not_a_real_built_in"}},
 	}
 
 	for i, tt := range tests {

integration/README.md

@@ -12,3 +12,6 @@ The integration tests have 2 modes of operating.
 
 > [!IMPORTANT]
 > Before running the tests locally without the "test existing" setting, compile ollama from the top of the source tree  `go build .` in addition to GPU support with cmake if applicable on your platform.  The integration tests expect to find an ollama binary at the top of the tree.
+
+
+Many tests use a default small model suitable to run on many systems.  You can override this default model by setting `OLLAMA_TEST_DEFAULT_MODEL`

integration/api_test.go

@@ -22,13 +22,12 @@ func TestAPIGenerate(t *testing.T) {
 	// Set up the test data
 	req := api.GenerateRequest{
 		Model:  smol,
-		Prompt: "why is the sky blue? be brief",
+		Prompt: blueSkyPrompt,
 		Options: map[string]interface{}{
 			"temperature": 0,
 			"seed":        123,
 		},
 	}
-	anyResp := []string{"rayleigh", "scattering"}
 
 	client, _, cleanup := InitServerConnection(ctx, t)
 	defer cleanup()
@@ -120,14 +119,14 @@ func TestAPIGenerate(t *testing.T) {
 				// Verify the response contains the expected data
 				response := buf.String()
 				atLeastOne := false
-				for _, resp := range anyResp {
+				for _, resp := range blueSkyExpected {
 					if strings.Contains(strings.ToLower(response), resp) {
 						atLeastOne = true
 						break
 					}
 				}
 				if !atLeastOne {
-					t.Errorf("none of %v found in %s", anyResp, response)
+					t.Errorf("none of %v found in %s", blueSkyExpected, response)
 				}
 			case <-ctx.Done():
 				t.Error("outer test context done while waiting for generate")
@@ -181,7 +180,7 @@ func TestAPIChat(t *testing.T) {
 		Messages: []api.Message{
 			{
 				Role:    "user",
-				Content: "why is the sky blue?  be brief",
+				Content: blueSkyPrompt,
 			},
 		},
 		Options: map[string]interface{}{
@@ -189,7 +188,6 @@ func TestAPIChat(t *testing.T) {
 			"seed":        123,
 		},
 	}
-	anyResp := []string{"rayleigh", "scattering"}
 
 	client, _, cleanup := InitServerConnection(ctx, t)
 	defer cleanup()
@@ -279,14 +277,14 @@ func TestAPIChat(t *testing.T) {
 				// Verify the response contains the expected data
 				response := buf.String()
 				atLeastOne := false
-				for _, resp := range anyResp {
+				for _, resp := range blueSkyExpected {
 					if strings.Contains(strings.ToLower(response), resp) {
 						atLeastOne = true
 						break
 					}
 				}
 				if !atLeastOne {
-					t.Errorf("none of %v found in %s", anyResp, response)
+					t.Errorf("none of %v found in %s", blueSkyExpected, response)
 				}
 			case <-ctx.Done():
 				t.Error("outer test context done while waiting for chat")

integration/basic_test.go

@@ -19,14 +19,14 @@ func TestBlueSky(t *testing.T) {
 	// Set up the test data
 	req := api.GenerateRequest{
 		Model:  smol,
-		Prompt: "why is the sky blue?",
+		Prompt: blueSkyPrompt,
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
 			"seed":        123,
 		},
 	}
-	GenerateTestHelper(ctx, t, req, []string{"rayleigh", "scattering"})
+	GenerateTestHelper(ctx, t, req, blueSkyExpected)
 }
 
 func TestUnicode(t *testing.T) {
@@ -110,12 +110,12 @@ func TestUnicodeModelDir(t *testing.T) {
 
 	req := api.GenerateRequest{
 		Model:  smol,
-		Prompt: "why is the sky blue?",
+		Prompt: blueSkyPrompt,
 		Stream: &stream,
 		Options: map[string]any{
 			"temperature": 0,
 			"seed":        123,
 		},
 	}
-	GenerateTestHelper(ctx, t, req, []string{"rayleigh", "scattering"})
+	GenerateTestHelper(ctx, t, req, blueSkyExpected)
 }

integration/context_test.go

@@ -63,11 +63,11 @@ func TestContextExhaustion(t *testing.T) {
 	if err := PullIfMissing(ctx, client, req.Model); err != nil {
 		t.Fatalf("PullIfMissing failed: %v", err)
 	}
-	DoGenerate(ctx, t, client, req, []string{"once", "upon", "lived", "sunny", "cloudy", "clear", "water"}, 120*time.Second, 10*time.Second)
+	DoGenerate(ctx, t, client, req, []string{"once", "upon", "lived", "sunny", "cloudy", "clear", "water", "time", "travel", "world"}, 120*time.Second, 10*time.Second)
 }
 
 // Send multiple generate requests with prior context and ensure the response is coherant and expected
-func TestGenerateWithHistory(t *testing.T) {
+func TestParallelGenerateWithHistory(t *testing.T) {
 	modelOverride := ollamaEngineChatModels[0] // Most recent ollama engine model
 	req, resp := GenerateRequests()
 	numParallel := 2
@@ -113,8 +113,48 @@ func TestGenerateWithHistory(t *testing.T) {
 	wg.Wait()
 }
 
+// Send generate requests with prior context and ensure the response is coherant and expected
+func TestGenerateWithHistory(t *testing.T) {
+	req := api.GenerateRequest{
+		Model:     smol,
+		Prompt:    rainbowPrompt,
+		Stream:    &stream,
+		KeepAlive: &api.Duration{Duration: 10 * time.Second},
+		Options: map[string]any{
+			"num_ctx": 16384,
+		},
+	}
+
+	softTimeout, hardTimeout := getTimeouts(t)
+	ctx, cancel := context.WithTimeout(context.Background(), hardTimeout)
+	defer cancel()
+	client, _, cleanup := InitServerConnection(ctx, t)
+	defer cleanup()
+
+	// Get the server running (if applicable) warm the model up with a single initial request
+	slog.Info("loading", "model", req.Model)
+	err := client.Generate(ctx,
+		&api.GenerateRequest{Model: req.Model, KeepAlive: &api.Duration{Duration: 10 * time.Second}, Options: req.Options},
+		func(response api.GenerateResponse) error { return nil },
+	)
+	if err != nil {
+		t.Fatalf("failed to load model %s: %s", req.Model, err)
+	}
+
+	req.Context = DoGenerate(ctx, t, client, req, rainbowExpected, 30*time.Second, 20*time.Second)
+
+	for i := 0; i < len(rainbowFollowups); i++ {
+		req.Prompt = rainbowFollowups[i]
+		if time.Now().Sub(started) > softTimeout {
+			slog.Info("exceeded soft timeout, winding down test")
+			return
+		}
+		req.Context = DoGenerate(ctx, t, client, req, rainbowExpected, 30*time.Second, 20*time.Second)
+	}
+}
+
 // Send multiple chat requests with prior context and ensure the response is coherant and expected
-func TestChatWithHistory(t *testing.T) {
+func TestParallelChatWithHistory(t *testing.T) {
 	modelOverride := ollamaEngineChatModels[0] // Most recent ollama engine model
 	req, resp := ChatRequests()
 	numParallel := 2
@@ -164,3 +204,55 @@ func TestChatWithHistory(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+// Send generate requests with prior context and ensure the response is coherant and expected
+func TestChatWithHistory(t *testing.T) {
+	req := api.ChatRequest{
+		Model:     smol,
+		Stream:    &stream,
+		KeepAlive: &api.Duration{Duration: 10 * time.Second},
+		Options: map[string]any{
+			"num_ctx": 16384,
+		},
+		Messages: []api.Message{
+			{
+				Role:    "user",
+				Content: rainbowPrompt,
+			},
+		},
+	}
+
+	softTimeout, hardTimeout := getTimeouts(t)
+	ctx, cancel := context.WithTimeout(context.Background(), hardTimeout)
+	defer cancel()
+	client, _, cleanup := InitServerConnection(ctx, t)
+	defer cleanup()
+
+	// Get the server running (if applicable) warm the model up with a single initial request
+	slog.Info("loading", "model", req.Model)
+	err := client.Generate(ctx,
+		&api.GenerateRequest{Model: req.Model, KeepAlive: &api.Duration{Duration: 10 * time.Second}, Options: req.Options},
+		func(response api.GenerateResponse) error { return nil },
+	)
+	if err != nil {
+		t.Fatalf("failed to load model %s: %s", req.Model, err)
+	}
+
+	assistant := DoChat(ctx, t, client, req, rainbowExpected, 30*time.Second, 20*time.Second)
+
+	for i := 0; i < len(rainbowFollowups); i++ {
+		if time.Now().Sub(started) > softTimeout {
+			slog.Info("exceeded soft timeout, winding down test")
+			return
+		}
+		req.Messages = append(req.Messages,
+			*assistant,
+			api.Message{Role: "user", Content: rainbowFollowups[i]},
+		)
+
+		assistant = DoChat(ctx, t, client, req, rainbowExpected, 30*time.Second, 20*time.Second)
+		if assistant == nil {
+			t.Fatalf("didn't get an assistant response for context")
+		}
+	}
+}

integration/library_models_test.go

@@ -4,7 +4,9 @@ package integration
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
+	"os"
 	"testing"
 	"time"
 
@@ -20,6 +22,7 @@ func TestLibraryModelsGenerate(t *testing.T) {
 	defer cancel()
 	client, _, cleanup := InitServerConnection(ctx, t)
 	defer cleanup()
+	targetArch := os.Getenv("OLLAMA_TEST_ARCHITECTURE")
 
 	chatModels := libraryChatModels
 	for _, model := range chatModels {
@@ -30,16 +33,26 @@ func TestLibraryModelsGenerate(t *testing.T) {
 			if err := PullIfMissing(ctx, client, model); err != nil {
 				t.Fatalf("pull failed %s", err)
 			}
+			if targetArch != "" {
+				resp, err := client.Show(ctx, &api.ShowRequest{Name: model})
+				if err != nil {
+					t.Fatalf("unable to show model: %s", err)
+				}
+				arch := resp.ModelInfo["general.architecture"].(string)
+				if arch != targetArch {
+					t.Skip(fmt.Sprintf("Skipping %s architecture %s != %s", model, arch, targetArch))
+				}
+			}
 			req := api.GenerateRequest{
 				Model:     model,
-				Prompt:    "why is the sky blue?",
+				Prompt:    blueSkyPrompt,
 				KeepAlive: &api.Duration{Duration: 10 * time.Second},
 				Options: map[string]interface{}{
 					"temperature": 0.1,
 					"seed":        123,
 				},
 			}
-			anyResp := []string{"rayleigh", "scatter", "atmosphere", "nitrogen", "oxygen", "wavelength"}
+			anyResp := blueSkyExpected
 			// Special cases
 			if model == "duckdb-nsql" {
 				anyResp = []string{"select", "from"}

integration/model_arch_test.go

@@ -68,14 +68,13 @@ func TestModelsGenerate(t *testing.T) {
 			// TODO - fiddle with context size
 			req := api.GenerateRequest{
 				Model:  model,
-				Prompt: "why is the sky blue?",
+				Prompt: blueSkyPrompt,
 				Options: map[string]interface{}{
 					"temperature": 0,
 					"seed":        123,
 				},
 			}
-			anyResp := []string{"rayleigh", "scattering", "atmosphere", "nitrogen", "oxygen"}
-			DoGenerate(ctx, t, client, req, anyResp, 120*time.Second, 30*time.Second)
+			DoGenerate(ctx, t, client, req, blueSkyExpected, 120*time.Second, 30*time.Second)
 		})
 	}
 }

integration/model_perf_test.go

@@ -40,6 +40,18 @@ var (
 // cat int.log | grep MODEL_PERF_HEADER | head -1| cut -f2- -d: > perf.csv
 // cat int.log | grep MODEL_PERF_DATA | cut -f2- -d: >> perf.csv
 func TestModelsPerf(t *testing.T) {
+	if s := os.Getenv("OLLAMA_NEW_ENGINE"); s != "" {
+		doModelPerfTest(t, ollamaEngineChatModels)
+	} else {
+		doModelPerfTest(t, append(ollamaEngineChatModels, llamaRunnerChatModels...))
+	}
+}
+
+func TestLibraryModelsPerf(t *testing.T) {
+	doModelPerfTest(t, libraryChatModels)
+}
+
+func doModelPerfTest(t *testing.T, chatModels []string) {
 	softTimeout, hardTimeout := getTimeouts(t)
 	slog.Info("Setting timeouts", "soft", softTimeout, "hard", hardTimeout)
 	ctx, cancel := context.WithTimeout(context.Background(), hardTimeout)
@@ -65,14 +77,12 @@ func TestModelsPerf(t *testing.T) {
 	}
 	longPrompt := "summarize the following: " + string(data)
 
-	var chatModels []string
-	if s := os.Getenv("OLLAMA_NEW_ENGINE"); s != "" {
-		chatModels = ollamaEngineChatModels
-	} else {
-		chatModels = append(ollamaEngineChatModels, llamaRunnerChatModels...)
-	}
+	targetArch := os.Getenv("OLLAMA_TEST_ARCHITECTURE")
 
 	for _, model := range chatModels {
+		if !strings.Contains(model, ":") {
+			model = model + ":latest"
+		}
 		t.Run(model, func(t *testing.T) {
 			if time.Now().Sub(started) > softTimeout {
 				t.Skip("skipping remaining tests to avoid excessive runtime")
@@ -88,6 +98,9 @@ func TestModelsPerf(t *testing.T) {
 			}
 			arch := resp.ModelInfo["general.architecture"].(string)
 			maxContext = int(resp.ModelInfo[fmt.Sprintf("%s.context_length", arch)].(float64))
+			if targetArch != "" && arch != targetArch {
+				t.Skip(fmt.Sprintf("Skipping %s architecture %s != %s", model, arch, targetArch))
+			}
 
 			if maxVram > 0 {
 				resp, err := client.List(ctx)
@@ -151,8 +164,8 @@ func TestModelsPerf(t *testing.T) {
 					prompt  string
 					anyResp []string
 				}{
-					{"why is the sky blue?", []string{"rayleigh", "scattering", "atmosphere", "nitrogen", "oxygen"}},
-					{maxPrompt, []string{"shakespeare", "oppression", "sorrows", "gutenberg", "child", "license", "sonnet", "melancholy"}},
+					{blueSkyPrompt, blueSkyExpected},
+					{maxPrompt, []string{"shakespeare", "oppression", "sorrows", "gutenberg", "child", "license", "sonnet", "melancholy", "love", "sorrow", "beauty"}},
 				}
 				var gpuPercent int
 				for _, tc := range testCases {
@@ -241,11 +254,12 @@ func TestModelsPerf(t *testing.T) {
 							}
 						}
 					}
+					// Round the logged prompt count for comparisons across versions/configurations which can vary slightly
 					fmt.Fprintf(os.Stderr, "MODEL_PERF_HEADER:%s,%s,%s,%s,%s,%s,%s\n",
 						"MODEL",
 						"CONTEXT",
 						"GPU PERCENT",
-						"PROMPT COUNT",
+						"APPROX PROMPT COUNT",
 						"LOAD TIME",
 						"PROMPT EVAL TPS",
 						"EVAL TPS",
@@ -254,7 +268,7 @@ func TestModelsPerf(t *testing.T) {
 						model,
 						numCtx,
 						gpuPercent,
-						resp.PromptEvalCount,
+						(resp.PromptEvalCount/10)*10,
 						float64(resp.LoadDuration)/1000000000.0,
 						float64(resp.PromptEvalCount)/(float64(resp.PromptEvalDuration)/1000000000.0),
 						float64(resp.EvalCount)/(float64(resp.EvalDuration)/1000000000.0),

integration/quantization_test.go

@@ -76,7 +76,7 @@ func TestQuantization(t *testing.T) {
 				stream := true
 				genReq := api.GenerateRequest{
 					Model:     newName,
-					Prompt:    "why is the sky blue?",
+					Prompt:    blueSkyPrompt,
 					KeepAlive: &api.Duration{Duration: 3 * time.Second},
 					Options: map[string]any{
 						"seed":        42,
@@ -88,14 +88,13 @@ func TestQuantization(t *testing.T) {
 
 				// Some smaller quantizations can cause models to have poor quality
 				// or get stuck in repetition loops, so we stop as soon as we have any matches
-				anyResp := []string{"rayleigh", "scattering", "day", "sun", "moon", "color", "nitrogen", "oxygen"}
 				reqCtx, reqCancel := context.WithCancel(ctx)
 				atLeastOne := false
 				var buf bytes.Buffer
 				genfn := func(response api.GenerateResponse) error {
 					buf.Write([]byte(response.Response))
 					fullResp := strings.ToLower(buf.String())
-					for _, resp := range anyResp {
+					for _, resp := range blueSkyExpected {
 						if strings.Contains(fullResp, resp) {
 							atLeastOne = true
 							t.Log(fullResp)

integration/utils_test.go

@@ -256,13 +256,29 @@ var (
 		"snowflake-arctic-embed",
 		"snowflake-arctic-embed2",
 	}
+
+	blueSkyPrompt   = "why is the sky blue? Be brief but factual in your reply"
+	blueSkyExpected = []string{"rayleigh", "scatter", "atmosphere", "nitrogen", "oxygen", "wavelength", "interact"}
+
+	rainbowPrompt    = "how do rainbows form? Be brief but factual in your reply"
+	rainbowFollowups = []string{
+		"Explain the physics involved in them.  Be breif in your reply",
+		"Explain the chemistry involved in them.  Be breif in your reply",
+		"Explain the quantum mechanics involved in them. Be breif in your reply",
+		"What are common myths related to them? Be brief in your reply",
+		"What are common fairytales related to them? Be brief in your reply",
+		"Can they form if there is no rain?  Be breif in your reply",
+		"Can they form if there are no clouds?  Be breif in your reply",
+		"Do they happen on other planets? Be brief in your reply",
+	}
+	rainbowExpected = []string{"water", "droplet", "mist", "glow", "refracted", "reflect", "color", "spectrum", "frequency", "end", "gold", "fortune", "blessing", "prosperity"}
 )
 
 func init() {
 	lifecycle.InitLogging()
-	custom := os.Getenv("OLLAMA_TEST_SMOL_MODEL")
+	custom := os.Getenv("OLLAMA_TEST_DEFAULT_MODEL")
 	if custom != "" {
-		slog.Info("setting smol test model to " + custom)
+		slog.Info("setting default test model to " + custom)
 		smol = custom
 	}
 }
@@ -577,11 +593,11 @@ func GenerateRequests() ([]api.GenerateRequest, [][]string) {
 			},
 		},
 		[][]string{
-			{"sunlight", "scattering", "interact", "color", "surface", "depth", "red", "orange", "yellow", "absorbs", "wavelength"},
-			{"soil", "organic", "earth", "black", "tan", "chemical", "processes", "pigments", "particles", "iron oxide", "rust", "air", "water", "mixture", "mixing"},
-			{"water", "droplet", "refracted", "reflect", "color", "spectrum"},
+			{"sunlight", "scatter", "interact", "color", "surface", "depth", "red", "orange", "yellow", "absorb", "wavelength", "water", "molecule"},
+			{"soil", "organic", "earth", "black", "tan", "chemical", "processes", "pigment", "particle", "iron oxide", "rust", "air", "water", "wet", "mixture", "mixing", "mineral", "element", "decomposed", "matter", "wavelength"},
+			{"water", "droplet", "refract", "reflect", "color", "spectrum", "raindrop"},
 			{"fourth", "july", "declaration", "independence"},
-			{"nitrogen", "oxygen", "carbon", "dioxide", "water", "vapor"},
+			{"nitrogen", "oxygen", "carbon", "dioxide", "water", "vapor", "fluid", "particles", "gas"},
 		}
 }
 

llama/patches/0014-graph-memory-reporting-on-failure.patch

@@ -4,48 +4,38 @@ Date: Fri, 18 Apr 2025 15:58:19 -0700
 Subject: [PATCH] graph memory reporting on failure
 
 ---
- ggml/include/ggml-alloc.h   |  6 ++++++
- ggml/include/ggml-backend.h |  6 ++++++
- ggml/src/ggml-alloc.c       | 38 +++++++++++++++++++++++++++++++++----
- ggml/src/ggml-backend.cpp   | 10 ++++++++++
- 4 files changed, 56 insertions(+), 4 deletions(-)
+ ggml/include/ggml-alloc.h   |  1 +
+ ggml/include/ggml-backend.h |  1 +
+ ggml/src/ggml-alloc.c       | 36 ++++++++++++++++++++++++++++++++----
+ ggml/src/ggml-backend.cpp   |  7 +++++++
+ 4 files changed, 41 insertions(+), 4 deletions(-)
 
 diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h
-index 2cb150fd..781b1e10 100644
+index 2cb150fd2..7ab3f0192 100644
 --- a/ggml/include/ggml-alloc.h
 +++ b/ggml/include/ggml-alloc.h
-@@ -66,6 +66,12 @@ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph
+@@ -65,6 +65,7 @@ GGML_API bool ggml_gallocr_reserve_n(
+ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
  
  GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
++GGML_API size_t ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id);
  
-+struct ggml_allocr_buffer_status {
-+    size_t size;
-+    bool allocated;
-+};
-+GGML_API struct ggml_allocr_buffer_status ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id);
-+
  // Utils
  // Create a buffer and allocate all the tensors in a ggml_context
- GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft);
 diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h
-index a2977ea2..8a91b381 100644
+index a2977ea2e..e8cf30841 100644
 --- a/ggml/include/ggml-backend.h
 +++ b/ggml/include/ggml-backend.h
-@@ -304,6 +304,12 @@ extern "C" {
+@@ -303,6 +303,7 @@ extern "C" {
+     GGML_API int                  ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched);
  
      GGML_API size_t               ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
++    GGML_API size_t               ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
  
-+    struct ggml_backend_buffer_status {
-+        size_t size;
-+        bool allocated;
-+    };
-+    GGML_API struct ggml_backend_buffer_status ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
-+
      GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
      GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
- 
 diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c
-index 8b6e6028..41c8c4a2 100644
+index 8b6e60283..b58bd671d 100644
 --- a/ggml/src/ggml-alloc.c
 +++ b/ggml/src/ggml-alloc.c
 @@ -350,6 +350,7 @@ struct node_alloc {
@@ -108,11 +98,11 @@ index 8b6e6028..41c8c4a2 100644
  }
  
  bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph *graph) {
-@@ -920,6 +932,24 @@ size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
+@@ -920,6 +932,22 @@ size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
      return ggml_backend_buffer_get_size(galloc->buffers[buffer_id]);
  }
  
-+struct ggml_allocr_buffer_status ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
++size_t ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
 +    GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers);
 +
 +    for (int i = 0; i < buffer_id; i++) {
@@ -121,34 +111,29 @@ index 8b6e6028..41c8c4a2 100644
 +            // (See above.) However, we need a different check because multiple buffers might be NULL in our
 +            // case and we still want to know the attempted size.
 +
-+            struct ggml_allocr_buffer_status status = {0, true};
-+            return status;
++            return 0;
 +        }
 +    }
 +
-+    struct ggml_allocr_buffer_status status = {galloc->buffer_sizes[buffer_id], galloc->buffers[buffer_id] != NULL};
-+    return status;
++    return galloc->buffer_sizes[buffer_id];
 +}
 +
  // utils
  
  static void free_buffers(ggml_backend_buffer_t ** buffers, const size_t * n_buffers) {
 diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
-index 97f47abd..eded0291 100644
+index 97f47abd2..d02a40e60 100644
 --- a/ggml/src/ggml-backend.cpp
 +++ b/ggml/src/ggml-backend.cpp
-@@ -1631,6 +1631,16 @@ size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backe
+@@ -1631,6 +1631,13 @@ size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backe
      return ggml_gallocr_get_buffer_size(sched->galloc, backend_index);
  }
  
-+struct ggml_backend_buffer_status ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) {
++size_t ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) {
 +    int backend_index = ggml_backend_sched_backend_id(sched, backend);
 +    GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
 +
-+    struct ggml_allocr_buffer_status allocr_status = ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
-+    struct ggml_backend_buffer_status status = {allocr_status.size, allocr_status.allocated};
-+
-+    return status;
++    return ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
 +}
 +
  void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) {

llama/patches/0022-ggml-No-alloc-mode.patch

@@ -3,35 +3,45 @@ From: Jesse Gross <jesse@ollama.com>
 Date: Wed, 23 Jul 2025 11:58:49 -0700
 Subject: [PATCH] ggml: No-alloc mode
 
-Callers can set a backend buffer type to be no-alloc, meaning that
+Callers can set a scheduler to be no-alloc, meaning that
 it does not allocate memory for tensors or operations. This can
 be used for calculating memory requirements. Tensors and graphs
 must be recreated with no-alloc set to false before loading data.
-
-Defaults to false for newly created backend buffer types.
 ---
- ggml/include/ggml-backend.h  |  1 +
- ggml/src/ggml-backend-impl.h |  2 ++
- ggml/src/ggml-backend.cpp    | 19 ++++++++++++++++++-
- 3 files changed, 21 insertions(+), 1 deletion(-)
+ ggml/include/ggml-backend.h     |   1 +
+ ggml/src/ggml-backend-impl.h    |  16 +++
+ ggml/src/ggml-backend.cpp       |  72 ++++++++++-
+ ggml/src/ggml-cuda/common.cuh   |  48 ++++++-
+ ggml/src/ggml-cuda/ggml-cuda.cu | 217 ++++++++++++++++++++++++++------
+ 5 files changed, 310 insertions(+), 44 deletions(-)
 
 diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h
-index 9424394e..b602a7c7 100644
+index 2773cc310..ae94887dd 100644
 --- a/ggml/include/ggml-backend.h
 +++ b/ggml/include/ggml-backend.h
-@@ -35,6 +35,7 @@ extern "C" {
-     //
+@@ -291,6 +291,7 @@ extern "C" {
  
-     GGML_API const char *          ggml_backend_buft_name          (ggml_backend_buffer_type_t buft);
-+    GGML_API void                  ggml_backend_buft_set_alloc     (ggml_backend_buffer_type_t buft, bool alloc);
-     GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer  (ggml_backend_buffer_type_t buft, size_t size);
-     GGML_API size_t                ggml_backend_buft_get_alignment (ggml_backend_buffer_type_t buft);
-     GGML_API size_t                ggml_backend_buft_get_max_size  (ggml_backend_buffer_type_t buft);
+     // Initialize a backend scheduler, backends with low index are given priority over backends with high index
+     GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload);
++    GGML_API ggml_backend_sched_t ggml_backend_sched_new_ext(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload, bool alloc_buffers);
+     GGML_API void                 ggml_backend_sched_free(ggml_backend_sched_t sched);
+ 
+     // Initialize backend buffers from a measure graph
 diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h
-index c36c12d6..81749a5a 100644
+index c36c12d65..369e9e25a 100644
 --- a/ggml/src/ggml-backend-impl.h
 +++ b/ggml/src/ggml-backend-impl.h
-@@ -32,6 +32,7 @@ extern "C" {
+@@ -26,12 +26,17 @@ extern "C" {
+         size_t                (*get_alloc_size)(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor);
+         // (optional) check if tensor data is in host memory and uses standard ggml tensor layout (defaults to false)
+         bool                  (*is_host)       (ggml_backend_buffer_type_t buft);
++
++        // (optional) returns a dummy buffer that is equivalent to one created by alloc_buffer but without actually being backed
++        // by memory
++        ggml_backend_buffer_t (*noalloc_buffer)(ggml_backend_buffer_type_t buft, size_t size);
+     };
+ 
+     struct ggml_backend_buffer_type {
          struct ggml_backend_buffer_type_i  iface;
          ggml_backend_dev_t device;
          void * context;
@@ -39,7 +49,7 @@ index c36c12d6..81749a5a 100644
      };
  
      //
-@@ -63,6 +64,7 @@ extern "C" {
+@@ -63,6 +68,7 @@ extern "C" {
          void * context;
          size_t size;
          enum ggml_backend_buffer_usage usage;
@@ -47,26 +57,40 @@ index c36c12d6..81749a5a 100644
      };
  
      GGML_API ggml_backend_buffer_t ggml_backend_buffer_init(
+@@ -114,6 +120,16 @@ extern "C" {
+         void (*event_record)(ggml_backend_t backend, ggml_backend_event_t event);
+         // wait for an event on on a different stream
+         void (*event_wait)  (ggml_backend_t backend, ggml_backend_event_t event);
++
++        // (optional) reserves intermediate buffers needed for the compution
++        // if alloc is true, memory is actually allocated, otherwise the required amount is just returned by buffer_size
++        enum ggml_status          (*graph_reserve)     (ggml_backend_t backend, struct ggml_cgraph * cgraph, bool alloc);
++
++        // (optional) returns the memory needed after calling graph_reserve
++        size_t                    (*buffer_size)       (ggml_backend_t backend);
++
++        // (optional) frees memory from intermediate buffers that was allocated either by graph_compute or graph_reserve
++        void                      (*reset)             (ggml_backend_t backend);
+     };
+ 
+     struct ggml_backend {
 diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
-index eded0291..05a842ed 100644
+index d02a40e60..6b4dee4c7 100644
 --- a/ggml/src/ggml-backend.cpp
 +++ b/ggml/src/ggml-backend.cpp
-@@ -35,12 +35,22 @@ const char * ggml_backend_buft_name(ggml_backend_buffer_type_t buft) {
-     return buft->iface.get_name(buft);
- }
- 
-+void ggml_backend_buft_set_alloc(ggml_backend_buffer_type_t buft, bool alloc) {
-+    buft->no_alloc = !alloc;
-+}
-+
- ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
-     if (size == 0) {
-         // return a dummy buffer for zero-sized allocations
+@@ -41,6 +41,19 @@ ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t
          return ggml_backend_buffer_init(buft, {}, NULL, 0);
      }
  
 +    if (buft->no_alloc) {
-+        ggml_backend_buffer_t buf = ggml_backend_buffer_init(buft, {}, NULL, size);
++        ggml_backend_buffer_t buf;
++
++        if (buft->iface.noalloc_buffer != NULL) {
++            buf = buft->iface.noalloc_buffer(buft, size);
++        } else {
++            buf = ggml_backend_buffer_init(buft, {}, NULL, size);
++        }
++
 +        buf->no_alloc = true;
 +        return buf;
 +    }
@@ -74,7 +98,7 @@ index eded0291..05a842ed 100644
      return buft->iface.alloc_buffer(buft, size);
  }
  
-@@ -89,7 +99,8 @@ ggml_backend_buffer_t ggml_backend_buffer_init(
+@@ -89,7 +102,8 @@ ggml_backend_buffer_t ggml_backend_buffer_init(
          /* .buft      = */ buft,
          /* .context   = */ context,
          /* .size      = */ size,
@@ -84,7 +108,7 @@ index eded0291..05a842ed 100644
      };
  
      return buffer;
-@@ -119,6 +130,12 @@ void * ggml_backend_buffer_get_base(ggml_backend_buffer_t buffer) {
+@@ -119,6 +133,12 @@ void * ggml_backend_buffer_get_base(ggml_backend_buffer_t buffer) {
          return NULL;
      }
  
@@ -97,3 +121,532 @@ index eded0291..05a842ed 100644
      void * base = buffer->iface.get_base(buffer);
  
      GGML_ASSERT(base != NULL && "backend buffer base cannot be NULL");
+@@ -663,6 +683,12 @@ struct ggml_backend_sched {
+     bool op_offload;
+ 
+     int debug;
++
++    // allocate buffers on attached ggml_backend_buffer_type_t's and during reservation
++    // if false, dummy buffers are used for faster memory sizing calculations
++    // the scheduler needs to be recreated with allocated buffers before it can be used
++    // for computation
++    bool alloc_buffers;
+ };
+ 
+ #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
+@@ -1449,6 +1475,17 @@ ggml_backend_sched_t ggml_backend_sched_new(
+         size_t graph_size,
+         bool parallel,
+         bool op_offload) {
++            return ggml_backend_sched_new_ext(backends, bufts, n_backends, graph_size, parallel, op_offload, true);
++        }
++
++ggml_backend_sched_t ggml_backend_sched_new_ext(
++        ggml_backend_t * backends,
++        ggml_backend_buffer_type_t * bufts,
++        int n_backends,
++        size_t graph_size,
++        bool parallel,
++        bool op_offload,
++        bool alloc_buffers) {
+     GGML_ASSERT(n_backends > 0);
+     GGML_ASSERT(n_backends <= GGML_SCHED_MAX_BACKENDS);
+     GGML_ASSERT(ggml_backend_dev_type(ggml_backend_get_device(backends[n_backends - 1])) == GGML_BACKEND_DEVICE_TYPE_CPU);
+@@ -1490,10 +1527,13 @@ ggml_backend_sched_t ggml_backend_sched_new(
+                 sched->events[b][c] = ggml_backend_event_new(backends[b]->device);
+             }
+         }
++
++        sched->bufts[b]->no_alloc = !alloc_buffers;
+     }
+ 
+     sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
+     sched->op_offload = op_offload;
++    sched->alloc_buffers = alloc_buffers;
+ 
+     ggml_backend_sched_reset(sched);
+ 
+@@ -1508,6 +1548,10 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
+         for (int c = 0; c < sched->n_copies; c++) {
+             ggml_backend_event_free(sched->events[b][c]);
+         }
++
++        if (sched->backends[b]->iface.reset != NULL) {
++            sched->backends[b]->iface.reset(sched->backends[b]);
++        }
+     }
+     ggml_gallocr_free(sched->galloc);
+     ggml_free(sched->ctx);
+@@ -1547,6 +1591,24 @@ bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph *
+         return false;
+     }
+ 
++    if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
++        return false;
++    }
++
++    struct ggml_backend_sched_split * splits = sched->splits;
++    for (int i = 0; i < sched->n_splits; i++) {
++        struct ggml_backend_sched_split * split = &splits[i];
++        int split_backend_id = split->backend_id;
++        ggml_backend_t split_backend = sched->backends[split_backend_id];
++
++        if (split_backend->iface.graph_reserve != NULL) {
++            enum ggml_status ec = split_backend->iface.graph_reserve(split_backend, &split->graph, sched->alloc_buffers);
++            if (ec != GGML_STATUS_SUCCESS) {
++                return false;
++            }
++        }
++    }
++
+     ggml_backend_sched_reset(sched);
+ 
+     return true;
+@@ -1635,7 +1697,13 @@ size_t ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched,
+     int backend_index = ggml_backend_sched_backend_id(sched, backend);
+     GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
+ 
+-    return ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
++    size_t size = ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
++
++    if (backend->iface.buffer_size != NULL) {
++        size += backend->iface.buffer_size(backend);
++    }
++
++    return size;
+ }
+ 
+ void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) {
+diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
+index 2e5d48797..b915ee1b8 100644
+--- a/ggml/src/ggml-cuda/common.cuh
++++ b/ggml/src/ggml-cuda/common.cuh
+@@ -35,6 +35,31 @@
+ #include "vendors/cuda.h"
+ #endif // defined(GGML_USE_HIP)
+ 
++extern bool reserving_graph;
++
++// If we are reserving the graph, pointers might be invalid and will fail if cudaMemcpyAsync tries to validate them.
++// However, since we don't actually expect a result, we don't need to actually do the memcpy.
++static cudaError_t cudaMemcpyAsyncReserve ( void* dst, const void* src, size_t count, cudaMemcpyKind kind, cudaStream_t stream = 0 ) {
++    if (!reserving_graph) {
++        return cudaMemcpyAsync(dst, src, count, kind, stream);
++    } else {
++        return cudaSuccess;
++    }
++}
++
++static cudaError_t cudaMemcpy2DAsyncReserve ( void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, cudaMemcpyKind kind, cudaStream_t stream = 0 ) {
++    if (!reserving_graph) {
++        return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, kind, stream);
++    } else {
++        return cudaSuccess;
++    }
++}
++
++#undef cudaMemcpyAsync
++#define cudaMemcpyAsync cudaMemcpyAsyncReserve
++#undef cudaMemcpy2DAsync
++#define cudaMemcpy2DAsync cudaMemcpy2DAsyncReserve
++
+ #define STRINGIZE_IMPL(...) #__VA_ARGS__
+ #define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__)
+ 
+@@ -771,6 +796,9 @@ struct ggml_cuda_pool {
+ 
+     virtual void * alloc(size_t size, size_t * actual_size) = 0;
+     virtual void free(void * ptr, size_t size) = 0;
++
++    virtual bool alloc_memory() = 0;
++    virtual size_t alloc_size() = 0;
+ };
+ 
+ template<typename T>
+@@ -914,11 +942,11 @@ struct ggml_backend_cuda_context {
+     // pool
+     std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES];
+ 
+-    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device);
++    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, bool alloc);
+ 
+     ggml_cuda_pool & pool(int device) {
+         if (pools[device] == nullptr) {
+-            pools[device] = new_pool_for_device(device);
++            pools[device] = new_pool_for_device(device, true);
+         }
+         return *pools[device];
+     }
+@@ -926,4 +954,20 @@ struct ggml_backend_cuda_context {
+     ggml_cuda_pool & pool() {
+         return pool(device);
+     }
++
++    void pool_set_alloc(bool alloc) {
++        GGML_ASSERT(pools[device] == nullptr || pools[device]->alloc_memory() == alloc);
++
++        if (pools[device] == nullptr) {
++            pools[device] = new_pool_for_device(device, alloc);
++        }
++    }
++
++    size_t pool_get_alloc_size() {
++        if (pools[device] == nullptr) {
++            return 0;
++        }
++
++        return pools[device]->alloc_size();
++    }
+ };
+diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
+index c7f9dc3a5..d5abe09e0 100644
+--- a/ggml/src/ggml-cuda/ggml-cuda.cu
++++ b/ggml/src/ggml-cuda/ggml-cuda.cu
+@@ -350,6 +350,8 @@ const ggml_cuda_device_info & ggml_cuda_info() {
+ 
+ // #define DEBUG_CUDA_MALLOC
+ 
++#define CUDA_ALIGNMENT 128
++
+ // buffer pool for cuda (legacy)
+ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
+     static const int MAX_BUFFERS = 256;
+@@ -362,9 +364,12 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
+ 
+     ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {};
+     size_t pool_size = 0;
++    bool allocate = true;
++    size_t last_alloc = 0;
+ 
+-    explicit ggml_cuda_pool_leg(int device) :
+-        device(device) {
++    explicit ggml_cuda_pool_leg(int device, bool alloc) :
++        device(device),
++        allocate(alloc) {
+     }
+ 
+     ~ggml_cuda_pool_leg() {
+@@ -372,7 +377,9 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
+         for (int i = 0; i < MAX_BUFFERS; ++i) {
+             ggml_cuda_buffer & b = buffer_pool[i];
+             if (b.ptr != nullptr) {
+-                CUDA_CHECK(cudaFree(b.ptr));
++                if (allocate) {
++                    CUDA_CHECK(cudaFree(b.ptr));
++                }
+                 pool_size -= b.size;
+             }
+         }
+@@ -420,8 +427,15 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
+         void * ptr;
+         size_t look_ahead_size = (size_t) (1.05 * size);
+         look_ahead_size = 256 * ((look_ahead_size + 255)/256);
+-        ggml_cuda_set_device(device);
+-        CUDA_CHECK(ggml_cuda_device_malloc(&ptr, look_ahead_size, device));
++        if (allocate) {
++            ggml_cuda_set_device(device);
++            if (ggml_cuda_device_malloc(&ptr, look_ahead_size, device) != cudaSuccess) {
++                    last_alloc = look_ahead_size;
++                    throw std::bad_alloc();
++            }
++        } else {
++            ptr = (void *)CUDA_ALIGNMENT;
++        }
+         *actual_size = look_ahead_size;
+         pool_size += look_ahead_size;
+ #ifdef DEBUG_CUDA_MALLOC
+@@ -441,10 +455,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
+             }
+         }
+         GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n");
+-        ggml_cuda_set_device(device);
+-        CUDA_CHECK(cudaFree(ptr));
++        if (allocate) {
++            ggml_cuda_set_device(device);
++            CUDA_CHECK(cudaFree(ptr));
++        }
+         pool_size -= size;
+     }
++
++    bool alloc_memory() override {
++        return allocate;
++    }
++
++    size_t alloc_size() override {
++        return pool_size + last_alloc;
++    }
+ };
+ 
+ // pool with virtual memory
+@@ -456,18 +480,24 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
+     CUdeviceptr pool_addr = 0;
+     size_t pool_used = 0;
+     size_t pool_size = 0;
++    bool allocate = true;
++    size_t last_alloc = 0;
+     size_t granularity;
+ #if defined(GGML_USE_HIP)
+     std::vector<std::pair<CUdeviceptr, size_t>> mappings;
+ #endif
+ 
+-    explicit ggml_cuda_pool_vmm(int device) :
++    explicit ggml_cuda_pool_vmm(int device, bool alloc) :
+         device(device),
+-        granularity(ggml_cuda_info().devices[device].vmm_granularity) {
++        granularity(ggml_cuda_info().devices[device].vmm_granularity),
++        allocate(alloc) {
++        if (!allocate) {
++            pool_addr = (CUdeviceptr)CUDA_ALIGNMENT;
++        }
+     }
+ 
+     ~ggml_cuda_pool_vmm() {
+-        if (pool_addr != 0) {
++        if (pool_addr != 0 && allocate) {
+ #if defined(GGML_USE_HIP)
+             // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285
+             for (std::pair<CUdeviceptr, size_t> & mapping : mappings) {
+@@ -494,35 +524,49 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
+ 
+             GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE);
+ 
+-            // allocate more physical memory
+-            CUmemAllocationProp prop = {};
+-            prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+-            prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+-            prop.location.id = device;
+-            CUmemGenericAllocationHandle handle;
+-            CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0));
+-
+-            // reserve virtual address space (if not already reserved)
+-            if (pool_addr == 0) {
+-                CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0));
+-            }
++            if (allocate) {
++                // allocate more physical memory
++                CUmemAllocationProp prop = {};
++                prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
++                prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
++                prop.location.id = device;
++                CUmemGenericAllocationHandle handle;
++                if (cuMemCreate(&handle, reserve_size, &prop, 0) != CUDA_SUCCESS) {
++                    last_alloc = reserve_size;
++                    throw std::bad_alloc();
++                }
+ 
+-            // map at the end of the pool
+-            CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size);
+-            CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0));
+-#if defined(GGML_USE_HIP)
+-            mappings.push_back({start_ptr, reserve_size});
+-#endif
++                // reserve virtual address space (if not already reserved)
++                if (pool_addr == 0) {
++                    CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0));
++                }
+ 
+-            // the memory allocation handle is no longer needed after mapping
+-            CU_CHECK(cuMemRelease(handle));
++                // map at the end of the pool
++                CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size);
++                if (cuMemMap(start_ptr, reserve_size, 0, handle, 0) != CUDA_SUCCESS) {
++                    last_alloc = reserve_size;
++                    CU_CHECK(cuMemRelease(handle));
++                    throw std::bad_alloc();
++                }
++
++                // the memory allocation handle is no longer needed after mapping
++                CU_CHECK(cuMemRelease(handle));
++
++                // set access
++                CUmemAccessDesc access = {};
++                access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
++                access.location.id = device;
++                access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
++                if (cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1) != CUDA_SUCCESS) {
++                    CU_CHECK(cuMemUnmap(start_ptr, reserve_size));
++                    last_alloc = reserve_size;
++                    throw std::bad_alloc();
++                }
+ 
+-            // set access
+-            CUmemAccessDesc access = {};
+-            access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+-            access.location.id = device;
+-            access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+-            CU_CHECK(cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1));
++    #if defined(GGML_USE_HIP)
++                mappings.push_back({start_ptr, reserve_size});
++    #endif
++            }
+ 
+             // add to the pool
+             pool_size += reserve_size;
+@@ -555,16 +599,24 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
+         // all deallocations must be in reverse order of the allocations
+         GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used));
+     }
++
++    bool alloc_memory() override {
++        return allocate;
++    }
++
++    size_t alloc_size() override {
++        return pool_size + last_alloc;
++    }
+ };
+ #endif // defined(GGML_USE_VMM)
+ 
+-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device) {
++std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, bool alloc) {
+ #if defined(GGML_USE_VMM)
+     if (ggml_cuda_info().devices[device].vmm) {
+-        return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device));
++        return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device, alloc));
+     }
+ #endif // defined(GGML_USE_VMM)
+-    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device));
++    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device, alloc));
+ }
+ 
+ // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
+@@ -748,11 +800,20 @@ static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_bac
+ }
+ 
+ static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
+-    return 128;
++    return CUDA_ALIGNMENT;
+ 
+     GGML_UNUSED(buft);
+ }
+ 
++static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_noalloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
++    ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context;
++
++    void * dev_ptr = (void *)ggml_backend_cuda_buffer_type_get_alignment(buft);
++    ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr);
++
++    return ggml_backend_buffer_init(buft, {}, ctx, size);
++}
++
+ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
+     size_t size = ggml_nbytes(tensor);
+     int64_t ne0 = tensor->ne[0];
+@@ -776,6 +837,7 @@ static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface
+     /* .get_max_size     = */ NULL, // defaults to SIZE_MAX
+     /* .get_alloc_size   = */ ggml_backend_cuda_buffer_type_get_alloc_size,
+     /* .is_host          = */ NULL,
++    /* .noalloc_buffer   = */ ggml_backend_cuda_buffer_type_noalloc_buffer,
+ };
+ 
+ ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) {
+@@ -2936,6 +2998,7 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx,
+ 
+ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph,
+     bool & graph_evaluated_or_captured, bool & use_cuda_graph, bool & cuda_graph_update_required) {
++
+     // flag used to determine whether it is an integrated_gpu
+     const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated;
+ 
+@@ -2951,6 +3014,11 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
+                     continue;
+                 }
+ 
++                // When reserving, we are forcing CUDA graphs but this operation is not graph-safe so we need to skip it
++                if (reserving_graph && node->op == GGML_OP_MUL_MAT_ID && node->ne[2] != 1) {
++                    continue;
++                }
++
+                 static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr);
+                 if (!disable_fusion) {
+                     if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
+@@ -3022,6 +3090,7 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
+ 
+ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
+     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
++    cuda_ctx->pool_set_alloc(true);
+ 
+     ggml_cuda_set_device(cuda_ctx->device);
+ 
+@@ -3101,6 +3170,71 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
+     return GGML_STATUS_SUCCESS;
+ }
+ 
++// This is used to skip operations that are not graph safe during the reservation process.
++bool reserving_graph = false;
++
++static enum ggml_status ggml_backend_cuda_graph_reserve(ggml_backend_t backend, ggml_cgraph * cgraph, bool alloc) {
++    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
++    cuda_ctx->pool_set_alloc(alloc);
++
++    #ifdef USE_CUDA_GRAPH
++    if (cuda_ctx->cuda_graph == nullptr) {
++        cuda_ctx->cuda_graph.reset(new ggml_cuda_graph());
++    }
++    #endif
++
++    ggml_cuda_set_device(cuda_ctx->device);
++
++    {
++        std::lock_guard<std::mutex> lock(ggml_cuda_lock);
++        ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed);
++    }
++
++    reserving_graph = true;
++
++    // Create CuBLAS handles early to avoid synchronous allocations during graph capture.
++    cuda_ctx->cublas_handle();
++
++    CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed));
++
++    enum ggml_status result = GGML_STATUS_SUCCESS;
++
++    try {
++        bool use_cuda_graph = false;
++        bool cuda_graph_update_required = false;
++        bool graph_evaluated_or_captured = false;
++
++        evaluate_and_capture_cuda_graph(cuda_ctx, cgraph, graph_evaluated_or_captured, use_cuda_graph, cuda_graph_update_required);
++    } catch (const std::exception &e) {
++        result = GGML_STATUS_FAILED;
++    }
++
++    cudaGraph_t graph;
++    CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph));
++    CUDA_CHECK(cudaGraphDestroy(graph));
++
++    reserving_graph = false;
++
++    {
++        std::lock_guard<std::mutex> lock(ggml_cuda_lock);
++        if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) {
++            ggml_cuda_lock_cv.notify_all();
++        }
++    }
++
++    return result;
++}
++
++static size_t ggml_backend_cuda_buffer_size(ggml_backend_t backend) {
++    ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *)backend->context;
++    return ctx->pool_get_alloc_size();
++}
++
++static void ggml_backend_cuda_reset(ggml_backend_t backend) {
++    ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *)backend->context;
++    ctx->pools[ctx->device] = NULL;
++}
++
+ static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) {
+     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
+ 
+@@ -3140,6 +3274,9 @@ static const ggml_backend_i ggml_backend_cuda_interface = {
+     /* .graph_compute           = */ ggml_backend_cuda_graph_compute,
+     /* .event_record            = */ ggml_backend_cuda_event_record,
+     /* .event_wait              = */ ggml_backend_cuda_event_wait,
++    /* .graph_reserve           = */ ggml_backend_cuda_graph_reserve,
++    /* .buffer_size             = */ ggml_backend_cuda_buffer_size,
++    /* .reset                   = */ ggml_backend_cuda_reset,
+ };
+ 
+ static ggml_guid_t ggml_backend_cuda_guid() {

llama/patches/0026-ggml-Backport-scale-kernel-fixes.patch

@@ -0,0 +1,57 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: Jesse Gross <jesse@ollama.com>
+Date: Tue, 23 Sep 2025 15:41:58 -0700
+Subject: [PATCH] ggml: Backport scale kernel fixes
+
+The GGML scale kernel uses signed 32-bit ints to represent
+the number of elements in the tensor. For large images,
+mistral-small3.2 overflows this, triggering CUDA errors due
+to negative arguments.
+
+Currently, this can happen when the user passes a large image
+to mistral-small3.2. However, with upcoming changes to reserve
+CUDA memory, it happens every time mistral-small is loaded as
+we reserve using a worst case batch.
+
+This patch is part of an upstream GGML commit and should be removed
+after GGML is updated past 0a1b398 "ggml: add ops for WAN video model
+(cuda && cpu) (#15669)".
+
+Fixes #10388
+---
+ ggml/src/ggml-cuda/scale.cu | 19 ++++++++++---------
+ 1 file changed, 10 insertions(+), 9 deletions(-)
+
+diff --git a/ggml/src/ggml-cuda/scale.cu b/ggml/src/ggml-cuda/scale.cu
+index 2ee9e5889..0ddeff6a1 100644
+--- a/ggml/src/ggml-cuda/scale.cu
++++ b/ggml/src/ggml-cuda/scale.cu
+@@ -1,18 +1,19 @@
+ #include "scale.cuh"
+ 
+-static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int k) {
+-    const int i = blockDim.x*blockIdx.x + threadIdx.x;
++#define MAX_GRIDDIM_X 0x7FFFFFFF
+ 
+-    if (i >= k) {
+-        return;
+-    }
++static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int64_t nelements) {
++    int64_t tid = (int64_t)blockIdx.x * (int64_t)blockDim.x + (int64_t)threadIdx.x;
++    int64_t stride = (int64_t)blockDim.x * (int64_t)gridDim.x;
+ 
+-    dst[i] = scale * x[i] + bias;
++    for (int64_t i = tid; i < nelements; i += stride) {
++        dst[i] = scale * x[i] + bias;
++    }
+ }
+ 
+-static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int k, cudaStream_t stream) {
+-    const int num_blocks = (k + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
+-    scale_f32<<<num_blocks, CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, k);
++static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
++    const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
++    scale_f32<<<MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, nelements);
+ }
+ 
+ void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

llm/server.go

@@ -853,19 +853,19 @@ func (s *ollamaServer) createLayout(systemInfo discover.SystemInfo, systemGPUs d
 
 	if memory == nil {
 		memory = &ml.BackendMemory{CPU: ml.DeviceMemory{
-			Weights: make([]ml.Memory, s.totalLayers),
-			Cache:   make([]ml.Memory, s.totalLayers),
+			Weights: make([]uint64, s.totalLayers),
+			Cache:   make([]uint64, s.totalLayers),
 		}}
 	}
 
 	layers := make([]uint64, len(memory.CPU.Weights))
 	for i := range layers {
 		for j := range memory.GPUs {
-			layers[i] += memory.GPUs[j].Weights[i].Size
-			layers[i] += memory.GPUs[j].Cache[i].Size
+			layers[i] += memory.GPUs[j].Weights[i]
+			layers[i] += memory.GPUs[j].Cache[i]
 		}
-		layers[i] += memory.CPU.Weights[i].Size
-		layers[i] += memory.CPU.Cache[i].Size
+		layers[i] += memory.CPU.Weights[i]
+		layers[i] += memory.CPU.Cache[i]
 		logutil.Trace("layer to assign", "layer", i, "size", format.HumanBytes2(layers[i]))
 	}
 
@@ -880,11 +880,11 @@ func (s *ollamaServer) createLayout(systemInfo discover.SystemInfo, systemGPUs d
 			found := false
 			for j := range memory.GPUs {
 				if gl[i].ID == memory.GPUs[j].ID {
-					if memory.GPUs[j].Graph.Size != 0 {
+					if memory.GPUs[j].Graph != 0 {
 						lastUsedGPU = i
 					}
 
-					reserved := uint64(float32(gl[i].FreeMemory)*backoff) + gl[i].MinimumMemory + envconfig.GpuOverhead() + memory.GPUs[j].Graph.Size
+					reserved := uint64(float32(gl[i].FreeMemory)*backoff) + gl[i].MinimumMemory + envconfig.GpuOverhead() + memory.GPUs[j].Graph
 					if gl[i].FreeMemory > reserved {
 						gl[i].FreeMemory -= reserved
 					} else {
@@ -895,7 +895,7 @@ func (s *ollamaServer) createLayout(systemInfo discover.SystemInfo, systemGPUs d
 						"available layer vram", format.HumanBytes2(gl[i].FreeMemory),
 						"backoff", fmt.Sprintf("%.2f", backoff), "minimum", format.HumanBytes2(gl[i].MinimumMemory),
 						"overhead", format.HumanBytes2(envconfig.GpuOverhead()),
-						"graph", format.HumanBytes2(memory.GPUs[j].Graph.Size))
+						"graph", format.HumanBytes2(memory.GPUs[j].Graph))
 
 					found = true
 					break
@@ -914,12 +914,12 @@ func (s *ollamaServer) createLayout(systemInfo discover.SystemInfo, systemGPUs d
 	}
 
 	// These sizes will only increase as we go through additional iterations and get additional information.
-	cpuSize := memory.InputWeights.Size + memory.CPU.Graph.Size
+	cpuSize := memory.InputWeights + memory.CPU.Graph
 	var vramSize uint64
 	for _, gl := range gpuLayers {
 		for _, gpu := range memory.GPUs {
 			if gl.ID == gpu.ID {
-				vramSize += gpu.Graph.Size
+				vramSize += gpu.Graph
 				break
 			}
 		}
@@ -1723,21 +1723,21 @@ func (s *ollamaServer) VRAMSize() uint64 {
 	var mem uint64
 
 	for _, g := range s.mem.GPUs {
-		mem += g.Allocated()
+		mem += g.Size()
 	}
 
 	// Some elements are always on CPU. However, if we have allocated all layers
 	// on the GPU then include the CPU components as well, to represent complete offloading.
 	noCPULayers := true
 	for i := range s.mem.CPU.Weights {
-		if s.mem.CPU.Weights[i].Size != 0 || s.mem.CPU.Cache[i].Size != 0 {
+		if s.mem.CPU.Weights[i] != 0 || s.mem.CPU.Cache[i] != 0 {
 			noCPULayers = false
 			break
 		}
 	}
 	if noCPULayers {
-		mem += s.mem.InputWeights.Size
-		mem += s.mem.CPU.Graph.Size
+		mem += s.mem.InputWeights
+		mem += s.mem.CPU.Graph
 	}
 
 	return mem
@@ -1748,10 +1748,10 @@ func (s *ollamaServer) TotalSize() uint64 {
 		return 0
 	}
 
-	mem := s.mem.InputWeights.Size
-	mem += s.mem.CPU.Allocated()
+	mem := s.mem.InputWeights
+	mem += s.mem.CPU.Size()
 	for _, g := range s.mem.GPUs {
-		mem += g.Allocated()
+		mem += g.Size()
 	}
 
 	return mem
@@ -1764,7 +1764,7 @@ func (s *ollamaServer) VRAMByGPU(gpuID string) uint64 {
 
 	for _, g := range s.mem.GPUs {
 		if g.ID == gpuID {
-			return g.Allocated()
+			return g.Size()
 		}
 	}
 

llm/server_test.go

@@ -155,18 +155,18 @@ func TestLLMServerFitGPU(t *testing.T) {
 			}
 
 			s.mem = &ml.BackendMemory{CPU: ml.DeviceMemory{
-				Weights: make([]ml.Memory, s.totalLayers),
-				Cache:   make([]ml.Memory, s.totalLayers),
+				Weights: make([]uint64, s.totalLayers),
+				Cache:   make([]uint64, s.totalLayers),
 			}, GPUs: make([]ml.DeviceMemory, len(gpus))}
 
 			for i := range tt.layers {
-				s.mem.CPU.Weights[i].Size = uint64(tt.layers[i])
+				s.mem.CPU.Weights[i] = uint64(tt.layers[i])
 			}
 
 			for i := range s.mem.GPUs {
 				s.mem.GPUs[i].ID = fmt.Sprintf("gpu%d", i)
-				s.mem.GPUs[i].Weights = make([]ml.Memory, s.totalLayers)
-				s.mem.GPUs[i].Cache = make([]ml.Memory, s.totalLayers)
+				s.mem.GPUs[i].Weights = make([]uint64, s.totalLayers)
+				s.mem.GPUs[i].Cache = make([]uint64, s.totalLayers)
 			}
 
 			gpuLayers, err := s.createLayout(systemInfo, gpus, s.mem, tt.requireFull, 0)

ml/backend.go

@@ -158,40 +158,6 @@ func (e ErrNoMem) Error() string {
 	return fmt.Sprintf("insufficient memory - required allocations: %+v", e.BackendMemory)
 }
 
-type AllocationStatus int
-
-const (
-	// Unallocated memory - have not yet attempted to allocate
-	Unallocated AllocationStatus = iota
-
-	// Failed memory - tried to allocate the memory and did not succeed
-	Failed
-
-	// Allocated memory = tried and succeeded to allocate memory
-	Allocated
-)
-
-// Memory is the size of an allocation and whether it was successful.
-type Memory struct {
-	Size   uint64
-	Status AllocationStatus
-}
-
-func (m Memory) String() string {
-	s := fmt.Sprint(m.Size)
-
-	switch m.Status {
-	case Unallocated:
-		s += "U"
-	case Failed:
-		s += "F"
-	case Allocated:
-		s += "A"
-	}
-
-	return s
-}
-
 // DeviceMemory provides a breakdown of the memory needed
 // per device, such as a CPU or GPU.
 type DeviceMemory struct {
@@ -204,39 +170,32 @@ type DeviceMemory struct {
 	ID string
 
 	// Weights is the per-layer memory needed for the model weights.
-	Weights []Memory
+	Weights []uint64
 
 	// Cache is the per-layer memory needed for the KV cache.
-	Cache []Memory
+	Cache []uint64
 
 	// Graph is the size of the compute graph. It is not per-layer.
-	Graph Memory
+	Graph uint64
 }
 
-// Allocated returns the total size of the memory that has been successfully
-// allocated on this device
-func (m DeviceMemory) Allocated() uint64 {
-	var mem uint64
+func sumMemory(mem []uint64) uint64 {
+	var sum uint64
 
-	for _, w := range m.Weights {
-		if w.Status == Allocated {
-			mem += w.Size
-		}
-	}
-	for _, c := range m.Cache {
-		if c.Status == Allocated {
-			mem += c.Size
-		}
-	}
-	if m.Graph.Status == Allocated {
-		mem += m.Graph.Size
+	for _, m := range mem {
+		sum += m
 	}
 
-	return mem
+	return sum
 }
 
-func memoryPresent(mem []Memory) bool {
-	return slices.ContainsFunc(mem, func(m Memory) bool { return m.Size != 0 })
+// Size returns the total size of the memory required by this device
+func (m DeviceMemory) Size() uint64 {
+	return sumMemory(m.Weights) + sumMemory(m.Cache) + m.Graph
+}
+
+func memoryPresent(mem []uint64) bool {
+	return slices.ContainsFunc(mem, func(m uint64) bool { return m != 0 })
 }
 
 func (m DeviceMemory) LogValue() slog.Value {
@@ -249,7 +208,7 @@ func (m DeviceMemory) LogValue() slog.Value {
 		attrs = append(attrs, slog.Any("Cache", m.Cache))
 	}
 
-	if m.Graph.Size != 0 {
+	if m.Graph != 0 {
 		attrs = append(attrs, slog.Any("Graph", m.Graph))
 	}
 
@@ -267,7 +226,7 @@ func (m DeviceMemory) LogValue() slog.Value {
 // accommodate that to make forward progress.
 type BackendMemory struct {
 	// InputWeights are always located on the CPU and cannot be moved
-	InputWeights Memory
+	InputWeights uint64
 
 	// CPU model components are located in system memory. This does not
 	// include unified memory allocated through the GPU.
@@ -279,7 +238,7 @@ type BackendMemory struct {
 
 func (m BackendMemory) LogValue() slog.Value {
 	var attrs []slog.Attr
-	if m.InputWeights.Size != 0 {
+	if m.InputWeights != 0 {
 		attrs = append(attrs, slog.Any("InputWeights", m.InputWeights))
 	}
 
@@ -291,17 +250,7 @@ func (m BackendMemory) LogValue() slog.Value {
 	return slog.GroupValue(attrs...)
 }
 
-func sumMemory(mem []Memory) uint64 {
-	var sum uint64
-
-	for _, m := range mem {
-		sum += m.Size
-	}
-
-	return sum
-}
-
-// Log prints a high level summary of the memory (allocated or not)
+// Log prints a high level summary of the memory
 func (m BackendMemory) Log(level slog.Level) {
 	var total uint64
 
@@ -311,7 +260,7 @@ func (m BackendMemory) Log(level slog.Level) {
 			total += sum
 		}
 	}
-	if sum := m.InputWeights.Size + sumMemory(m.CPU.Weights); sum > 0 {
+	if sum := m.InputWeights + sumMemory(m.CPU.Weights); sum > 0 {
 		slog.Log(context.TODO(), level, "model weights", "device", m.CPU.Name, "size", format.HumanBytes2(sum))
 		total += sum
 	}
@@ -328,12 +277,12 @@ func (m BackendMemory) Log(level slog.Level) {
 	}
 
 	for _, gpu := range m.GPUs {
-		if sum := gpu.Graph.Size; sum > 0 {
+		if sum := gpu.Graph; sum > 0 {
 			slog.Log(context.TODO(), level, "compute graph", "device", gpu.Name, "size", format.HumanBytes2(sum))
 			total += sum
 		}
 	}
-	if sum := m.CPU.Graph.Size; sum > 0 {
+	if sum := m.CPU.Graph; sum > 0 {
 		slog.Log(context.TODO(), level, "compute graph", "device", m.CPU.Name, "size", format.HumanBytes2(sum))
 		total += sum
 	}

ml/backend/ggml/ggml.go

@@ -159,7 +159,6 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 			C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
 			bt := C.ggml_backend_dev_buffer_type(d)
 			cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, bt)
-			C.ggml_backend_buft_set_alloc(bt, C.bool(params.AllocMemory))
 
 			btDeviceMemory[C.ggml_backend_dev_buffer_type(d)] = &requiredMemory.CPU
 		}
@@ -169,8 +168,8 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 	var props C.struct_ggml_backend_dev_props
 	C.ggml_backend_dev_get_props(cpuDeviceBufferType.d, &props)
 	requiredMemory.CPU.ID = C.GoString(props.id)
-	requiredMemory.CPU.Weights = make([]ml.Memory, blocks+1)
-	requiredMemory.CPU.Cache = make([]ml.Memory, blocks+1)
+	requiredMemory.CPU.Weights = make([]uint64, blocks+1)
+	requiredMemory.CPU.Cache = make([]uint64, blocks+1)
 
 	// create list of buffer types for each gpu
 	var gpuDeviceBufferTypes []deviceBufferType
@@ -181,15 +180,14 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 			d:   d,
 			bts: append([]C.ggml_backend_buffer_type_t{bt}, cpuDeviceBufferType.bts...),
 		})
-		C.ggml_backend_buft_set_alloc(bt, C.bool(params.AllocMemory))
 
 		btDeviceMemory[bt] = &requiredMemory.GPUs[i]
 		requiredMemory.GPUs[i].Name = C.GoString(C.ggml_backend_dev_name(d))
 		var props C.struct_ggml_backend_dev_props
 		C.ggml_backend_dev_get_props(d, &props)
 		requiredMemory.GPUs[i].ID = C.GoString(props.id)
-		requiredMemory.GPUs[i].Weights = make([]ml.Memory, blocks+1)
-		requiredMemory.GPUs[i].Cache = make([]ml.Memory, blocks+1)
+		requiredMemory.GPUs[i].Weights = make([]uint64, blocks+1)
+		requiredMemory.GPUs[i].Cache = make([]uint64, blocks+1)
 	}
 
 	// inputs always use cpu
@@ -275,13 +273,9 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 
 			size := pad(C.ggml_backend_buft_get_alloc_size(bt, tt), C.ggml_backend_buft_get_alignment(bt))
 			if layer == -1 {
-				// Assume that InputWeights can be allocated - they're always in system memory and can't be moved in any case
-				if params.AllocMemory {
-					requiredMemory.InputWeights.Status = ml.Allocated
-				}
-				requiredMemory.InputWeights.Size += uint64(size)
+				requiredMemory.InputWeights += uint64(size)
 			} else {
-				btDeviceMemory[bt].Weights[layer].Size += uint64(size)
+				btDeviceMemory[bt].Weights[layer] += uint64(size)
 			}
 
 			//nolint:staticcheck // TODO: check if buffer type supports this tensor
@@ -341,47 +335,6 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 		}
 	}
 
-	// allocate buffers for each context
-	bbs := make(map[*C.struct_ggml_context]C.ggml_backend_buffer_t, len(ctxs))
-	for bt, c := range ctxs {
-		if C.ggml_get_first_tensor(c) == nil {
-			continue
-		}
-
-		b := C.ggml_backend_alloc_ctx_tensors_from_buft(c, bt)
-		if params.AllocMemory {
-			for i := range btDeviceMemory[bt].Weights {
-				if btDeviceMemory[bt].Weights[i].Size != 0 {
-					if b != nil {
-						btDeviceMemory[bt].Weights[i].Status = ml.Allocated
-					} else {
-						btDeviceMemory[bt].Weights[i].Status = ml.Failed
-					}
-				}
-			}
-		}
-
-		if b == nil {
-			for _, b := range bbs {
-				C.ggml_backend_buffer_free(b)
-			}
-
-			for _, ctx := range ctxs {
-				C.ggml_free(ctx)
-			}
-
-			panic(ml.ErrNoMem{BackendMemory: requiredMemory})
-		}
-
-		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
-		bbs[c] = b
-	}
-
-	for bs := range maps.Values(bbs) {
-		logutil.Trace("model weights", "buffer", C.GoString(C.ggml_backend_buffer_name(bs)),
-			"size", format.HumanBytes2(uint64(C.ggml_backend_buffer_get_size(bs))))
-	}
-
 	// map tensor names to tensors for easy lookup later
 	tensors := make(map[string]*C.struct_ggml_tensor)
 	for _, c := range ctxs {
@@ -419,6 +372,46 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 	}
 
 	maxGraphNodes := max(8192, len(meta.Tensors().Items())*5)
+
+	sched := C.ggml_backend_sched_new_ext(
+		(*C.ggml_backend_t)(unsafe.Pointer(&schedBackends[0])),
+		(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&schedBufts[0])),
+		C.int(len(schedBackends)),
+		C.size_t(maxGraphNodes),
+		C._Bool(false),
+		C._Bool(false),
+		C._Bool(params.AllocMemory),
+	)
+
+	// allocate buffers for each context
+	bbs := make(map[*C.struct_ggml_context]C.ggml_backend_buffer_t, len(ctxs))
+	for bt, c := range ctxs {
+		if C.ggml_get_first_tensor(c) == nil {
+			continue
+		}
+
+		b := C.ggml_backend_alloc_ctx_tensors_from_buft(c, bt)
+		if b == nil {
+			for _, b := range bbs {
+				C.ggml_backend_buffer_free(b)
+			}
+
+			for _, ctx := range ctxs {
+				C.ggml_free(ctx)
+			}
+
+			panic(ml.ErrNoMem{BackendMemory: requiredMemory})
+		}
+
+		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
+		bbs[c] = b
+	}
+
+	for bs := range maps.Values(bbs) {
+		logutil.Trace("model weights", "buffer", C.GoString(C.ggml_backend_buffer_name(bs)),
+			"size", format.HumanBytes2(uint64(C.ggml_backend_buffer_get_size(bs))))
+	}
+
 	return &Backend{
 		modelPath:         modelPath,
 		allocMemory:       params.AllocMemory,
@@ -426,18 +419,11 @@ func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
 		meta:              meta,
 		tensorLoadTargets: targets,
 		tensors:           tensors,
-		sched: C.ggml_backend_sched_new(
-			(*C.ggml_backend_t)(unsafe.Pointer(&schedBackends[0])),
-			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&schedBufts[0])),
-			C.int(len(schedBackends)),
-			C.size_t(maxGraphNodes),
-			C._Bool(false),
-			C._Bool(false),
-		),
-		schedBackends: schedBackends,
-		schedBufts:    schedBufts,
-		input:         deviceBufferTypes[input.d],
-		output:        output.d,
+		sched:             sched,
+		schedBackends:     schedBackends,
+		schedBufts:        schedBufts,
+		input:             deviceBufferTypes[input.d],
+		output:            output.d,
 		layers: func() map[int]layerDevice {
 			m := make(map[int]layerDevice)
 			for i, layer := range layers {
@@ -795,24 +781,15 @@ func (c *Context) Reserve() {
 
 	// Reserve may get called multiple times for different graphs - we just want the last run, which will contain the max allocations
 	for _, bt := range c.b.schedBufts {
-		c.b.btDeviceMemory[bt].Graph = ml.Memory{}
+		c.b.btDeviceMemory[bt].Graph = 0
 	}
 
 	for i := range c.b.schedBackends {
-		bufferStatus := C.ggml_backend_sched_get_attempted_buffer_size(c.b.sched, c.b.schedBackends[i])
-
-		graph := &c.b.btDeviceMemory[c.b.schedBufts[i]].Graph
-		graph.Size += uint64(bufferStatus.size)
-		if c.b.allocMemory {
-			if bufferStatus.allocated && graph.Status != ml.Failed {
-				graph.Status = ml.Allocated
-			} else {
-				graph.Status = ml.Failed
-			}
-		}
+		bufferSize := C.ggml_backend_sched_get_attempted_buffer_size(c.b.sched, c.b.schedBackends[i])
+		c.b.btDeviceMemory[c.b.schedBufts[i]].Graph += uint64(bufferSize)
 
 		logutil.Trace("compute graph", "backend", C.GoString(C.ggml_backend_name(c.b.schedBackends[i])),
-			"buffer_type", C.GoString(C.ggml_backend_buft_name(c.b.schedBufts[i])), "size", format.HumanBytes2(uint64(bufferStatus.size)))
+			"buffer_type", C.GoString(C.ggml_backend_buft_name(c.b.schedBufts[i])), "size", format.HumanBytes2(uint64(bufferSize)))
 	}
 
 	if !reserved {
@@ -862,16 +839,7 @@ func (c *Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
 
 	b := C.ggml_backend_buft_alloc_buffer(c.buft, size)
 	if c.layer >= 0 {
-		cache := &c.b.btDeviceMemory[c.buft].Cache[c.layer]
-
-		cache.Size += uint64(size)
-		if c.b.allocMemory {
-			if b != nil {
-				cache.Status = ml.Allocated
-			} else {
-				cache.Status = ml.Failed
-			}
-		}
+		c.b.btDeviceMemory[c.buft].Cache[c.layer] += uint64(size)
 	}
 
 	if b == nil {

ml/backend/ggml/ggml/include/ggml-alloc.h

@@ -65,12 +65,7 @@ GGML_API bool ggml_gallocr_reserve_n(
 GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
 
 GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
-
-struct ggml_allocr_buffer_status {
-    size_t size;
-    bool allocated;
-};
-GGML_API struct ggml_allocr_buffer_status ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id);
+GGML_API size_t ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id);
 
 // Utils
 // Create a buffer and allocate all the tensors in a ggml_context

ml/backend/ggml/ggml/include/ggml-backend.h

@@ -35,7 +35,6 @@ extern "C" {
     //
 
     GGML_API const char *          ggml_backend_buft_name          (ggml_backend_buffer_type_t buft);
-    GGML_API void                  ggml_backend_buft_set_alloc     (ggml_backend_buffer_type_t buft, bool alloc);
     GGML_API ggml_backend_buffer_t ggml_backend_buft_alloc_buffer  (ggml_backend_buffer_type_t buft, size_t size);
     GGML_API size_t                ggml_backend_buft_get_alignment (ggml_backend_buffer_type_t buft);
     GGML_API size_t                ggml_backend_buft_get_max_size  (ggml_backend_buffer_type_t buft);
@@ -293,6 +292,7 @@ extern "C" {
 
     // Initialize a backend scheduler, backends with low index are given priority over backends with high index
     GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload);
+    GGML_API ggml_backend_sched_t ggml_backend_sched_new_ext(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload, bool alloc_buffers);
     GGML_API void                 ggml_backend_sched_free(ggml_backend_sched_t sched);
 
     // Initialize backend buffers from a measure graph
@@ -306,12 +306,7 @@ extern "C" {
     GGML_API int                  ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched);
 
     GGML_API size_t               ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
-
-    struct ggml_backend_buffer_status {
-        size_t size;
-        bool allocated;
-    };
-    GGML_API struct ggml_backend_buffer_status ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
+    GGML_API size_t               ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
 
     GGML_API void                 ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
     GGML_API ggml_backend_t       ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);

ml/backend/ggml/ggml/src/ggml-alloc.c

@@ -932,7 +932,7 @@ size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
     return ggml_backend_buffer_get_size(galloc->buffers[buffer_id]);
 }
 
-struct ggml_allocr_buffer_status ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
+size_t ggml_gallocr_get_attempted_buffer_size(ggml_gallocr_t galloc, int buffer_id) {
     GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers);
 
     for (int i = 0; i < buffer_id; i++) {
@@ -941,13 +941,11 @@ struct ggml_allocr_buffer_status ggml_gallocr_get_attempted_buffer_size(ggml_gal
             // (See above.) However, we need a different check because multiple buffers might be NULL in our
             // case and we still want to know the attempted size.
 
-            struct ggml_allocr_buffer_status status = {0, true};
-            return status;
+            return 0;
         }
     }
 
-    struct ggml_allocr_buffer_status status = {galloc->buffer_sizes[buffer_id], galloc->buffers[buffer_id] != NULL};
-    return status;
+    return galloc->buffer_sizes[buffer_id];
 }
 
 // utils

ml/backend/ggml/ggml/src/ggml-backend-impl.h

@@ -26,6 +26,10 @@ extern "C" {
         size_t                (*get_alloc_size)(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor);
         // (optional) check if tensor data is in host memory and uses standard ggml tensor layout (defaults to false)
         bool                  (*is_host)       (ggml_backend_buffer_type_t buft);
+
+        // (optional) returns a dummy buffer that is equivalent to one created by alloc_buffer but without actually being backed
+        // by memory
+        ggml_backend_buffer_t (*noalloc_buffer)(ggml_backend_buffer_type_t buft, size_t size);
     };
 
     struct ggml_backend_buffer_type {
@@ -116,6 +120,16 @@ extern "C" {
         void (*event_record)(ggml_backend_t backend, ggml_backend_event_t event);
         // wait for an event on on a different stream
         void (*event_wait)  (ggml_backend_t backend, ggml_backend_event_t event);
+
+        // (optional) reserves intermediate buffers needed for the compution
+        // if alloc is true, memory is actually allocated, otherwise the required amount is just returned by buffer_size
+        enum ggml_status          (*graph_reserve)     (ggml_backend_t backend, struct ggml_cgraph * cgraph, bool alloc);
+
+        // (optional) returns the memory needed after calling graph_reserve
+        size_t                    (*buffer_size)       (ggml_backend_t backend);
+
+        // (optional) frees memory from intermediate buffers that was allocated either by graph_compute or graph_reserve
+        void                      (*reset)             (ggml_backend_t backend);
     };
 
     struct ggml_backend {

ml/backend/ggml/ggml/src/ggml-backend.cpp

@@ -35,10 +35,6 @@ const char * ggml_backend_buft_name(ggml_backend_buffer_type_t buft) {
     return buft->iface.get_name(buft);
 }
 
-void ggml_backend_buft_set_alloc(ggml_backend_buffer_type_t buft, bool alloc) {
-    buft->no_alloc = !alloc;
-}
-
 ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
     if (size == 0) {
         // return a dummy buffer for zero-sized allocations
@@ -46,7 +42,14 @@ ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t
     }
 
     if (buft->no_alloc) {
-        ggml_backend_buffer_t buf = ggml_backend_buffer_init(buft, {}, NULL, size);
+        ggml_backend_buffer_t buf;
+
+        if (buft->iface.noalloc_buffer != NULL) {
+            buf = buft->iface.noalloc_buffer(buft, size);
+        } else {
+            buf = ggml_backend_buffer_init(buft, {}, NULL, size);
+        }
+
         buf->no_alloc = true;
         return buf;
     }
@@ -688,6 +691,12 @@ struct ggml_backend_sched {
     bool op_offload;
 
     int debug;
+
+    // allocate buffers on attached ggml_backend_buffer_type_t's and during reservation
+    // if false, dummy buffers are used for faster memory sizing calculations
+    // the scheduler needs to be recreated with allocated buffers before it can be used
+    // for computation
+    bool alloc_buffers;
 };
 
 #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
@@ -1474,6 +1483,17 @@ ggml_backend_sched_t ggml_backend_sched_new(
         size_t graph_size,
         bool parallel,
         bool op_offload) {
+            return ggml_backend_sched_new_ext(backends, bufts, n_backends, graph_size, parallel, op_offload, true);
+        }
+
+ggml_backend_sched_t ggml_backend_sched_new_ext(
+        ggml_backend_t * backends,
+        ggml_backend_buffer_type_t * bufts,
+        int n_backends,
+        size_t graph_size,
+        bool parallel,
+        bool op_offload,
+        bool alloc_buffers) {
     GGML_ASSERT(n_backends > 0);
     GGML_ASSERT(n_backends <= GGML_SCHED_MAX_BACKENDS);
     GGML_ASSERT(ggml_backend_dev_type(ggml_backend_get_device(backends[n_backends - 1])) == GGML_BACKEND_DEVICE_TYPE_CPU);
@@ -1515,10 +1535,13 @@ ggml_backend_sched_t ggml_backend_sched_new(
                 sched->events[b][c] = ggml_backend_event_new(backends[b]->device);
             }
         }
+
+        sched->bufts[b]->no_alloc = !alloc_buffers;
     }
 
     sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
     sched->op_offload = op_offload;
+    sched->alloc_buffers = alloc_buffers;
 
     ggml_backend_sched_reset(sched);
 
@@ -1533,6 +1556,10 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
         for (int c = 0; c < sched->n_copies; c++) {
             ggml_backend_event_free(sched->events[b][c]);
         }
+
+        if (sched->backends[b]->iface.reset != NULL) {
+            sched->backends[b]->iface.reset(sched->backends[b]);
+        }
     }
     ggml_gallocr_free(sched->galloc);
     ggml_free(sched->ctx);
@@ -1572,6 +1599,24 @@ bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph *
         return false;
     }
 
+    if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
+        return false;
+    }
+
+    struct ggml_backend_sched_split * splits = sched->splits;
+    for (int i = 0; i < sched->n_splits; i++) {
+        struct ggml_backend_sched_split * split = &splits[i];
+        int split_backend_id = split->backend_id;
+        ggml_backend_t split_backend = sched->backends[split_backend_id];
+
+        if (split_backend->iface.graph_reserve != NULL) {
+            enum ggml_status ec = split_backend->iface.graph_reserve(split_backend, &split->graph, sched->alloc_buffers);
+            if (ec != GGML_STATUS_SUCCESS) {
+                return false;
+            }
+        }
+    }
+
     ggml_backend_sched_reset(sched);
 
     return true;
@@ -1656,14 +1701,17 @@ size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backe
     return ggml_gallocr_get_buffer_size(sched->galloc, backend_index);
 }
 
-struct ggml_backend_buffer_status ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) {
+size_t ggml_backend_sched_get_attempted_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) {
     int backend_index = ggml_backend_sched_backend_id(sched, backend);
     GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
 
-    struct ggml_allocr_buffer_status allocr_status = ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
-    struct ggml_backend_buffer_status status = {allocr_status.size, allocr_status.allocated};
+    size_t size = ggml_gallocr_get_attempted_buffer_size(sched->galloc, backend_index);
 
-    return status;
+    if (backend->iface.buffer_size != NULL) {
+        size += backend->iface.buffer_size(backend);
+    }
+
+    return size;
 }
 
 void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) {

ml/backend/ggml/ggml/src/ggml-cuda/common.cuh

@@ -35,6 +35,31 @@
 #include "vendors/cuda.h"
 #endif // defined(GGML_USE_HIP)
 
+extern bool reserving_graph;
+
+// If we are reserving the graph, pointers might be invalid and will fail if cudaMemcpyAsync tries to validate them.
+// However, since we don't actually expect a result, we don't need to actually do the memcpy.
+static cudaError_t cudaMemcpyAsyncReserve ( void* dst, const void* src, size_t count, cudaMemcpyKind kind, cudaStream_t stream = 0 ) {
+    if (!reserving_graph) {
+        return cudaMemcpyAsync(dst, src, count, kind, stream);
+    } else {
+        return cudaSuccess;
+    }
+}
+
+static cudaError_t cudaMemcpy2DAsyncReserve ( void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, cudaMemcpyKind kind, cudaStream_t stream = 0 ) {
+    if (!reserving_graph) {
+        return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, kind, stream);
+    } else {
+        return cudaSuccess;
+    }
+}
+
+#undef cudaMemcpyAsync
+#define cudaMemcpyAsync cudaMemcpyAsyncReserve
+#undef cudaMemcpy2DAsync
+#define cudaMemcpy2DAsync cudaMemcpy2DAsyncReserve
+
 #define STRINGIZE_IMPL(...) #__VA_ARGS__
 #define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__)
 
@@ -771,6 +796,9 @@ struct ggml_cuda_pool {
 
     virtual void * alloc(size_t size, size_t * actual_size) = 0;
     virtual void free(void * ptr, size_t size) = 0;
+
+    virtual bool alloc_memory() = 0;
+    virtual size_t alloc_size() = 0;
 };
 
 template<typename T>
@@ -914,11 +942,11 @@ struct ggml_backend_cuda_context {
     // pool
     std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES];
 
-    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device);
+    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, bool alloc);
 
     ggml_cuda_pool & pool(int device) {
         if (pools[device] == nullptr) {
-            pools[device] = new_pool_for_device(device);
+            pools[device] = new_pool_for_device(device, true);
         }
         return *pools[device];
     }
@@ -926,4 +954,20 @@ struct ggml_backend_cuda_context {
     ggml_cuda_pool & pool() {
         return pool(device);
     }
+
+    void pool_set_alloc(bool alloc) {
+        GGML_ASSERT(pools[device] == nullptr || pools[device]->alloc_memory() == alloc);
+
+        if (pools[device] == nullptr) {
+            pools[device] = new_pool_for_device(device, alloc);
+        }
+    }
+
+    size_t pool_get_alloc_size() {
+        if (pools[device] == nullptr) {
+            return 0;
+        }
+
+        return pools[device]->alloc_size();
+    }
 };

ml/backend/ggml/ggml/src/ggml-cuda/ggml-cuda.cu

@@ -355,6 +355,8 @@ const ggml_cuda_device_info & ggml_cuda_info() {
 
 // #define DEBUG_CUDA_MALLOC
 
+#define CUDA_ALIGNMENT 128
+
 // buffer pool for cuda (legacy)
 struct ggml_cuda_pool_leg : public ggml_cuda_pool {
     static const int MAX_BUFFERS = 256;
@@ -367,9 +369,12 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
 
     ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {};
     size_t pool_size = 0;
+    bool allocate = true;
+    size_t last_alloc = 0;
 
-    explicit ggml_cuda_pool_leg(int device) :
-        device(device) {
+    explicit ggml_cuda_pool_leg(int device, bool alloc) :
+        device(device),
+        allocate(alloc) {
     }
 
     ~ggml_cuda_pool_leg() {
@@ -377,7 +382,9 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         for (int i = 0; i < MAX_BUFFERS; ++i) {
             ggml_cuda_buffer & b = buffer_pool[i];
             if (b.ptr != nullptr) {
-                CUDA_CHECK(cudaFree(b.ptr));
+                if (allocate) {
+                    CUDA_CHECK(cudaFree(b.ptr));
+                }
                 pool_size -= b.size;
             }
         }
@@ -425,8 +432,15 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         void * ptr;
         size_t look_ahead_size = (size_t) (1.05 * size);
         look_ahead_size = 256 * ((look_ahead_size + 255)/256);
-        ggml_cuda_set_device(device);
-        CUDA_CHECK(ggml_cuda_device_malloc(&ptr, look_ahead_size, device));
+        if (allocate) {
+            ggml_cuda_set_device(device);
+            if (ggml_cuda_device_malloc(&ptr, look_ahead_size, device) != cudaSuccess) {
+                    last_alloc = look_ahead_size;
+                    throw std::bad_alloc();
+            }
+        } else {
+            ptr = (void *)CUDA_ALIGNMENT;
+        }
         *actual_size = look_ahead_size;
         pool_size += look_ahead_size;
 #ifdef DEBUG_CUDA_MALLOC
@@ -446,10 +460,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
             }
         }
         GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n");
-        ggml_cuda_set_device(device);
-        CUDA_CHECK(cudaFree(ptr));
+        if (allocate) {
+            ggml_cuda_set_device(device);
+            CUDA_CHECK(cudaFree(ptr));
+        }
         pool_size -= size;
     }
+
+    bool alloc_memory() override {
+        return allocate;
+    }
+
+    size_t alloc_size() override {
+        return pool_size + last_alloc;
+    }
 };
 
 // pool with virtual memory
@@ -461,18 +485,24 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
     CUdeviceptr pool_addr = 0;
     size_t pool_used = 0;
     size_t pool_size = 0;
+    bool allocate = true;
+    size_t last_alloc = 0;
     size_t granularity;
 #if defined(GGML_USE_HIP)
     std::vector<std::pair<CUdeviceptr, size_t>> mappings;
 #endif
 
-    explicit ggml_cuda_pool_vmm(int device) :
+    explicit ggml_cuda_pool_vmm(int device, bool alloc) :
         device(device),
-        granularity(ggml_cuda_info().devices[device].vmm_granularity) {
+        granularity(ggml_cuda_info().devices[device].vmm_granularity),
+        allocate(alloc) {
+        if (!allocate) {
+            pool_addr = (CUdeviceptr)CUDA_ALIGNMENT;
+        }
     }
 
     ~ggml_cuda_pool_vmm() {
-        if (pool_addr != 0) {
+        if (pool_addr != 0 && allocate) {
 #if defined(GGML_USE_HIP)
             // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285
             for (std::pair<CUdeviceptr, size_t> & mapping : mappings) {
@@ -499,36 +529,50 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 
             GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE);
 
-            // allocate more physical memory
-            CUmemAllocationProp prop = {};
-            prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
-            prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-            prop.location.id = device;
-            CUmemGenericAllocationHandle handle;
-            CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0));
+            if (allocate) {
+                // allocate more physical memory
+                CUmemAllocationProp prop = {};
+                prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+                prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+                prop.location.id = device;
+                CUmemGenericAllocationHandle handle;
+                if (cuMemCreate(&handle, reserve_size, &prop, 0) != CUDA_SUCCESS) {
+                    last_alloc = reserve_size;
+                    throw std::bad_alloc();
+                }
 
-            // reserve virtual address space (if not already reserved)
-            if (pool_addr == 0) {
-                CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0));
+                // reserve virtual address space (if not already reserved)
+                if (pool_addr == 0) {
+                    CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0));
+                }
+
+                // map at the end of the pool
+                CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size);
+                if (cuMemMap(start_ptr, reserve_size, 0, handle, 0) != CUDA_SUCCESS) {
+                    last_alloc = reserve_size;
+                    CU_CHECK(cuMemRelease(handle));
+                    throw std::bad_alloc();
+                }
+
+                // the memory allocation handle is no longer needed after mapping
+                CU_CHECK(cuMemRelease(handle));
+
+                // set access
+                CUmemAccessDesc access = {};
+                access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+                access.location.id = device;
+                access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+                if (cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1) != CUDA_SUCCESS) {
+                    CU_CHECK(cuMemUnmap(start_ptr, reserve_size));
+                    last_alloc = reserve_size;
+                    throw std::bad_alloc();
+                }
+
+    #if defined(GGML_USE_HIP)
+                mappings.push_back({start_ptr, reserve_size});
+    #endif
             }
 
-            // map at the end of the pool
-            CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size);
-            CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0));
-#if defined(GGML_USE_HIP)
-            mappings.push_back({start_ptr, reserve_size});
-#endif
-
-            // the memory allocation handle is no longer needed after mapping
-            CU_CHECK(cuMemRelease(handle));
-
-            // set access
-            CUmemAccessDesc access = {};
-            access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-            access.location.id = device;
-            access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
-            CU_CHECK(cuMemSetAccess((CUdeviceptr)((char *)(pool_addr) + pool_size), reserve_size, &access, 1));
-
             // add to the pool
             pool_size += reserve_size;
 
@@ -560,16 +604,24 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
         // all deallocations must be in reverse order of the allocations
         GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used));
     }
+
+    bool alloc_memory() override {
+        return allocate;
+    }
+
+    size_t alloc_size() override {
+        return pool_size + last_alloc;
+    }
 };
 #endif // defined(GGML_USE_VMM)
 
-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device) {
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, bool alloc) {
 #if defined(GGML_USE_VMM)
     if (ggml_cuda_info().devices[device].vmm) {
-        return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device));
+        return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device, alloc));
     }
 #endif // defined(GGML_USE_VMM)
-    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device));
+    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device, alloc));
 }
 
 // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
@@ -753,11 +805,20 @@ static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_bac
 }
 
 static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
-    return 128;
+    return CUDA_ALIGNMENT;
 
     GGML_UNUSED(buft);
 }
 
+static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_noalloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
+    ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context;
+
+    void * dev_ptr = (void *)ggml_backend_cuda_buffer_type_get_alignment(buft);
+    ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr);
+
+    return ggml_backend_buffer_init(buft, {}, ctx, size);
+}
+
 static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
     size_t size = ggml_nbytes(tensor);
     int64_t ne0 = tensor->ne[0];
@@ -781,6 +842,7 @@ static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface
     /* .get_max_size     = */ NULL, // defaults to SIZE_MAX
     /* .get_alloc_size   = */ ggml_backend_cuda_buffer_type_get_alloc_size,
     /* .is_host          = */ NULL,
+    /* .noalloc_buffer   = */ ggml_backend_cuda_buffer_type_noalloc_buffer,
 };
 
 ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) {
@@ -2941,6 +3003,7 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx,
 
 static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph,
     bool & graph_evaluated_or_captured, bool & use_cuda_graph, bool & cuda_graph_update_required) {
+
     // flag used to determine whether it is an integrated_gpu
     const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated;
 
@@ -2956,6 +3019,11 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
                     continue;
                 }
 
+                // When reserving, we are forcing CUDA graphs but this operation is not graph-safe so we need to skip it
+                if (reserving_graph && node->op == GGML_OP_MUL_MAT_ID && node->ne[2] != 1) {
+                    continue;
+                }
+
                 static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr);
                 if (!disable_fusion) {
                     if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
@@ -3027,6 +3095,7 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
 
 static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
+    cuda_ctx->pool_set_alloc(true);
 
     ggml_cuda_set_device(cuda_ctx->device);
 
@@ -3106,6 +3175,71 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
     return GGML_STATUS_SUCCESS;
 }
 
+// This is used to skip operations that are not graph safe during the reservation process.
+bool reserving_graph = false;
+
+static enum ggml_status ggml_backend_cuda_graph_reserve(ggml_backend_t backend, ggml_cgraph * cgraph, bool alloc) {
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
+    cuda_ctx->pool_set_alloc(alloc);
+
+    #ifdef USE_CUDA_GRAPH
+    if (cuda_ctx->cuda_graph == nullptr) {
+        cuda_ctx->cuda_graph.reset(new ggml_cuda_graph());
+    }
+    #endif
+
+    ggml_cuda_set_device(cuda_ctx->device);
+
+    {
+        std::lock_guard<std::mutex> lock(ggml_cuda_lock);
+        ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed);
+    }
+
+    reserving_graph = true;
+
+    // Create CuBLAS handles early to avoid synchronous allocations during graph capture.
+    cuda_ctx->cublas_handle();
+
+    CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed));
+
+    enum ggml_status result = GGML_STATUS_SUCCESS;
+
+    try {
+        bool use_cuda_graph = false;
+        bool cuda_graph_update_required = false;
+        bool graph_evaluated_or_captured = false;
+
+        evaluate_and_capture_cuda_graph(cuda_ctx, cgraph, graph_evaluated_or_captured, use_cuda_graph, cuda_graph_update_required);
+    } catch (const std::exception &e) {
+        result = GGML_STATUS_FAILED;
+    }
+
+    cudaGraph_t graph;
+    CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph));
+    CUDA_CHECK(cudaGraphDestroy(graph));
+
+    reserving_graph = false;
+
+    {
+        std::lock_guard<std::mutex> lock(ggml_cuda_lock);
+        if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) {
+            ggml_cuda_lock_cv.notify_all();
+        }
+    }
+
+    return result;
+}
+
+static size_t ggml_backend_cuda_buffer_size(ggml_backend_t backend) {
+    ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *)backend->context;
+    return ctx->pool_get_alloc_size();
+}
+
+static void ggml_backend_cuda_reset(ggml_backend_t backend) {
+    ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *)backend->context;
+    ctx->pools[ctx->device] = NULL;
+}
+
 static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) {
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
 
@@ -3145,6 +3279,9 @@ static const ggml_backend_i ggml_backend_cuda_interface = {
     /* .graph_compute           = */ ggml_backend_cuda_graph_compute,
     /* .event_record            = */ ggml_backend_cuda_event_record,
     /* .event_wait              = */ ggml_backend_cuda_event_wait,
+    /* .graph_reserve           = */ ggml_backend_cuda_graph_reserve,
+    /* .buffer_size             = */ ggml_backend_cuda_buffer_size,
+    /* .reset                   = */ ggml_backend_cuda_reset,
 };
 
 static ggml_guid_t ggml_backend_cuda_guid() {

ml/backend/ggml/ggml/src/ggml-cuda/scale.cu

@@ -1,18 +1,19 @@
 #include "scale.cuh"
 
-static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int k) {
-    const int i = blockDim.x*blockIdx.x + threadIdx.x;
+#define MAX_GRIDDIM_X 0x7FFFFFFF
 
-    if (i >= k) {
-        return;
+static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int64_t nelements) {
+    int64_t tid = (int64_t)blockIdx.x * (int64_t)blockDim.x + (int64_t)threadIdx.x;
+    int64_t stride = (int64_t)blockDim.x * (int64_t)gridDim.x;
+
+    for (int64_t i = tid; i < nelements; i += stride) {
+        dst[i] = scale * x[i] + bias;
     }
-
-    dst[i] = scale * x[i] + bias;
 }
 
-static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int k, cudaStream_t stream) {
-    const int num_blocks = (k + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
-    scale_f32<<<num_blocks, CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, k);
+static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
+    const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
+    scale_f32<<<MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream>>>(x, dst, scale, bias, nelements);
 }
 
 void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

model/bytepairencoding.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"iter"
 	"log/slog"
+	"slices"
 	"strings"
 
 	"github.com/dlclark/regexp2"
@@ -13,16 +14,28 @@ import (
 )
 
 type BytePairEncoding struct {
-	pre   *regexp2.Regexp
-	vocab *Vocabulary
+	vocab   *Vocabulary
+	regexps []*regexp2.Regexp
 }
 
 var _ TextProcessor = (*BytePairEncoding)(nil)
 
-func NewBytePairEncoding(pre string, vocab *Vocabulary) BytePairEncoding {
+func NewBytePairEncoding(vocab *Vocabulary, pretokenizers ...string) BytePairEncoding {
+	if len(pretokenizers) == 0 {
+		// set default byte-level pretokenizer if none provided, e.g.
+		// https://github.com/huggingface/tokenizers/blob/main/tokenizers/src/pre_tokenizers/byte_level.rs#L44
+		pretokenizers = []string{`'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`}
+	}
+
 	return BytePairEncoding{
-		pre:   regexp2.MustCompile(pre, regexp2.None),
 		vocab: vocab,
+		regexps: slices.Collect(func(yield func(*regexp2.Regexp) bool) {
+			for _, p := range pretokenizers {
+				if !yield(regexp2.MustCompile(p, regexp2.RE2)) {
+					return
+				}
+			}
+		}),
 	}
 }
 
@@ -35,13 +48,36 @@ func (bpe BytePairEncoding) Is(id int32, special Special) bool {
 }
 
 func (bpe *BytePairEncoding) split(s string) iter.Seq[string] {
-	return func(yield func(string) bool) {
-		for m, _ := bpe.pre.FindStringMatch(s); m != nil; m, _ = bpe.pre.FindNextMatch(m) {
-			if !yield(m.String()) {
-				break
+	parts := []string{s}
+	for _, re := range bpe.regexps {
+		parts = slices.Collect(func(yield func(string) bool) {
+			for _, part := range parts {
+				r := []rune(part)
+				var offset int
+				for m, _ := re.FindRunesMatch(r); m != nil; m, _ = re.FindNextMatch(m) {
+					if offset-m.Index != 0 {
+						if !yield(string(r[:m.Index])) {
+							return
+						}
+					}
+
+					if !yield(m.String()) {
+						return
+					}
+
+					offset = m.Index + m.Length
+				}
+
+				if offset < len(r) {
+					if !yield(string(r[offset:])) {
+						return
+					}
+				}
 			}
-		}
+		})
 	}
+
+	return slices.Values(parts)
 }
 
 // fragment is a string fragment and their corresponding token IDs

model/bytepairencoding_test.go

@@ -59,12 +59,12 @@ func llama(t testing.TB) BytePairEncoding {
 	}
 
 	return NewBytePairEncoding(
-		`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		&Vocabulary{
 			Values: tokens,
 			Types:  types,
 			Merges: merges,
 		},
+		"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
 	)
 }
 
@@ -282,3 +282,41 @@ func BenchmarkBytePairEncoding(b *testing.B) {
 		})
 	}
 }
+
+func TestSplit(t *testing.T) {
+	cases := []struct {
+		name string
+		patterns,
+		want []string
+	}{
+		{
+			name: "default",
+			want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " 123", " 一二三"},
+		},
+		{
+			name: "unicode",
+			patterns: []string{
+				"\\p{N}{1,3}",
+				`[一-龥぀-ゟ゠-ヿ]+`,
+				"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+",
+			},
+			want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " ", "123", " ", "一二三"},
+		},
+		{
+			name: "individual digits",
+			patterns: []string{
+				"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+			},
+			want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " ", "1", "2", "3", " 一二三"},
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			tokenizer := NewBytePairEncoding(nil, tt.patterns...)
+			if diff := cmp.Diff(tt.want, slices.Collect(tokenizer.split("Hello, WORLD!! How's it going? 123 一二三"))); diff != "" {
+				t.Errorf("no match (-theirs +ours):\n%s", diff)
+			}
+		})
+	}
+}

model/model.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	_ "image/jpeg"
 	_ "image/png"
+	"log/slog"
 	"os"
 	"reflect"
 	"strconv"
@@ -171,35 +172,44 @@ func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
 			// make a copy
 			tagsCopy := tags
 			if tag := t.Field(i).Tag.Get("gguf"); tag != "" {
-				tagsCopy = append(tagsCopy, ParseTags(tag))
+				tagsCopy = append(tagsCopy, parseTag(tag))
 			}
 
 			if tt == reflect.TypeOf((*Base)(nil)).Elem() {
 				vv.Set(reflect.ValueOf(base))
 			} else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() {
-				var fn func([]Tag) [][]string
-				fn = func(tags []Tag) (names [][]string) {
+				var fn func([]Tag, string, string) [][]string
+				fn = func(tags []Tag, prefix, suffix string) (fullNames [][]string) {
 					if len(tags) > 0 {
-						localNames := []string{tags[0].Name}
-						localNames = append(localNames, tags[0].Alternate...)
-
-						for _, localName := range localNames {
-							fullName := []string{localName}
-							nested := fn(tags[1:])
-							if len(nested) > 0 {
-								for _, rest := range nested {
-									names = append(names, append(fullName, rest...))
+						var names []string
+						if tags[0].name != "" {
+							for _, n := range append([]string{tags[0].name}, tags[0].alternatives...) {
+								names = append(names, prefix+n+suffix)
+							}
+						}
+						childNames := fn(tags[1:], tags[0].prefix, tags[0].suffix)
+						if len(names) == 0 {
+							// current tag has no name, use child names only
+							fullNames = append(fullNames, childNames...)
+						} else if len(childNames) == 0 {
+							// current tag has names but no children, create branches for each name
+							for _, name := range names {
+								fullNames = append(fullNames, []string{name})
+							}
+						} else {
+							// merge each name with each child
+							for _, name := range names {
+								for _, childName := range childNames {
+									fullNames = append(fullNames, append([]string{name}, childName...))
 								}
-							} else {
-								names = append(names, fullName)
 							}
 						}
 					}
 
-					return names
+					return fullNames
 				}
 
-				names := fn(tagsCopy)
+				names := fn(tagsCopy, "", "")
 				for _, name := range names {
 					if tensor := base.Backend().Get(strings.Join(name, ".")); tensor != nil {
 						logutil.Trace("found tensor", "", tensor)
@@ -213,9 +223,9 @@ func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value {
 				for i := range vv.Len() {
 					vvv := vv.Index(i)
 					if vvv.Kind() == reflect.Pointer || vvv.Kind() == reflect.Interface {
-						setPointer(base, vvv, append(tagsCopy, Tag{Name: strconv.Itoa(i)}))
+						setPointer(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)}))
 					} else {
-						vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{Name: strconv.Itoa(i)})...))
+						vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)})...))
 					}
 				}
 			}
@@ -254,18 +264,31 @@ func setPointer(base Base, v reflect.Value, tags []Tag) {
 }
 
 type Tag struct {
-	Name      string
-	Alternate []string
+	name,
+	// prefix and suffix are applied to child tags
+	prefix,
+	suffix string
+	alternatives []string
 }
 
-func ParseTags(s string) (tag Tag) {
+func parseTag(s string) (tag Tag) {
 	parts := strings.Split(s, ",")
 	if len(parts) > 0 {
-		tag.Name = parts[0]
+		tag.name = parts[0]
 
 		for _, part := range parts[1:] {
-			if value, ok := strings.CutPrefix(part, "alt:"); ok {
-				tag.Alternate = append(tag.Alternate, value)
+			if value, ok := strings.CutPrefix(part, "alt:"); ok && tag.name == "" {
+				// elevate alternative to primary if no primary given
+				tag.name = value
+				slog.Warn("gguf tag has alt: but no primary name", "tag", s)
+			} else if ok {
+				tag.alternatives = append(tag.alternatives, value)
+			}
+			if value, ok := strings.CutPrefix(part, "pre:"); ok {
+				tag.prefix = value
+			}
+			if value, ok := strings.CutPrefix(part, "suf:"); ok {
+				tag.suffix = value
 			}
 		}
 	}

model/model_test.go

@@ -22,14 +22,14 @@ func TestParseTags(t *testing.T) {
 		{
 			value: "output",
 			want: Tag{
-				Name: "output",
+				name: "output",
 			},
 		},
 		{
 			value: "output,alt:token_embd",
 			want: Tag{
-				Name: "output",
-				Alternate: []string{
+				name: "output",
+				alternatives: []string{
 					"token_embd",
 				},
 			},
@@ -38,8 +38,8 @@ func TestParseTags(t *testing.T) {
 
 	for _, tt := range cases {
 		t.Run(tt.value, func(t *testing.T) {
-			got := ParseTags(tt.value)
-			if diff := cmp.Diff(tt.want, got); diff != "" {
+			got := parseTag(tt.value)
+			if diff := cmp.Diff(tt.want, got, cmp.AllowUnexported((Tag{}))); diff != "" {
 				t.Errorf("ParseTags() returned unexpected values (-want +got):\n%s", diff)
 			}
 		})
@@ -125,6 +125,7 @@ func TestPopulateFieldsAlternateName(t *testing.T) {
 		Input  *nn.Embedding `gguf:"input"`
 		Output *nn.Linear    `gguf:"output,alt:input"`
 		Nested *nested       `gguf:"nested"`
+		Tensor ml.Tensor     `gguf:"leaf,alt:tensor"`
 	}
 
 	var m fakeModel
@@ -133,6 +134,7 @@ func TestPopulateFieldsAlternateName(t *testing.T) {
 		names: []string{
 			"input.weight",
 			"nested.b.weight",
+			"leaf",
 		},
 	}}, v.Elem()))
 
@@ -142,6 +144,58 @@ func TestPopulateFieldsAlternateName(t *testing.T) {
 		Nested: &nested{
 			Weight: &nn.Linear{Weight: &fakeTensor{Name: "nested.b.weight"}},
 		},
+		Tensor: &fakeTensor{Name: "leaf"},
+	}, m); diff != "" {
+		t.Errorf("populateFields() set incorrect values (-want +got):\n%s", diff)
+	}
+}
+
+func TestPopulateFieldsPrefixSuffixName(t *testing.T) {
+	type fakeBlock struct {
+		A  *nn.Linear `gguf:"a"`
+		B  *nn.Linear `gguf:",pre:b_"`
+		C  *nn.Linear `gguf:",suf:_c"`
+		XY *nn.Linear `gguf:",pre:x_,suf:_y"`
+	}
+
+	type fakeModel struct {
+		Blocks []fakeBlock `gguf:"blk"`
+	}
+
+	m := fakeModel{
+		Blocks: make([]fakeBlock, 2),
+	}
+	v := reflect.ValueOf(&m)
+	v.Elem().Set(populateFields(Base{b: &fakeBackend{
+		names: []string{
+			"blk.0.a.weight",
+			"blk.0.b_weight",
+			"blk.0.b_bias",
+			"blk.0.weight_c",
+			"blk.0.x_weight_y",
+			"blk.1.a.weight",
+			"blk.1.b_weight",
+			"blk.1.b_bias",
+			"blk.1.weight_c",
+			"blk.1.x_weight_y",
+		},
+	}}, v.Elem()))
+
+	if diff := cmp.Diff(fakeModel{
+		Blocks: []fakeBlock{
+			{
+				A:  &nn.Linear{Weight: &fakeTensor{Name: "blk.0.a.weight"}},
+				B:  &nn.Linear{Weight: &fakeTensor{Name: "blk.0.b_weight"}, Bias: &fakeTensor{Name: "blk.0.b_bias"}},
+				C:  &nn.Linear{Weight: &fakeTensor{Name: "blk.0.weight_c"}},
+				XY: &nn.Linear{Weight: &fakeTensor{Name: "blk.0.x_weight_y"}},
+			},
+			{
+				A:  &nn.Linear{Weight: &fakeTensor{Name: "blk.1.a.weight"}},
+				B:  &nn.Linear{Weight: &fakeTensor{Name: "blk.1.b_weight"}, Bias: &fakeTensor{Name: "blk.1.b_bias"}},
+				C:  &nn.Linear{Weight: &fakeTensor{Name: "blk.1.weight_c"}},
+				XY: &nn.Linear{Weight: &fakeTensor{Name: "blk.1.x_weight_y"}},
+			},
+		},
 	}, m); diff != "" {
 		t.Errorf("populateFields() set incorrect values (-want +got):\n%s", diff)
 	}

model/models/deepseek2/model.go

@@ -0,0 +1,324 @@
+package deepseek2
+
+// uses deepseek 2 architecture but written based on deepseek 3 model
+
+import (
+	"math"
+
+	"github.com/ollama/ollama/fs"
+	"github.com/ollama/ollama/kvcache"
+	"github.com/ollama/ollama/ml"
+	"github.com/ollama/ollama/ml/nn"
+	"github.com/ollama/ollama/ml/nn/fast"
+	"github.com/ollama/ollama/ml/nn/rope"
+	"github.com/ollama/ollama/model"
+	"github.com/ollama/ollama/model/input"
+)
+
+type Options struct {
+	numExpertsUsed      int
+	numExperts          int
+	normTopKProb        bool
+	routedScalingFactor float32
+
+	kvLoraRank,
+	qkNopeHeadDim,
+	qkRopeHeadDim,
+	kqNopeHeadDim,
+	qkHeadDim int
+	qLoraRank int
+	vHeadDim  int
+
+	hiddenSize,
+	numHeads,
+	numKVHeads,
+	keyLength,
+	valueLength,
+	originalContextLength int
+
+	eps,
+	ropeBase,
+	ropeScale float32
+	kqScale float64
+}
+
+func (o Options) RoPEOptions() []func(*rope.Options) {
+	attnFactor := float32(1.0 / (1.0 + 0.1*math.Log(float64(o.ropeScale))))
+	return []func(*rope.Options){
+		rope.WithOriginalContextLength(o.originalContextLength),
+		rope.WithExtrapolationFactor(1.),
+		rope.WithAttentionFactor(attnFactor),
+	}
+}
+
+type Attention struct {
+	Q *nn.Linear `gguf:"attn_q"`
+
+	QA     *nn.Linear  `gguf:"attn_q_a"`
+	QANorm *nn.RMSNorm `gguf:"attn_q_a_norm"`
+	QB     *nn.Linear  `gguf:"attn_q_b"`
+
+	KVA     *nn.Linear  `gguf:"attn_kv_a_mqa"`
+	KVANorm *nn.RMSNorm `gguf:"attn_kv_a_norm"`
+	KVB     *nn.Linear  `gguf:"attn_kv_b"`
+
+	Output *nn.Linear `gguf:"attn_out,alt:attn_output"`
+}
+
+func (attn *Attention) Forward(ctx ml.Context, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
+	seqLength := hiddenStates.Dim(1)
+
+	var query ml.Tensor
+	if opts.qLoraRank == 0 { // nil {
+		query = attn.Q.Forward(ctx, hiddenStates)
+	} else {
+		query = attn.QA.Forward(ctx, hiddenStates)
+		query = attn.QANorm.Forward(ctx, query, opts.eps)
+		query = attn.QB.Forward(ctx, query)
+	}
+
+	query = query.Reshape(ctx, query.Dim(0)/opts.numHeads, opts.numHeads, seqLength)
+
+	qPass := query.View(ctx, 0,
+		opts.qkNopeHeadDim, query.Stride(1),
+		query.Dim(1), query.Stride(2),
+		query.Dim(2))
+
+	qRot := query.View(ctx, opts.qkNopeHeadDim*query.Stride(0),
+		opts.qkRopeHeadDim, query.Stride(1),
+		query.Dim(1), query.Stride(2),
+		query.Dim(2))
+
+	compressedKV := attn.KVA.Forward(ctx, hiddenStates)
+
+	kPass := compressedKV.View(ctx, 0, opts.kvLoraRank, compressedKV.Stride(1), compressedKV.Dim(1))
+	kRot := compressedKV.View(ctx, opts.kvLoraRank*compressedKV.Stride(0),
+		opts.qkRopeHeadDim, compressedKV.Stride(1),
+		1, compressedKV.Stride(1),
+		compressedKV.Dim(1))
+
+	kPass = attn.KVANorm.Forward(ctx, kPass, opts.eps)
+	kPass = attn.KVB.Forward(ctx, kPass)
+
+	kv := kPass.Reshape(ctx, kPass.Dim(0)/opts.numKVHeads, opts.numKVHeads, seqLength)
+	kPass = kv.View(ctx, 0, opts.kqNopeHeadDim, kv.Stride(1), kv.Dim(1), kv.Stride(2), kv.Dim(2))
+	value := kv.View(ctx, opts.kqNopeHeadDim*kv.Stride(0),
+		opts.vHeadDim, kv.Stride(1),
+		kv.Dim(1), kv.Stride(2),
+		kv.Dim(2)).Contiguous(ctx)
+
+	qRot = fast.RoPE(ctx, qRot, positions, opts.qkRopeHeadDim, opts.ropeBase, 1./opts.ropeScale, opts.RoPEOptions()...)
+	kRot = fast.RoPE(ctx, kRot, positions, opts.qkRopeHeadDim, opts.ropeBase, 1./opts.ropeScale, opts.RoPEOptions()...)
+
+	kRot = kRot.Repeat(ctx, 1, qPass.Dim(1))
+
+	query = qRot.Concat(ctx, qPass, 0)
+	key := kRot.Concat(ctx, kPass, 0)
+
+	attention := nn.Attention(ctx, query, key, value, opts.kqScale, cache)
+	attention = attention.Reshape(ctx, attention.Dim(0)*attention.Dim(1), seqLength)
+	return attn.Output.Forward(ctx, attention)
+}
+
+type MLP interface {
+	Forward(ml.Context, ml.Tensor, *Options) ml.Tensor
+}
+
+type sparse struct {
+	Router       *nn.Linear `gguf:"ffn_gate_inp"`
+	Gate         *nn.Linear `gguf:"ffn_gate_exps"`
+	Up           *nn.Linear `gguf:"ffn_up_exps"`
+	Down         *nn.Linear `gguf:"ffn_down_exps"`
+	SharedExpert *dense     `gguf:",suf:_shexp"`
+	ExpProbsBias ml.Tensor  `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
+}
+
+func (moe *sparse) Moe(ctx ml.Context, hiddenStates, topKIndices, topKWeights ml.Tensor, opts *Options) ml.Tensor {
+	hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
+
+	upStates := moe.Up.Weight.MulmatID(ctx, hiddenStates, topKIndices)
+	hiddenStates = moe.Gate.Weight.MulmatID(ctx, hiddenStates, topKIndices)
+	hiddenStates = hiddenStates.SILU(ctx, upStates)
+
+	experts := moe.Down.Weight.MulmatID(ctx, hiddenStates, topKIndices)
+	experts = experts.Mul(ctx, topKWeights)
+	nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
+	for i := 1; i < opts.numExpertsUsed; i++ {
+		nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
+	}
+	return nextStates
+}
+
+func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor {
+	scores = scores.Add(ctx, moe.ExpProbsBias)
+	topKIndices := scores.TopK(ctx, opts.numExpertsUsed)
+	return topKIndices
+}
+
+func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
+	residuals := hiddenStates
+
+	routerLogits := moe.Router.Forward(ctx, hiddenStates)
+	scores := routerLogits.Sigmoid(ctx)
+	topKIndices := moe.topKIndices(ctx, scores, opts)
+	topKWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, topKIndices)
+
+	if opts.normTopKProb {
+		topKWeights = topKWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1))
+		topKWeights = topKWeights.Div(ctx, topKWeights.SumRows(ctx))
+		topKWeights = topKWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1))
+	}
+
+	topKWeights = topKWeights.Scale(ctx, float64(opts.routedScalingFactor))
+	hiddenStates = moe.Moe(ctx, hiddenStates, topKIndices, topKWeights, opts)
+	sharedExpertResult := moe.SharedExpert.Forward(ctx, residuals, opts)
+
+	hiddenStates = hiddenStates.Add(ctx, sharedExpertResult)
+	return hiddenStates
+}
+
+type dense struct {
+	Gate *nn.Linear `gguf:"ffn_gate"`
+	Up   *nn.Linear `gguf:"ffn_up"`
+	Down *nn.Linear `gguf:"ffn_down"`
+}
+
+func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
+	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
+	return mlp.Down.Forward(ctx, hiddenStates)
+}
+
+type Layer struct {
+	AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
+	Attention     *Attention
+
+	MLPNorm *nn.RMSNorm `gguf:"ffn_norm"`
+	MLP     MLP
+}
+
+func (t *Layer) Forward(ctx ml.Context, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
+	residual := hiddenStates
+	hiddenStates = t.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
+	hiddenStates = t.Attention.Forward(ctx, hiddenStates, positions, cache, opts)
+
+	if outputs != nil {
+		hiddenStates = hiddenStates.Rows(ctx, outputs)
+		residual = residual.Rows(ctx, outputs)
+	}
+
+	hiddenStates = hiddenStates.Add(ctx, residual)
+	residual = hiddenStates
+
+	hiddenStates = t.MLPNorm.Forward(ctx, hiddenStates, opts.eps)
+	hiddenStates = t.MLP.Forward(ctx, hiddenStates, opts)
+	hiddenStates = hiddenStates.Add(ctx, residual)
+	return hiddenStates
+}
+
+type Model struct {
+	model.Base
+	model.BytePairEncoding
+
+	TokenEmbedding *nn.Embedding `gguf:"token_embd"`
+	Layers         []Layer       `gguf:"blk"`
+
+	OutputNorm *nn.RMSNorm `gguf:"output_norm"`
+	Output     *nn.Linear  `gguf:"output,alt:token_embd"`
+
+	*Options
+}
+
+func New(c fs.Config) (model.Model, error) {
+	layers := make([]Layer, c.Uint("block_count"))
+
+	firstDenseLayerIndex := int(c.Uint("leading_dense_block_count"))
+	for i := range layers {
+		if i < firstDenseLayerIndex {
+			layers[i].MLP = &dense{}
+		} else {
+			layers[i].MLP = &sparse{}
+		}
+	}
+
+	mScale := float32(1.0 + float64(c.Float("rope.scaling.yarn_log_multiplier"))*math.Log(float64(c.Float("rope.scaling.factor"))))
+	kqScale := float64(mScale) * float64(mScale) / math.Sqrt(float64(c.Uint("attention.key_length")))
+
+	m := Model{
+		BytePairEncoding: model.NewBytePairEncoding(
+			&model.Vocabulary{
+				Values: c.Strings("tokenizer.ggml.tokens"),
+				Types:  c.Ints("tokenizer.ggml.token_type"),
+				Merges: c.Strings("tokenizer.ggml.merges"),
+				AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
+				BOS:    []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
+				AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
+				EOS: append(
+					[]int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))},
+					c.Ints("tokenizer.ggml.eos_token_ids")...,
+				),
+			},
+			// Split regex into multiple parts (according to DeepSeek3's regex)
+			"\\p{N}{1,3}",
+			`[一-龥぀-ゟ゠-ヿ]+`,
+			"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+",
+		),
+		Layers: layers,
+		Options: &Options{
+			hiddenSize:     int(c.Uint("embedding_length")),
+			numHeads:       int(c.Uint("attention.head_count")),
+			numKVHeads:     int(c.Uint("attention.head_count_kv")),
+			keyLength:      int(c.Uint("attention.key_length")),
+			valueLength:    int(c.Uint("attention.value_length")),
+			eps:            c.Float("attention.layer_norm_rms_epsilon"),
+			ropeBase:       c.Float("rope.freq_base"),
+			ropeScale:      c.Float("rope.scaling.factor", 1),
+			numExperts:     int(c.Uint("expert_count")),
+			numExpertsUsed: int(c.Uint("expert_used_count")),
+			normTopKProb:   c.Bool("expert_weights_norm", true),
+
+			qLoraRank:     int(c.Uint("attention.q_lora_rank")), //&qLoraRankVal,
+			kvLoraRank:    int(c.Uint("attention.kv_lora_rank")),
+			qkHeadDim:     int(c.Uint("attention.key_length")),
+			vHeadDim:      int(c.Uint("attention.value_length")),
+			qkRopeHeadDim: int(c.Uint("rope.dimension_count")),
+			qkNopeHeadDim: int(c.Uint("attention.key_length")) - int(c.Uint("rope.dimension_count")),
+			kqNopeHeadDim: int(c.Uint("attention.key_length")) - int(c.Uint("rope.dimension_count")),
+
+			routedScalingFactor:   c.Float("expert_weights_scale"),
+			originalContextLength: int(c.Uint("rope.scaling.original_context_length")),
+
+			kqScale: kqScale,
+		},
+	}
+
+	m.Cache = kvcache.NewCausalCache(m.Shift)
+	return &m, nil
+}
+
+func (m Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
+	return fast.RoPE(ctx, key, shift, m.qkRopeHeadDim, m.ropeBase, 1./m.ropeScale, m.RoPEOptions()...), nil
+}
+
+func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
+	positions := ctx.Input().FromIntSlice(batch.Positions, len(batch.Positions))
+
+	hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
+
+	for i, layer := range m.Layers {
+		m.Cache.SetLayer(i)
+
+		var outputs ml.Tensor
+		if i == len(m.Layers)-1 {
+			outputs = batch.Outputs
+		}
+
+		hiddenStates = layer.Forward(ctx, hiddenStates, positions, outputs, m.Cache, m.Options)
+	}
+
+	hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
+	return m.Output.Forward(ctx, hiddenStates), nil
+}
+
+func init() {
+	model.Register("deepseek2", New)
+}

model/models/gptoss/model.go

@@ -227,17 +227,6 @@ func New(c fs.Config) (model.Model, error) {
 	m := Transformer{
 		TransformerBlocks: make([]TransformerBlock, c.Uint("block_count")),
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer",
-				strings.Join([]string{
-					`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?`,
-					`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?`,
-					`\p{N}{1,3}`,
-					` ?[^\s\p{L}\p{N}]+[\r\n/]*`,
-					`\s*[\r\n]+`,
-					`\s+(?!\S)`,
-					`\s+`,
-				}, "|"),
-			),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -250,6 +239,15 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			strings.Join([]string{
+				`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?`,
+				`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?`,
+				`\p{N}{1,3}`,
+				` ?[^\s\p{L}\p{N}]+[\r\n/]*`,
+				`\s*[\r\n]+`,
+				`\s+(?!\S)`,
+				`\s+`,
+			}, "|"),
 		),
 		Options: Options{
 			hiddenSize:            int(c.Uint("embedding_length")),

model/models/llama/model.go

@@ -54,10 +54,30 @@ func New(c fs.Config) (model.Model, error) {
 	}
 	switch c.String("tokenizer.ggml.model") {
 	case "gpt2":
-		processor = model.NewBytePairEncoding(
-			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
-			&vocabulary,
-		)
+		var pretokenizers []string
+		switch c.String("tokenizer.ggml.pre") {
+		case "default":
+			// no-op use the default bpe pretokenizer
+		case "qwen2":
+			pretokenizers = []string{
+				"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+			}
+		case "refact":
+			pretokenizers = []string{
+				`\p{N}`,
+				`'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`,
+			}
+		case "tekken":
+			pretokenizers = []string{
+				"[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+			}
+		default:
+			// use a llama-style pretokenizer
+			pretokenizers = []string{
+				"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+			}
+		}
+		processor = model.NewBytePairEncoding(&vocabulary, pretokenizers...)
 	case "llama":
 		processor = model.NewSentencePiece(&vocabulary)
 	default:

model/models/llama4/model.go

@@ -34,8 +34,6 @@ func (p *Projector) Forward(ctx ml.Context, visionOutputs ml.Tensor) ml.Tensor {
 func New(c fs.Config) (model.Model, error) {
 	m := Model{
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer",
-				`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -48,6 +46,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		ImageProcessor: newImageProcessor(c),
 		VisionModel:    newVisionModel(c),

model/models/llama4/model_text.go

@@ -88,22 +88,10 @@ func (e *TextExperts) Forward(ctx ml.Context, hiddenStates, routerLogits ml.Tens
 	return nextStates
 }
 
-// TextSharedExpert is TextMLP with different tensor names
-type TextSharedExpert struct {
-	Gate *nn.Linear `gguf:"ffn_gate_shexp"`
-	Up   *nn.Linear `gguf:"ffn_up_shexp"`
-	Down *nn.Linear `gguf:"ffn_down_shexp"`
-}
-
-func (mlp *TextSharedExpert) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
-	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
-	return mlp.Down.Forward(ctx, hiddenStates)
-}
-
 type TextMOE struct {
 	Router       *nn.Linear `gguf:"ffn_gate_inp"`
 	Experts      *TextExperts
-	SharedExpert *TextSharedExpert
+	SharedExpert *TextMLP `gguf:",suf:_shexp"`
 }
 
 func (moe *TextMOE) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {

model/models/mistral3/model.go

@@ -33,7 +33,6 @@ var _ model.TextProcessor = (*Model)(nil)
 func New(c fs.Config) (model.Model, error) {
 	m := &Model{
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer", `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -46,6 +45,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		TextModel:           newTextModel(c),
 		VisionModel:         newVisionModel(c),

model/models/mllama/model.go

@@ -33,7 +33,6 @@ const (
 func New(c fs.Config) (model.Model, error) {
 	m := Model{
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -46,6 +45,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		ImageProcessor: newImageProcessor(c),
 		VisionModel:    newVisionModel(c),

model/models/models.go

@@ -2,6 +2,7 @@ package models
 
 import (
 	_ "github.com/ollama/ollama/model/models/bert"
+	_ "github.com/ollama/ollama/model/models/deepseek2"
 	_ "github.com/ollama/ollama/model/models/gemma2"
 	_ "github.com/ollama/ollama/model/models/gemma3"
 	_ "github.com/ollama/ollama/model/models/gemma3n"

model/models/qwen2/model.go

@@ -139,7 +139,6 @@ func New(c fs.Config) (model.Model, error) {
 	m := Model{
 		Layers: make([]DecoderLayer, c.Uint("block_count")),
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -152,6 +151,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		Options: Options{
 			hiddenSize: int(c.Uint("embedding_length")),

model/models/qwen25vl/model.go

@@ -29,7 +29,6 @@ var _ model.MultimodalProcessor = (*Model)(nil)
 func New(c fs.Config) (model.Model, error) {
 	m := &Model{
 		BytePairEncoding: model.NewBytePairEncoding(
-			c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -42,6 +41,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		TextModel:      NewTextModel(c),
 		VisionModel:    newVisionModel(c),

model/models/qwen3/embed.go

@@ -35,7 +35,6 @@ func newEmbed(c fs.Config) (model.Model, error) {
 	}
 	m := embedModel{
 		BytePairEncoding: model.NewBytePairEncoding(
-			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -48,6 +47,7 @@ func newEmbed(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		Model: &Model{
 			Layers: layers,

model/models/qwen3/model.go

@@ -200,7 +200,6 @@ func New(c fs.Config) (model.Model, error) {
 
 	m := Model{
 		BytePairEncoding: model.NewBytePairEncoding(
-			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 			&model.Vocabulary{
 				Values: c.Strings("tokenizer.ggml.tokens"),
 				Types:  c.Ints("tokenizer.ggml.token_type"),
@@ -213,6 +212,7 @@ func New(c fs.Config) (model.Model, error) {
 					c.Ints("tokenizer.ggml.eos_token_ids")...,
 				),
 			},
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
 		),
 		Layers: layers,
 		Options: &Options{

model/parsers/parsers.go

@@ -2,10 +2,16 @@ package parsers
 
 import (
 	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/harmony"
 )
 
 type Parser interface {
-	Add(s string, tools []api.Tool) (content string, thinking string, calls []api.ToolCall, err error)
+	// Init initializes the parser with tools and optional last message for chat prefill
+	// Returns processed tools if the parser needs to modify them (e.g., harmony renames them)
+	Init(tools []api.Tool, lastMessage *api.Message) []api.Tool
+	// Add processes streamed content and returns parsed content, thinking, and tool calls
+	// The done flag indicates if this is the last chunk (used for draining accumulators)
+	Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error)
 	HasToolSupport() bool
 	HasThinkingSupport() bool
 }
@@ -17,6 +23,8 @@ func ParserForName(name string) Parser {
 		return parser
 	case "passthrough":
 		return &PassthroughParser{}
+	case "harmony":
+		return harmony.NewHarmonyMessageHandler()
 	default:
 		return nil
 	}
@@ -24,7 +32,11 @@ func ParserForName(name string) Parser {
 
 type PassthroughParser struct{}
 
-func (p *PassthroughParser) Add(s string, tools []api.Tool) (content string, thinking string, calls []api.ToolCall, err error) {
+func (p *PassthroughParser) Init(tools []api.Tool, lastMessage *api.Message) []api.Tool {
+	return tools // passthrough doesn't modify tools
+}
+
+func (p *PassthroughParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
 	return s, "", nil, nil
 }
 

model/parsers/qwen3coder.go

@@ -11,6 +11,7 @@ import (
 	"strconv"
 	"strings"
 	"unicode"
+	"unicode/utf8"
 
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/logutil"
@@ -31,6 +32,7 @@ const (
 type Qwen3CoderParser struct {
 	state qwenParserState
 	acc   strings.Builder
+	tools []api.Tool
 }
 
 func (p *Qwen3CoderParser) HasToolSupport() bool {
@@ -41,7 +43,12 @@ func (p *Qwen3CoderParser) HasThinkingSupport() bool {
 	return false
 }
 
-func (p *Qwen3CoderParser) Add(s string, tools []api.Tool) (content string, thinking string, calls []api.ToolCall, err error) {
+func (p *Qwen3CoderParser) Init(tools []api.Tool, lastMessage *api.Message) []api.Tool {
+	p.tools = tools
+	return tools // Qwen doesn't modify tools
+}
+
+func (p *Qwen3CoderParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
 	p.acc.WriteString(s)
 
 	events := p.parseEvents()
@@ -51,7 +58,7 @@ func (p *Qwen3CoderParser) Add(s string, tools []api.Tool) (content string, thin
 	for _, event := range events {
 		switch event := event.(type) {
 		case qwenEventRawToolCall:
-			toolCall, err := parseToolCall(event, tools)
+			toolCall, err := parseToolCall(event, p.tools)
 			if err != nil {
 				slog.Warn("qwen tool call parsing failed", "error", err)
 				return "", "", nil, err
@@ -198,12 +205,21 @@ func overlap(s, delim string) int {
 }
 
 func trailingWhitespaceLen(s string) int {
-	for i := len(s) - 1; i >= 0; i-- {
-		if !unicode.IsSpace(rune(s[i])) {
-			return len(s) - i - 1
+	remaining := s
+	total := 0
+	for len(remaining) > 0 {
+		r, size := utf8.DecodeLastRuneInString(remaining)
+		// if it's an invalid utf8 rune, assume it isn't whitespace
+		if r == utf8.RuneError && size == 1 {
+			break
 		}
+		if !unicode.IsSpace(r) {
+			break
+		}
+		total += size
+		remaining = remaining[:len(remaining)-size]
 	}
-	return len(s)
+	return total
 }
 
 type XMLFunctionCall struct {
@@ -359,7 +375,7 @@ func parseValue(raw string, paramType api.PropertyType) any {
 
 	// Try array
 	if typeSet["array"] {
-		var arr []interface{}
+		var arr []any
 		if err := json.Unmarshal([]byte(raw), &arr); err == nil {
 			return arr
 		}
@@ -371,7 +387,7 @@ func parseValue(raw string, paramType api.PropertyType) any {
 
 	// Try object
 	if typeSet["object"] {
-		var obj map[string]interface{}
+		var obj map[string]any
 		if err := json.Unmarshal([]byte(raw), &obj); err == nil {
 			return obj
 		}

model/parsers/qwen3coder_test.go

@@ -166,6 +166,137 @@ func TestQwenParserStreaming(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "unicode content",
+			steps: []step{
+				{
+					input: "你好 🌍<tool_call>test</tool_call>مرحبا",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "你好 🌍"},
+						qwenEventRawToolCall{raw: "test"},
+						qwenEventContent{content: "مرحبا"},
+					},
+				},
+			},
+		},
+		{
+			desc: "arabic text handling",
+			steps: []step{
+				{
+					input:      "مرحبا بالعالم",
+					wantEvents: []qwenEvent{qwenEventContent{content: "مرحبا بالعالم"}},
+				},
+			},
+		},
+		{
+			desc: "emoji passthrough",
+			steps: []step{
+				{
+					input:      "✅",
+					wantEvents: []qwenEvent{qwenEventContent{content: "✅"}},
+				},
+			},
+		},
+		{
+			desc: "emoji after tool call",
+			steps: []step{
+				{
+					input: "<tool_call>test</tool_call>完成 ✅",
+					wantEvents: []qwenEvent{
+						qwenEventRawToolCall{raw: "test"},
+						qwenEventContent{content: "完成 ✅"},
+					},
+				},
+			},
+		},
+		{
+			desc: "unicode streaming with whitespace handling",
+			steps: []step{
+				{
+					input: "مرحبا",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "مرحبا"},
+					},
+				},
+				{
+					input:      " \n",
+					wantEvents: []qwenEvent{},
+				},
+				{
+					input: "世界",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: " \n世界"},
+					},
+				},
+			},
+		},
+		{
+			desc: "non-breaking space withheld across chunks",
+			steps: []step{
+				{
+					input: "Hello\u00a0",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "Hello"},
+					},
+				},
+				{
+					input: "world",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "\u00a0world"},
+					},
+				},
+			},
+		},
+		{
+			desc: "ideographic space before partial tool",
+			steps: []step{
+				{
+					input: "Hello\u3000<tool",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "Hello"},
+					},
+				},
+				{
+					input:      "_call>abc",
+					wantEvents: []qwenEvent{},
+				},
+				{
+					input: "</tool_call>def",
+					wantEvents: []qwenEvent{
+						qwenEventRawToolCall{raw: "abc"},
+						qwenEventContent{content: "def"},
+					},
+				},
+			},
+		},
+		{
+			desc: "ideographic space before partial tool fakeout",
+			steps: []step{
+				{
+					input: "Hello\u3000<tool",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "Hello"},
+					},
+				},
+				{
+					input: "fakeout>abc",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "\u3000<toolfakeout>abc"},
+					},
+				},
+			},
+		},
+		{
+			desc: "unicode with partial tool tag",
+			steps: []step{
+				{
+					input: "测试🎯 <to",
+					wantEvents: []qwenEvent{
+						qwenEventContent{content: "测试🎯"},
+					},
+				},
+			},
+		},
 	}
 
 	anyOnlies := false
@@ -347,6 +478,27 @@ ls && echo "a > b and a < b"
 				},
 			},
 		},
+		{
+			name:  "unicode in function names and parameters",
+			tools: []api.Tool{},
+			rawToolCall: `<function=获取天气>
+<parameter=城市>
+北京
+</parameter>
+<parameter=message>
+Hello! 你好! 🌟 مرحبا
+</parameter>
+</function>`,
+			wantToolCall: api.ToolCall{
+				Function: api.ToolCallFunction{
+					Name: "获取天气",
+					Arguments: map[string]any{
+						"城市":      "北京",
+						"message": "Hello! 你好! 🌟 مرحبا",
+					},
+				},
+			},
+		},
 	}
 
 	for i, step := range steps {
@@ -360,6 +512,42 @@ ls && echo "a > b and a < b"
 	}
 }
 
+func TestTrailingWhitespaceLenUnicode(t *testing.T) {
+	cases := []struct {
+		name  string
+		input string
+		want  int
+	}{
+		{
+			name:  "ascii space",
+			input: "Hello ",
+			want:  1,
+		},
+		{
+			name:  "non-breaking space",
+			input: "Hello\u00a0",
+			want:  2,
+		},
+		{
+			name:  "ideographic space",
+			input: "Hello\u3000",
+			want:  3,
+		},
+		{
+			name:  "multiple runes of whitespace",
+			input: "Hi\u00a0\u3000",
+			want:  5,
+		},
+	}
+
+	for _, tc := range cases {
+		got := trailingWhitespaceLen(tc.input)
+		if got != tc.want {
+			t.Errorf("%s: trailingWhitespaceLen(%q) = %d, want %d", tc.name, tc.input, got, tc.want)
+		}
+	}
+}
+
 func TestQwenToolCallValueParsing(t *testing.T) {
 	cases := []struct {
 		desc      string
@@ -867,6 +1055,8 @@ func TestTrailingWhitespaceLen(t *testing.T) {
 		{desc: "trailing whitespace with newlines", s: "abc \n", want: 2},
 		{desc: "only whitespace", s: " \n  ", want: 4},
 		{desc: "leading whitespace doesn't count", s: " \n abc", want: 0},
+		{desc: "unicode with trailing space", s: "测试🎯 ", want: 1},
+		{desc: "unicode with trailing tab and newline", s: "مرحبا\t\n", want: 2},
 	}
 
 	for _, tc := range cases {
@@ -876,3 +1066,30 @@ func TestTrailingWhitespaceLen(t *testing.T) {
 		}
 	}
 }
+
+func TestOverlapFunction(t *testing.T) {
+	cases := []struct {
+		desc  string
+		s     string
+		delim string
+		want  int
+	}{
+		{desc: "no overlap", s: "hello", delim: "<tool", want: 0},
+		{desc: "full overlap", s: "hello<tool", delim: "<tool>", want: 5},
+		{desc: "partial overlap", s: "hello<to", delim: "<tool>", want: 3},
+		{desc: "unicode with partial overlap", s: "测试🎯<to", delim: "<tool>", want: 3},
+		{desc: "unicode string with no overlap", s: "مرحبا", delim: "<tool>", want: 0},
+		{desc: "unicode at boundary", s: "世界<", delim: "<tool>", want: 1},
+		{desc: "unicode delimiter single rune", s: "hello🔧", delim: "🔧工具", want: len("🔧")},
+		{desc: "unicode delimiter multiple runes", s: "hello🔧工", delim: "🔧工具", want: len("🔧工")},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.desc, func(t *testing.T) {
+			got := overlap(tc.s, tc.delim)
+			if got != tc.want {
+				t.Errorf("overlap(%q, %q) = %d, want %d", tc.s, tc.delim, got, tc.want)
+			}
+		})
+	}
+}

model/renderers/qwen3coder.go

@@ -99,9 +99,7 @@ func Qwen3CoderRenderer(messages []api.Message, tools []api.Tool, _ *api.ThinkVa
 					sb.WriteString("\n<name>" + name + "</name>")
 
 					if len(prop.Type) > 0 {
-						// TODO(!!!)(drifkin): we should match the reference implementation for
-						// more complex types here instead of using this format
-						sb.WriteString("\n<type>" + prop.ToTypeScriptType() + "</type>")
+						sb.WriteString("\n<type>" + formatToolDefinitionType(prop.Type) + "</type>")
 					}
 
 					if prop.Description != "" {
@@ -215,3 +213,24 @@ func formatToolCallArgument(value any) string {
 
 	return fmt.Sprintf("%v", value)
 }
+
+func formatToolDefinitionType(tp api.PropertyType) string {
+	if len(tp) == 0 {
+		return "[]"
+	}
+
+	if len(tp) == 1 {
+		return tp[0]
+	}
+
+	// TODO(drifkin): it would be nice to format the JSON here similarly to
+	// python's default json.dumps behavior (spaces after commas and colons).
+	// This would let us be byte-for-byte compatible with the reference
+	// implementation for most common inputs
+	jsonBytes, err := json.Marshal(tp)
+	if err != nil {
+		return "[]"
+	}
+
+	return string(jsonBytes)
+}

model/renderers/qwen3coder_test.go

@@ -336,3 +336,35 @@ func TestFormatToolCallArgument(t *testing.T) {
 		})
 	}
 }
+
+func TestQwen3ToolDefinitionTypes(t *testing.T) {
+	tests := []struct {
+		name         string
+		propertyType api.PropertyType
+		expected     string
+	}{
+		{
+			name:         "simple",
+			propertyType: api.PropertyType{"string"},
+			expected:     "string",
+		},
+		{
+			name:         "multiple",
+			propertyType: api.PropertyType{"string", "number"},
+			expected:     "[\"string\",\"number\"]",
+		},
+		{
+			name:         "empty",
+			propertyType: api.PropertyType{},
+			expected:     "[]",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := formatToolDefinitionType(tt.propertyType)
+			if got != tt.expected {
+				t.Errorf("formatToolDefinitionType() = %v, want %v", got, tt.expected)
+			}
+		})
+	}
+}

sample/samplers_test.go

@@ -82,7 +82,6 @@ func modelHelper(t testing.TB) model.BytePairEncoding {
 	merges := make([]string, 0, 1)
 	// Only need vocab for Grammar Test
 	return model.NewBytePairEncoding(
-		``,
 		&model.Vocabulary{
 			Values: tokens,
 			Types:  make([]int32, len(vocab)),

server/routes.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"cmp"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -34,7 +35,6 @@ import (
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/format"
 	"github.com/ollama/ollama/fs/ggml"
-	"github.com/ollama/ollama/harmony"
 	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/logutil"
 	"github.com/ollama/ollama/model/parsers"
@@ -49,6 +49,8 @@ import (
 	"github.com/ollama/ollama/version"
 )
 
+const signinURLStr = "https://ollama.com/connect?name=%s&key=%s"
+
 func shouldUseHarmony(model *Model) bool {
 	if slices.Contains([]string{"gptoss", "gpt-oss"}, model.Config.ModelFamily) {
 		// heuristic to check whether the template expects to be parsed via harmony:
@@ -74,9 +76,20 @@ var lowVRAMThreshold uint64 = 20 * format.GibiByte
 var mode string = gin.DebugMode
 
 type Server struct {
-	addr    net.Addr
-	sched   *Scheduler
-	lowVRAM bool
+	addr         net.Addr
+	sched        *Scheduler
+	lowVRAM      bool
+	cloudBaseURL *url.URL
+}
+
+func (s *Server) webServiceBase() *url.URL {
+	defaultWebServiceURL := url.URL{Scheme: "https", Host: "ollama.com"}
+	if s != nil && s.cloudBaseURL != nil {
+		u := *s.cloudBaseURL
+		return &u
+	}
+	u := defaultWebServiceURL
+	return &u
 }
 
 func init() {
@@ -151,6 +164,17 @@ func (s *Server) scheduleRunner(ctx context.Context, name string, caps []model.C
 	return runner.llama, model, &opts, nil
 }
 
+func signinURL() (string, error) {
+	pubKey, err := auth.GetPublicKey()
+	if err != nil {
+		return "", err
+	}
+
+	encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
+	h, _ := os.Hostname()
+	return fmt.Sprintf(signinURLStr, url.PathEscape(h), encKey), nil
+}
+
 func (s *Server) GenerateHandler(c *gin.Context) {
 	checkpointStart := time.Now()
 	var req api.GenerateRequest
@@ -251,18 +275,21 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		client := api.NewClient(remoteURL, http.DefaultClient)
 		err = client.Generate(c, &req, fn)
 		if err != nil {
-			var sErr api.AuthorizationError
-			if errors.As(err, &sErr) && sErr.StatusCode == http.StatusUnauthorized {
-				pk, pkErr := auth.GetPublicKey()
-				if pkErr != nil {
-					slog.Error("couldn't get public key", "error", pkErr)
-					c.JSON(http.StatusUnauthorized, gin.H{"error": "error getting public key"})
+			var authError api.AuthorizationError
+			if errors.As(err, &authError) {
+				sURL, sErr := signinURL()
+				if sErr != nil {
+					slog.Error(sErr.Error())
+					c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
 					return
 				}
-				c.JSON(http.StatusUnauthorized, gin.H{
-					"error":      "unauthorized",
-					"public_key": pk,
-				})
+
+				c.JSON(authError.StatusCode, gin.H{"error": "unauthorized", "signin_url": sURL})
+				return
+			}
+			var apiError api.StatusError
+			if errors.As(err, &apiError) {
+				c.JSON(apiError.StatusCode, apiError)
 				return
 			}
 			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
@@ -291,17 +318,21 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		return
 	}
 
-	useHarmony := shouldUseHarmony(m) && !req.Raw
-	var harmonyMessageHandler *harmony.HarmonyMessageHandler
-	var harmonyToolParser *harmony.HarmonyToolCallAccumulator
-	if useHarmony {
-		harmonyMessageHandler = harmony.NewHarmonyMessageHandler()
-		harmonyMessageHandler.HarmonyParser.AddImplicitStart()
-		harmonyToolParser = harmonyMessageHandler.CreateToolParser()
+	var builtinParser parsers.Parser
+	if shouldUseHarmony(m) && m.Config.Parser == "" {
+		m.Config.Parser = "harmony"
 	}
 
-	// Validate Think value: string values currently only allowed for gptoss models
-	if req.Think != nil && req.Think.IsString() && !useHarmony {
+	if !req.Raw && m.Config.Parser != "" {
+		builtinParser = parsers.ParserForName(m.Config.Parser)
+		if builtinParser != nil {
+			// no tools or last message for generate endpoint
+			builtinParser.Init(nil, nil)
+		}
+	}
+
+	// Validate Think value: string values currently only allowed for harmony/gptoss models
+	if req.Think != nil && req.Think.IsString() && m.Config.Parser != "harmony" {
 		c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("think value %q is not supported for this model", req.Think.String())})
 		return
 	}
@@ -425,7 +456,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 	}
 
 	var thinkingState *thinking.Parser
-	if !useHarmony {
+	if builtinParser == nil {
 		openingTag, closingTag := thinking.InferTags(m.Template.Template)
 		if req.Think != nil && req.Think.Bool() && openingTag != "" && closingTag != "" {
 			thinkingState = &thinking.Parser{
@@ -462,11 +493,17 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 				},
 			}
 
-			if useHarmony {
-				content, thinking, toolContent := harmonyMessageHandler.AddContent(cr.Content, harmonyToolParser)
+			if builtinParser != nil {
+				content, thinking, toolCalls, err := builtinParser.Add(cr.Content, cr.Done)
+				if err != nil {
+					ch <- gin.H{"error": err.Error()}
+					return
+				}
 				res.Response = content
 				res.Thinking = thinking
-				harmonyToolParser.Add(toolContent)
+				if cr.Done && len(toolCalls) > 0 {
+					res.ToolCalls = toolCalls
+				}
 			} else if thinkingState != nil {
 				thinking, content := thinkingState.AddContent(cr.Content)
 				res.Thinking = thinking
@@ -478,26 +515,6 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 			}
 
 			if cr.Done {
-				if useHarmony {
-					toolName, toolContent := harmonyToolParser.Drain()
-					if toolName != nil {
-						*toolName = strings.TrimPrefix(*toolName, "functions.")
-						var args api.ToolCallFunctionArguments
-						if err := json.Unmarshal([]byte(toolContent), &args); err != nil {
-							errStr := fmt.Sprintf("error parsing tool call: raw='%s', err=%s", toolContent, err.Error())
-							ch <- gin.H{"error": errStr}
-							return
-						}
-
-						res.ToolCalls = append(res.ToolCalls, api.ToolCall{
-							Function: api.ToolCallFunction{
-								Name:      *toolName,
-								Arguments: args,
-							},
-						})
-					}
-				}
-
 				res.DoneReason = cr.DoneReason.String()
 				res.TotalDuration = time.Since(checkpointStart)
 				res.LoadDuration = checkpointLoaded.Sub(checkpointStart)
@@ -512,7 +529,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 				}
 			}
 
-			if useHarmony {
+			if builtinParser != nil {
 				// only send messages with meaningful content (empty messages confuse clients)
 				if res.Response != "" || res.Thinking != "" || res.Done || len(res.ToolCalls) > 0 {
 					ch <- res
@@ -761,6 +778,105 @@ func (s *Server) EmbeddingsHandler(c *gin.Context) {
 	c.JSON(http.StatusOK, resp)
 }
 
+func (s *Server) WebSearchHandler(c *gin.Context) {
+	var req api.WebSearchRequest
+	if err := c.ShouldBindJSON(&req); errors.Is(err, io.EOF) {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
+		return
+	} else if err != nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	req.Query = strings.TrimSpace(req.Query)
+	if req.Query == "" {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "query is required"})
+		return
+	}
+
+	if req.MaxResults != 0 && (req.MaxResults < 1 || req.MaxResults > 10) {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "max_results must be between 1 and 10"})
+		return
+	}
+
+	webServiceClient := api.NewClient(s.webServiceBase(), http.DefaultClient)
+	resp, err := webServiceClient.WebSearch(c.Request.Context(), &req)
+	if err != nil {
+		var authError api.AuthorizationError
+		if errors.As(err, &authError) {
+			sURL, sErr := signinURL()
+			if sErr != nil {
+				slog.Error(sErr.Error())
+				c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
+				return
+			}
+
+			c.JSON(authError.StatusCode, gin.H{"error": "unauthorized", "signin_url": sURL})
+			return
+		}
+		var apiError api.StatusError
+		if errors.As(err, &apiError) {
+			c.JSON(apiError.StatusCode, apiError)
+			return
+		}
+		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
+		return
+	}
+
+	if resp == nil {
+		resp = &api.WebSearchResponse{}
+	}
+
+	c.JSON(http.StatusOK, resp)
+}
+
+func (s *Server) WebFetchHandler(c *gin.Context) {
+	var req api.WebFetchRequest
+	if err := c.ShouldBindJSON(&req); errors.Is(err, io.EOF) {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
+		return
+	} else if err != nil {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	req.URL = strings.TrimSpace(req.URL)
+	if req.URL == "" {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "url is required"})
+		return
+	}
+
+	webServiceClient := api.NewClient(s.webServiceBase(), http.DefaultClient)
+	resp, err := webServiceClient.WebFetch(c.Request.Context(), &req)
+	if err != nil {
+		var authError api.AuthorizationError
+		if errors.As(err, &authError) {
+			sURL, sErr := signinURL()
+			if sErr != nil {
+				slog.Error(sErr.Error())
+				c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
+				return
+			}
+
+			c.JSON(authError.StatusCode, gin.H{"error": "unauthorized", "signin_url": sURL})
+			return
+		}
+		var apiError api.StatusError
+		if errors.As(err, &apiError) {
+			c.JSON(apiError.StatusCode, apiError)
+			return
+		}
+		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
+		return
+	}
+
+	if resp == nil {
+		resp = &api.WebFetchResponse{}
+	}
+
+	c.JSON(http.StatusOK, resp)
+}
+
 func (s *Server) PullHandler(c *gin.Context) {
 	var req api.PullRequest
 	err := c.ShouldBindJSON(&req)
@@ -1423,9 +1539,12 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) {
 	r.POST("/api/show", s.ShowHandler)
 	r.DELETE("/api/delete", s.DeleteHandler)
 
-	r.DELETE("/api/user/keys/:encodedKey", s.SignoutHandler)
 	r.POST("/api/me", s.WhoamiHandler)
 
+	r.POST("/api/signout", s.SignoutHandler)
+	// deprecated
+	r.DELETE("/api/user/keys/:encodedKey", s.SignoutHandler)
+
 	// Create
 	r.POST("/api/create", s.CreateHandler)
 	r.POST("/api/blobs/:digest", s.CreateBlobHandler)
@@ -1438,6 +1557,8 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) {
 	r.POST("/api/chat", s.ChatHandler)
 	r.POST("/api/embed", s.EmbedHandler)
 	r.POST("/api/embeddings", s.EmbeddingsHandler)
+	r.POST("/api/web_search", s.WebSearchHandler)
+	r.POST("/api/web_fetch", s.WebFetchHandler)
 
 	// Inference (OpenAI compatibility)
 	r.POST("/v1/chat/completions", openai.ChatMiddleware(), s.ChatHandler)
@@ -1636,11 +1757,32 @@ func (s *Server) WhoamiHandler(c *gin.Context) {
 	if err != nil {
 		slog.Error(err.Error())
 	}
+
+	// user isn't signed in
+	if user != nil && user.Name == "" {
+		sURL, sErr := signinURL()
+		if sErr != nil {
+			slog.Error(sErr.Error())
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
+			return
+		}
+
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized", "signin_url": sURL})
+		return
+	}
+
 	c.JSON(http.StatusOK, user)
 }
 
 func (s *Server) SignoutHandler(c *gin.Context) {
-	encodedKey := c.Param("encodedKey")
+	pubKey, err := auth.GetPublicKey()
+	if err != nil {
+		slog.Error("couldn't get public key", "error", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "there was an error signing out"})
+		return
+	}
+
+	encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
 
 	// todo allow other hosts
 	u, err := url.Parse("https://ollama.com")
@@ -1651,11 +1793,11 @@ func (s *Server) SignoutHandler(c *gin.Context) {
 	}
 
 	client := api.NewClient(u, http.DefaultClient)
-	err = client.Signout(c, encodedKey)
+	err = client.Disconnect(c, encKey)
 	if err != nil {
-		slog.Error(err.Error())
-		if strings.Contains(err.Error(), "page not found") || strings.Contains(err.Error(), "invalid credentials") {
-			c.JSON(http.StatusNotFound, gin.H{"error": "you are not currently signed in"})
+		var authError api.AuthorizationError
+		if errors.As(err, &authError) {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "you are not currently signed in"})
 			return
 		}
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "there was an error signing out"})
@@ -1747,7 +1889,7 @@ func (s *Server) ChatHandler(c *gin.Context) {
 	}
 
 	// expire the runner
-	if len(req.Messages) == 0 && req.KeepAlive != nil && int(req.KeepAlive.Seconds()) == 0 {
+	if len(req.Messages) == 0 && req.KeepAlive != nil && req.KeepAlive.Duration == 0 {
 		s.sched.expireRunner(m)
 
 		c.JSON(http.StatusOK, api.ChatResponse{
@@ -1813,18 +1955,21 @@ func (s *Server) ChatHandler(c *gin.Context) {
 		client := api.NewClient(remoteURL, http.DefaultClient)
 		err = client.Chat(c, &req, fn)
 		if err != nil {
-			var sErr api.AuthorizationError
-			if errors.As(err, &sErr) && sErr.StatusCode == http.StatusUnauthorized {
-				pk, pkErr := auth.GetPublicKey()
-				if pkErr != nil {
-					slog.Error("couldn't get public key", "error", pkErr)
-					c.JSON(http.StatusUnauthorized, gin.H{"error": "error getting public key"})
+			var authError api.AuthorizationError
+			if errors.As(err, &authError) {
+				sURL, sErr := signinURL()
+				if sErr != nil {
+					slog.Error(sErr.Error())
+					c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
 					return
 				}
-				c.JSON(http.StatusUnauthorized, gin.H{
-					"error":      "unauthorized",
-					"public_key": pk,
-				})
+
+				c.JSON(authError.StatusCode, gin.H{"error": "unauthorized", "signin_url": sURL})
+				return
+			}
+			var apiError api.StatusError
+			if errors.As(err, &apiError) {
+				c.JSON(apiError.StatusCode, apiError)
 				return
 			}
 			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
@@ -1870,32 +2015,23 @@ func (s *Server) ChatHandler(c *gin.Context) {
 	}
 	msgs = filterThinkTags(msgs, m)
 
-	var builtinParser parsers.Parser
-	if m.Config.Parser != "" {
-		builtinParser = parsers.ParserForName(m.Config.Parser)
+	if shouldUseHarmony(m) && m.Config.Parser == "" {
+		m.Config.Parser = "harmony"
 	}
 
-	var harmonyMessageHandler *harmony.HarmonyMessageHandler
-	var harmonyToolParser *harmony.HarmonyToolCallAccumulator
-
-	useHarmony := shouldUseHarmony(m) || m.Config.Parser == "harmony"
-
+	var builtinParser parsers.Parser
 	processedTools := req.Tools
-	if useHarmony {
-		harmonyMessageHandler = harmony.NewHarmonyMessageHandler()
-		var lastMessage *api.Message
-		if len(msgs) > 0 {
-			lastMessage = &msgs[len(msgs)-1]
-		}
-		harmonyMessageHandler.HarmonyParser.AddImplicitStartOrPrefill(lastMessage)
-		harmonyToolParser = harmonyMessageHandler.CreateToolParser()
 
-		// make a copy of tools to pass to the chat prompt. Function names may be
-		// renamed to be valid Harmony function names.
-		processedTools = make([]api.Tool, len(req.Tools))
-		copy(processedTools, req.Tools)
-		for i, tool := range processedTools {
-			processedTools[i].Function.Name = harmonyMessageHandler.FunctionNameMap.ConvertAndAdd(tool.Function.Name)
+	if m.Config.Parser != "" {
+		builtinParser = parsers.ParserForName(m.Config.Parser)
+		if builtinParser != nil {
+			// Determine last message for chat prefill
+			var lastMessage *api.Message
+			if len(msgs) > 0 {
+				lastMessage = &msgs[len(msgs)-1]
+			}
+			// Initialize parser and get processed tools
+			processedTools = builtinParser.Init(req.Tools, lastMessage)
 		}
 	}
 
@@ -1919,8 +2055,8 @@ func (s *Server) ChatHandler(c *gin.Context) {
 		return
 	}
 
-	// Validate Think value: string values currently only allowed for gptoss models
-	if req.Think != nil && req.Think.IsString() && !useHarmony {
+	// Validate Think value: string values currently only allowed for harmony/gptoss models
+	if req.Think != nil && req.Think.IsString() && m.Config.Parser != "harmony" {
 		c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("think value %q is not supported for this model", req.Think.String())})
 		return
 	}
@@ -1939,7 +2075,7 @@ func (s *Server) ChatHandler(c *gin.Context) {
 	}
 
 	var toolParser *tools.Parser
-	if len(req.Tools) > 0 && !useHarmony {
+	if len(req.Tools) > 0 && (builtinParser == nil || !builtinParser.HasToolSupport()) {
 		toolParser = tools.NewParser(m.Template.Template, req.Tools)
 	}
 
@@ -1971,38 +2107,10 @@ func (s *Server) ChatHandler(c *gin.Context) {
 				res.LoadDuration = checkpointLoaded.Sub(checkpointStart)
 			}
 
-			// TODO(drifkin): fold this as much as possibleinto the generic m.Config.Parser logic
-			if useHarmony {
-				content, thinking, toolContent := harmonyMessageHandler.AddContent(r.Content, harmonyToolParser)
-				res.Message.Content = content
-				res.Message.Thinking = thinking
-				harmonyToolParser.Add(toolContent)
-
-				if r.Done {
-					toolName, toolContent := harmonyToolParser.Drain()
-					if toolName != nil {
-						*toolName = strings.TrimPrefix(*toolName, "functions.")
-						*toolName = harmonyMessageHandler.FunctionNameMap.OriginalFromConverted(*toolName)
-						var args api.ToolCallFunctionArguments
-						if err := json.Unmarshal([]byte(toolContent), &args); err != nil {
-							errStr := fmt.Sprintf("error parsing tool call: raw='%s', err=%s", toolContent, err.Error())
-							ch <- gin.H{"error": errStr}
-							return
-						}
-						res.Message.ToolCalls = []api.ToolCall{{Function: api.ToolCallFunction{Name: *toolName, Arguments: args}}}
-					}
-				}
-
-				// only send messages with meaningful content (empty messages confuse clients)
-				if res.Message.Content != "" || res.Message.Thinking != "" || len(res.Message.ToolCalls) > 0 || res.Done {
-					ch <- res
-				}
-
-				return
-			} else if builtinParser != nil {
+			if builtinParser != nil {
 				slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser input", "parser", m.Config.Parser, "content", r.Content)
 
-				content, thinking, toolCalls, err := builtinParser.Add(r.Content, req.Tools)
+				content, thinking, toolCalls, err := builtinParser.Add(r.Content, r.Done)
 				if err != nil {
 					ch <- gin.H{"error": err.Error()}
 					return

server/routes_test.go

@@ -13,6 +13,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -92,6 +93,9 @@ func (t *panicTransport) RoundTrip(r *http.Request) (*http.Response, error) {
 var panicOnRoundTrip = &http.Client{Transport: &panicTransport{}}
 
 func TestRoutes(t *testing.T) {
+	// Disable authentication for tests to avoid issues with missing private keys
+	t.Setenv("OLLAMA_AUTH", "false")
+
 	type testCase struct {
 		Name     string
 		Method   string
@@ -139,6 +143,11 @@ func TestRoutes(t *testing.T) {
 		}
 	}
 
+	var (
+		searchRequests []api.WebSearchRequest
+		fetchRequests  []api.WebFetchRequest
+	)
+
 	testCases := []testCase{
 		{
 			Name:   "Version Handler",
@@ -455,6 +464,69 @@ func TestRoutes(t *testing.T) {
 				}
 			},
 		},
+		{
+			Name:   "Web Search Handler",
+			Method: http.MethodPost,
+			Path:   "/api/web_search",
+			Setup: func(t *testing.T, req *http.Request) {
+				searchRequests = nil
+				payload := api.WebSearchRequest{Query: "cats", MaxResults: 2}
+				data, err := json.Marshal(payload)
+				if err != nil {
+					t.Fatalf("failed to marshal request: %v", err)
+				}
+				req.Body = io.NopCloser(bytes.NewReader(data))
+				req.Header.Set("Content-Type", "application/json")
+			},
+			Expected: func(t *testing.T, resp *http.Response) {
+				if resp.StatusCode != http.StatusOK {
+					t.Fatalf("expected status 200, got %d", resp.StatusCode)
+				}
+				var out api.WebSearchResponse
+				if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
+					t.Fatalf("failed to decode response: %v", err)
+				}
+				if len(out.Results) != 1 || out.Results[0].Title != "Result" {
+					t.Fatalf("unexpected response: %+v", out)
+				}
+				if len(searchRequests) != 1 {
+					t.Fatalf("expected 1 forwarded request, got %d", len(searchRequests))
+				}
+				if searchRequests[0].Query != "cats" || searchRequests[0].MaxResults != 2 {
+					t.Fatalf("unexpected forwarded request: %+v", searchRequests[0])
+				}
+			},
+		},
+		{
+			Name:   "Web Fetch Handler",
+			Method: http.MethodPost,
+			Path:   "/api/web_fetch",
+			Setup: func(t *testing.T, req *http.Request) {
+				fetchRequests = nil
+				payload := api.WebFetchRequest{URL: "https://example.com"}
+				data, err := json.Marshal(payload)
+				if err != nil {
+					t.Fatalf("failed to marshal request: %v", err)
+				}
+				req.Body = io.NopCloser(bytes.NewReader(data))
+				req.Header.Set("Content-Type", "application/json")
+			},
+			Expected: func(t *testing.T, resp *http.Response) {
+				if resp.StatusCode != http.StatusOK {
+					t.Fatalf("expected status 200, got %d", resp.StatusCode)
+				}
+				var out api.WebFetchResponse
+				if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
+					t.Fatalf("failed to decode response: %v", err)
+				}
+				if out.Title != "Example" || len(out.Links) != 1 {
+					t.Fatalf("unexpected response: %+v", out)
+				}
+				if len(fetchRequests) != 1 || fetchRequests[0].URL != "https://example.com" {
+					t.Fatalf("unexpected forwarded request: %+v", fetchRequests)
+				}
+			},
+		},
 		{
 			Name: "openai retrieve model handler",
 			Setup: func(t *testing.T, req *http.Request) {
@@ -513,7 +585,40 @@ func TestRoutes(t *testing.T) {
 		HTTPClient: panicOnRoundTrip,
 	}
 
-	s := &Server{}
+	remoteSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/api/web_search":
+			var req api.WebSearchRequest
+			if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+			searchRequests = append(searchRequests, req)
+			resp := api.WebSearchResponse{Results: []api.WebSearchResult{{Title: "Result", URL: "https://example.com", Content: "snippet"}}}
+			_ = json.NewEncoder(w).Encode(resp)
+		case "/api/web_fetch":
+			var req api.WebFetchRequest
+			if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+			fetchRequests = append(fetchRequests, req)
+			resp := api.WebFetchResponse{Title: "Example", Content: "content", Links: []string{"https://example.com"}}
+			_ = json.NewEncoder(w).Encode(resp)
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer remoteSrv.Close()
+
+	remoteURL, err := url.Parse(remoteSrv.URL)
+	if err != nil {
+		t.Fatalf("parse remote server URL: %v", err)
+	}
+
+	s := &Server{
+		cloudBaseURL: remoteURL,
+	}
 	router, err := s.GenerateRoutes(rc)
 	if err != nil {
 		t.Fatalf("failed to generate routes: %v", err)

tools/tools.go

@@ -273,9 +273,21 @@ func findArguments(buffer []byte) (map[string]any, int) {
 						if args, ok := obj["arguments"].(map[string]any); ok {
 							return args, true
 						}
+						if argsStr, ok := obj["arguments"].(string); ok {
+							var argsData map[string]interface{}
+							if err := json.Unmarshal([]byte(argsStr), &argsData); err == nil {
+								return argsData, ok
+							}
+						}
 						if args, ok := obj["parameters"].(map[string]any); ok {
 							return args, true
 						}
+						if argsStr, ok := obj["parameters"].(string); ok {
+							var argsData map[string]interface{}
+							if err := json.Unmarshal([]byte(argsStr), &argsData); err == nil {
+								return argsData, ok
+							}
+						}
 						return nil, true
 					}
 

tools/tools_test.go

@@ -1274,6 +1274,22 @@ func TestFindArguments(t *testing.T) {
 				"items": []any{"{", "}", map[string]any{"key": "value"}},
 			},
 		},
+		{
+			name:   "stringified arguments",
+			buffer: []byte(`{"name": "get_temperature", "arguments": "{\"format\": \"fahrenheit\", \"location\": \"San Francisco, CA\"}"}`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+		},
+		{
+			name:   "stringified parameters",
+			buffer: []byte(`{"name": "get_temperature", "parameters": "{\"format\": \"fahrenheit\", \"location\": \"San Francisco, CA\"}"}`),
+			want: map[string]any{
+				"format":   "fahrenheit",
+				"location": "San Francisco, CA",
+			},
+		},
 	}
 
 	for _, tt := range tests {