ci: test darwin xcode pin

Implements Cohere2MoeForCausalLM (e.g. CohereLabs/North-Mini-Code-1.0)

.github/workflows/test-darwin-xcode-pin.yaml

@@ -0,0 +1,98 @@
+name: test-darwin-xcode-pin
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - test/darwin-xcode-pin
+  pull_request:
+    paths:
+      - '.github/workflows/test-darwin-xcode-pin.yaml'
+      - 'scripts/build_darwin.sh'
+      - 'MLX_VERSION'
+      - 'MLX_C_VERSION'
+      - 'cmake/**'
+      - 'x/mlxrunner/**'
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+env:
+  CGO_CFLAGS: '-O3'
+  CGO_CXXFLAGS: '-O3'
+  PINNED_DEVELOPER_DIR: /Applications/Xcode_26.4.1.app/Contents/Developer
+
+jobs:
+  darwin-build:
+    runs-on: macos-26-xlarge
+    env:
+      CGO_CFLAGS: '-mmacosx-version-min=14.0 -O3'
+      CGO_CXXFLAGS: '-mmacosx-version-min=14.0 -O3'
+      CGO_LDFLAGS: '-mmacosx-version-min=14.0 -O3'
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set build environment
+        shell: bash
+        run: |
+          set -euo pipefail
+          VERSION="0.0.0-xcode-pin-${GITHUB_SHA::7}"
+          {
+            echo "VERSION=${VERSION}"
+            echo "GOFLAGS='-ldflags=-w -s \"-X=github.com/ollama/ollama/version.Version=${VERSION}\" \"-X=github.com/ollama/ollama/server.mode=release\"'"
+          } >>"${GITHUB_ENV}"
+
+      - name: Select Xcode 26.4.1
+        shell: bash
+        run: |
+          set -euo pipefail
+          if [ ! -d "${PINNED_DEVELOPER_DIR}" ]; then
+            echo "Missing ${PINNED_DEVELOPER_DIR}"
+            ls -1 /Applications | grep '^Xcode' || true
+            exit 1
+          fi
+
+          sudo xcode-select -s "${PINNED_DEVELOPER_DIR}"
+          echo "DEVELOPER_DIR=${PINNED_DEVELOPER_DIR}" >>"${GITHUB_ENV}"
+
+          sw_vers
+          xcodebuild -version
+          xcrun --sdk macosx --show-sdk-version
+          xcrun --find metal
+
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+          cache-dependency-path: |
+            go.sum
+            LLAMA_CPP_VERSION
+            MLX_VERSION
+            MLX_C_VERSION
+
+      - name: Build unsigned Darwin runtime
+        run: ./scripts/build_darwin.sh build package
+
+      - name: Verify MLX payload
+        shell: bash
+        run: |
+          set -euo pipefail
+          test -f dist/darwin/lib/ollama/mlx_metal_v3/libmlxc.dylib
+          test -f dist/darwin/lib/ollama/mlx_metal_v3/mlx.metallib
+          test -f dist/darwin/lib/ollama/mlx_metal_v4/libmlxc.dylib
+          test -f dist/darwin/lib/ollama/mlx_metal_v4/mlx.metallib
+          find dist/darwin/lib/ollama -maxdepth 3 -type f \( -name 'libmlx*.dylib' -o -name '*.metallib' \) -print
+          lipo -archs dist/darwin/lib/ollama/mlx_metal_v3/libmlxc.dylib
+          lipo -archs dist/darwin/lib/ollama/mlx_metal_v4/libmlxc.dylib
+
+      - name: Log build results
+        run: ls -l dist/
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: ollama-darwin-xcode-pin
+          path: dist/ollama-darwin.tgz
+          compression-level: 0

LLAMA_CPP_VERSION

@@ -1 +1 @@
-b9509
+b9672

integration/context_test.go

@@ -17,9 +17,10 @@ import (
 func TestLongInputContext(t *testing.T) {
 	// Setting NUM_PARALLEL to 1 ensures the allocated context is exactly what
 	// we asked for and there is nothing extra that we could spill over into.
-	// Older runners silently truncate oversized prompts, while llama-server
-	// rejects them with a client error. Accept both behaviors so this test can
-	// run against main and the llama-server branch.
+	// Context shift happens after a prompt has been admitted to a slot. Initial
+	// prompts that fill or exceed the slot are still rejected by llama-server.
+	// Accept a context-limit error here because older runners may truncate this
+	// prompt while llama-server reports it as too large to admit.
 	t.Setenv("OLLAMA_NUM_PARALLEL", "1")
 
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
@@ -74,6 +75,7 @@ func isContextLimitError(err string) bool {
 	return strings.Contains(err, "context") &&
 		(strings.Contains(err, "exceed") ||
 			strings.Contains(err, "too large") ||
+			strings.Contains(err, "longer") ||
 			strings.Contains(err, "too long"))
 }
 

integration/create_test.go

@@ -48,10 +48,8 @@ func skipIfRemote(t *testing.T) {
 
 // findHFCLI returns the path to the HuggingFace CLI, or "" if not found.
 func findHFCLI() string {
-	for _, name := range []string{"huggingface-cli", "hf"} {
-		if p, err := exec.LookPath(name); err == nil {
-			return p
-		}
+	if p, err := exec.LookPath("hf"); err == nil {
+		return p
 	}
 	return ""
 }

llama/compat/README.md

@@ -23,11 +23,14 @@ intentionally skipped so a developer can iterate on a local llama.cpp tree.
 - `llama-ollama-compat-util.h`, `llama-ollama-compat-util.cpp` - helpers for
   KV edits, tensor renames, skip-prefix tracking, tensor load operations, and
   small tensor repacking primitives.
-- `llama-cpp-hooks.patch` - small additive call-site edits in llama.cpp files.
+- `001-llama-cpp-hooks.patch` - small additive call-site edits in llama.cpp files.
   It currently touches `src/llama-model-loader.cpp` and `tools/mtmd/clip.cpp`.
+- `002-llama-cpp-ui-empty-assets.patch` - lets the llama.cpp UI embed helper
+  generate an empty asset table when no UI assets are present.
 - `compat.cmake`, `apply-patch.cmake` - CMake glue and an idempotent applier
   (used by `llama/server/CMakeLists.txt`) that applies every `*.patch` under
-  this directory — the hooks patch plus each `models/` architecture patch.
+  this directory by numeric filename order — the hooks patch plus each
+  `models/` architecture patch.
 - `models/` - the sibling **new-architecture** layer: implementations of
   architectures llama.cpp doesn't support yet, each added via a small
   registration patch. (Those files *add* archs; the files above *translate*
@@ -116,7 +119,7 @@ cd /path/to/llama.cpp
 git diff -- \
     src/llama-model-loader.cpp \
     tools/mtmd/clip.cpp \
-    > /path/to/ollama/llama/compat/llama-cpp-hooks.patch
+    > /path/to/ollama/llama/compat/001-llama-cpp-hooks.patch
 ```
 
 ## Implementation Notes

llama/compat/apply-patch.cmake

@@ -3,10 +3,11 @@
 # Invocation (from a CMake PATCH_COMMAND):
 #   cmake -DPATCH_DIR=<dir of *.patch> -P apply-patch.cmake
 #
-# Every *.patch under PATCH_DIR is applied in the current working directory
-# (which ExternalProject / FetchContent sets to the fetched source's
-# SOURCE_DIR). A patch already applied — detected via `git apply --reverse
-# --check` — is skipped. This makes re-configuring and re-building safe.
+# Every *.patch under PATCH_DIR is applied in numeric filename order in the
+# current working directory (which ExternalProject / FetchContent sets to the
+# fetched source's SOURCE_DIR). A patch already applied — detected via
+# `git apply --reverse --check` — is skipped. This makes re-configuring and
+# re-building safe.
 
 if(NOT DEFINED PATCH_DIR)
     message(FATAL_ERROR "apply-patch.cmake: PATCH_DIR not set")
@@ -19,8 +20,17 @@ get_filename_component(_git_ceiling "${_patch_workdir}" DIRECTORY)
 set(_git_apply_env GIT_CEILING_DIRECTORIES=${_git_ceiling})
 
 file(GLOB_RECURSE _patches "${PATCH_DIR}/*.patch")
-list(SORT _patches)
+set(_patch_entries)
 foreach(PATCH_FILE IN LISTS _patches)
+    get_filename_component(_patch_name "${PATCH_FILE}" NAME)
+    list(APPEND _patch_entries "${_patch_name}|${PATCH_FILE}")
+endforeach()
+
+list(SORT _patch_entries)
+foreach(_patch_entry IN LISTS _patch_entries)
+    string(REGEX REPLACE "^[^|]*\\|" "" PATCH_FILE "${_patch_entry}")
+    file(RELATIVE_PATH _patch_rel "${PATCH_DIR}" "${PATCH_FILE}")
+
     # If the patch can be REVERSED cleanly, it's already applied. Skip.
     execute_process(
         COMMAND ${CMAKE_COMMAND} -E env ${_git_apply_env}
@@ -29,7 +39,7 @@ foreach(PATCH_FILE IN LISTS _patches)
         OUTPUT_QUIET ERROR_QUIET
     )
     if(_reverse_check EQUAL 0)
-        message(STATUS "llama/compat: patch already applied, skipping")
+        message(STATUS "llama/compat: ${_patch_rel} already applied, skipping")
         continue()
     endif()
 
@@ -41,10 +51,10 @@ foreach(PATCH_FILE IN LISTS _patches)
     )
     if(NOT _apply_result EQUAL 0)
         message(FATAL_ERROR
-            "llama/compat: failed to apply ${PATCH_FILE}\n"
+            "llama/compat: failed to apply ${_patch_rel}\n"
             "This usually means the pinned llama.cpp source has changed. "
             "Regenerate the patch against the pinned LLAMA_CPP_VERSION and retry.")
     endif()
 
-    message(STATUS "llama/compat: applied patch")
+    message(STATUS "llama/compat: applied ${_patch_rel}")
 endforeach()

llama/compat/compat.cmake

@@ -18,7 +18,7 @@
 # The compat layer consists of:
 #   1. Ollama-owned compat source files linked into the fetched llama.cpp
 #      targets from this directory.
-#   2. A small patch file that adds call-sites in llama.cpp loaders.
+#   2. A small ordered patch set that adds call-sites in llama.cpp loaders.
 
 set(_compat_dir ${CMAKE_CURRENT_LIST_DIR})
 
@@ -46,7 +46,7 @@ set(OLLAMA_LLAMA_CPP_COMPAT_DIR
 # Also export the individual paths in case callers want to do something
 # custom (e.g. emit a dependency on the patch so reconfigures re-apply).
 set(OLLAMA_LLAMA_CPP_COMPAT_PATCH_FILE
-    "${_compat_dir}/llama-cpp-hooks.patch"
+    "${_compat_dir}/001-llama-cpp-hooks.patch"
     CACHE INTERNAL "Path to the llama.cpp compat patch")
 
 set(OLLAMA_LLAMA_CPP_COMPAT_SOURCES

llama/compat/models/003-llama-cpp-laguna.patch

@@ -16,8 +16,8 @@ diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
 diff --git a/src/llama-arch.h b/src/llama-arch.h
 --- a/src/llama-arch.h
 +++ b/src/llama-arch.h
-@@ -140,2 +140,3 @@ enum llm_arch {
-     LLM_ARCH_MELLUM,
+@@ -145,2 +145,3 @@ enum llm_arch {
+     LLM_ARCH_EAGLE3,
 +    LLM_ARCH_LAGUNA,
      LLM_ARCH_UNKNOWN,
 @@ -382,2 +383,3 @@ enum llm_tensor {

llama/compat/models/laguna.cpp

@@ -14,7 +14,7 @@ void llama_model_laguna::load_arch_hparams(llama_model_loader & ml) {
 
     ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
     hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
-    ml.get_key_or_arr("laguna.attention.layer_types", hparams.is_swa_impl, hparams.n_layer, false);
+    ml.get_key_or_arr("laguna.attention.layer_types", hparams.is_swa_impl, hparams.n_layer(), false);
 
     ml.get_key("laguna.rope.swa.dimension_count", hparams.n_rot_swa,                false);
     ml.get_key("laguna.rope.swa.freq_base",       hparams.rope_freq_base_train_swa, false);

llm/llama_server.go

@@ -273,15 +273,18 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req
 		return nil, err
 	}
 
-	// llama-server rejects prompts that fill the entire slot context, while the
-	// old runner could accept exactly num_ctx prompt tokens. Keep one token of
-	// headroom so token-level truncation preserves old behavior as closely as
-	// llama-server allows.
 	limit := s.options.NumCtx - 1
 	if len(tokens) <= limit {
 		return prompt, nil
 	}
 
+	if !s.launch.config.ContextShift {
+		return nil, api.StatusError{
+			StatusCode:   http.StatusBadRequest,
+			ErrorMessage: "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length",
+		}
+	}
+
 	nKeep := req.Options.NumKeep
 	if nKeep < 0 {
 		nKeep = len(tokens)
@@ -293,7 +296,7 @@ func (s *llamaServerRunner) completionPromptForRequest(ctx context.Context, req
 	truncated = append(truncated, tokens[:nKeep]...)
 	truncated = append(truncated, tokens[nKeep+discard:]...)
 
-	slog.Warn("truncating input prompt", "limit", s.options.NumCtx, "prompt", len(tokens), "keep", nKeep, "new", len(truncated))
+	slog.Warn("truncating input prompt", "limit", limit, "prompt", len(tokens), "keep", nKeep, "new", len(truncated))
 	return truncated, nil
 }
 

llm/llama_server_test.go

@@ -488,14 +488,16 @@ func TestLlamaServerCompletionForwardsRepeatLastNZero(t *testing.T) {
 	}
 }
 
-func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
-	var completionReq llamaServerCompletionRequest
+func TestLlamaServerCompletionRejectsPromptOverContext(t *testing.T) {
+	const wantError = "the prompt is longer than the context length currently available to the model; shorten the prompt, adjust the context length in settings, or use a model with a longer context length"
+
 	var tokenizeReq struct {
 		Content      string `json:"content"`
 		AddSpecial   bool   `json:"add_special"`
 		ParseSpecial *bool  `json:"parse_special"`
 	}
 
+	completionCalled := false
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/health":
@@ -507,10 +509,7 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 			}
 			fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`)
 		case "/completion":
-			if err := json.NewDecoder(r.Body).Decode(&completionReq); err != nil {
-				t.Errorf("invalid completion request body: %v", err)
-				return
-			}
+			completionCalled = true
 			w.Header().Set("Content-Type", "text/event-stream")
 			fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":7,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
 		default:
@@ -537,8 +536,15 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 		Options:  &opts,
 		Truncate: true,
 	}, func(cr CompletionResponse) {})
-	if err != nil {
-		t.Fatalf("Completion error: %v", err)
+	var statusErr api.StatusError
+	if !errors.As(err, &statusErr) {
+		t.Fatalf("Completion error = %T %v, want api.StatusError", err, err)
+	}
+	if statusErr.StatusCode != http.StatusBadRequest {
+		t.Fatalf("StatusCode = %d, want %d", statusErr.StatusCode, http.StatusBadRequest)
+	}
+	if statusErr.ErrorMessage != wantError {
+		t.Fatalf("ErrorMessage = %q, want %q", statusErr.ErrorMessage, wantError)
 	}
 
 	if tokenizeReq.Content != strings.Repeat("long prompt ", 2) {
@@ -547,10 +553,149 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 	if !tokenizeReq.AddSpecial {
 		t.Fatal("expected tokenize request to add special tokens")
 	}
+	if completionCalled {
+		t.Fatal("completion endpoint was called")
+	}
+}
 
-	got, ok := completionReq.Prompt.([]any)
+func TestLlamaServerCompletionContextShiftAllowsPromptWithHeadroom(t *testing.T) {
+	var capturedReq llamaServerCompletionRequest
+	var tokenizeReq struct {
+		Content      string `json:"content"`
+		AddSpecial   bool   `json:"add_special"`
+		ParseSpecial *bool  `json:"parse_special"`
+	}
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/health":
+			fmt.Fprint(w, `{"status":"ok"}`)
+		case "/tokenize":
+			if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil {
+				t.Errorf("invalid tokenize request body: %v", err)
+				return
+			}
+			fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6]}`)
+		case "/completion":
+			if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil {
+				t.Errorf("invalid completion request body: %v", err)
+				return
+			}
+			w.Header().Set("Content-Type", "text/event-stream")
+			fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":10,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
+		default:
+			t.Errorf("unexpected path: %s", r.URL.Path)
+		}
+	}))
+	defer srv.Close()
+
+	parts := strings.Split(srv.URL, ":")
+	var portInt int
+	fmt.Sscanf(parts[len(parts)-1], "%d", &portInt)
+
+	runner := &llamaServerRunner{
+		port:    portInt,
+		cmd:     fakeRunningCmd(),
+		sem:     semaphore.NewWeighted(1),
+		options: api.Options{Runner: api.Runner{NumCtx: 8}},
+		launch: llamaServerLaunchConfig{
+			config: LlamaServerConfig{ContextShift: true},
+		},
+	}
+
+	opts := api.DefaultOptions()
+	opts.NumKeep = 3
+	prompt := strings.Repeat("long prompt ", 2)
+	err := runner.Completion(t.Context(), CompletionRequest{
+		Prompt:   prompt,
+		Options:  &opts,
+		Truncate: true,
+	}, func(cr CompletionResponse) {})
+	if err != nil {
+		t.Fatalf("Completion error: %v", err)
+	}
+
+	if tokenizeReq.Content != prompt {
+		t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt)
+	}
+	if !tokenizeReq.AddSpecial {
+		t.Fatal("expected tokenize request to add special tokens")
+	}
+	if capturedReq.Prompt != prompt {
+		t.Fatalf("prompt = %q, want %q", capturedReq.Prompt, prompt)
+	}
+	if capturedReq.NKeep != opts.NumKeep {
+		t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep)
+	}
+}
+
+func TestLlamaServerCompletionContextShiftTruncatesPromptOverContext(t *testing.T) {
+	var capturedReq llamaServerCompletionRequest
+	var tokenizeReq struct {
+		Content      string `json:"content"`
+		AddSpecial   bool   `json:"add_special"`
+		ParseSpecial *bool  `json:"parse_special"`
+	}
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/health":
+			fmt.Fprint(w, `{"status":"ok"}`)
+		case "/tokenize":
+			if err := json.NewDecoder(r.Body).Decode(&tokenizeReq); err != nil {
+				t.Errorf("invalid tokenize request body: %v", err)
+				return
+			}
+			fmt.Fprint(w, `{"tokens":[0,1,2,3,4,5,6,7,8,9]}`)
+		case "/completion":
+			if err := json.NewDecoder(r.Body).Decode(&capturedReq); err != nil {
+				t.Errorf("invalid completion request body: %v", err)
+				return
+			}
+			w.Header().Set("Content-Type", "text/event-stream")
+			fmt.Fprintln(w, `data: {"content":"ok","stop":true,"timings":{"prompt_n":7,"prompt_ms":1,"predicted_n":1,"predicted_ms":1}}`)
+		default:
+			t.Errorf("unexpected path: %s", r.URL.Path)
+		}
+	}))
+	defer srv.Close()
+
+	parts := strings.Split(srv.URL, ":")
+	var portInt int
+	fmt.Sscanf(parts[len(parts)-1], "%d", &portInt)
+
+	runner := &llamaServerRunner{
+		port:    portInt,
+		cmd:     fakeRunningCmd(),
+		sem:     semaphore.NewWeighted(1),
+		options: api.Options{Runner: api.Runner{NumCtx: 8}},
+		launch: llamaServerLaunchConfig{
+			config: LlamaServerConfig{ContextShift: true},
+		},
+	}
+
+	opts := api.DefaultOptions()
+	opts.NumKeep = 3
+	prompt := strings.Repeat("long prompt ", 2)
+	err := runner.Completion(t.Context(), CompletionRequest{
+		Prompt:   prompt,
+		Options:  &opts,
+		Truncate: true,
+	}, func(cr CompletionResponse) {})
+	if err != nil {
+		t.Fatalf("Completion error: %v", err)
+	}
+
+	if tokenizeReq.Content != prompt {
+		t.Fatalf("tokenize content = %q, want %q", tokenizeReq.Content, prompt)
+	}
+	if !tokenizeReq.AddSpecial {
+		t.Fatal("expected tokenize request to add special tokens")
+	}
+
+	got, ok := capturedReq.Prompt.([]any)
 	if !ok {
-		t.Fatalf("completion prompt = %T, want token array", completionReq.Prompt)
+		t.Fatalf("completion prompt = %T, want token array", capturedReq.Prompt)
 	}
 	want := []int{0, 1, 2, 6, 7, 8, 9}
 	if len(got) != len(want) {
@@ -562,6 +707,9 @@ func TestLlamaServerCompletionTruncatesPromptAsTokens(t *testing.T) {
 			t.Fatalf("token prompt[%d] = %#v, want %d", i, got[i], wantToken)
 		}
 	}
+	if capturedReq.NKeep != opts.NumKeep {
+		t.Fatalf("n_keep = %d, want %d", capturedReq.NKeep, opts.NumKeep)
+	}
 }
 
 func TestLlamaServerCompletionWithMediaUsesRunnerMarker(t *testing.T) {

model/parsers/cohere.go

@@ -0,0 +1,315 @@
+package parsers
+
+import (
+	"encoding/json"
+	"log/slog"
+	"strings"
+	"unicode"
+
+	"github.com/ollama/ollama/api"
+)
+
+// CohereParser parses output from Cohere North / Command A 2026 models
+// (e.g. North-Mini-Code-1.0). The generation prompt ends with
+// <|START_THINKING|> (reasoning on) or <|START_THINKING|><|END_THINKING|>
+// (reasoning off), so output begins inside the thinking block when reasoning
+// is enabled. After thinking, the model emits either
+// <|START_TEXT|>content<|END_TEXT|> or an <|START_ACTION|>[...]<|END_ACTION|>
+// tool call array, then <|END_OF_TURN_TOKEN|>.
+type CohereParser struct {
+	state     cohereParserState
+	buffer    strings.Builder
+	callIndex int
+}
+
+type cohereParserState int
+
+const (
+	cohereCollectingThinking cohereParserState = iota
+	cohereAwaitingBlock
+	cohereCollectingContent
+	cohereCollectingAction
+)
+
+const (
+	cohereEndThinking = "<|END_THINKING|>"
+	cohereStartText   = "<|START_TEXT|>"
+	cohereEndText     = "<|END_TEXT|>"
+	cohereStartAction = "<|START_ACTION|>"
+	cohereEndAction   = "<|END_ACTION|>"
+	cohereEndOfTurn   = "<|END_OF_TURN_TOKEN|>"
+
+	// Legacy response markers from the older Command A chat template, which
+	// also ships in these models' tokenizer_config and shows up in sampled
+	// output; treated as aliases for START_TEXT / END_TEXT.
+	cohereStartResponse = "<|START_RESPONSE|>"
+	cohereEndResponse   = "<|END_RESPONSE|>"
+)
+
+func (p *CohereParser) HasToolSupport() bool {
+	return true
+}
+
+func (p *CohereParser) HasThinkingSupport() bool {
+	return true
+}
+
+func (p *CohereParser) PreservedTokens() []string {
+	return []string{
+		"<|START_THINKING|>", cohereEndThinking,
+		cohereStartText, cohereEndText,
+		cohereStartAction, cohereEndAction,
+		cohereStartResponse, cohereEndResponse,
+	}
+}
+
+func (p *CohereParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
+	p.buffer.Reset()
+	p.callIndex = 0
+
+	// The template enables reasoning by default; nil means default.
+	thinkingEnabled := thinkValue == nil || thinkValue.Bool()
+
+	assistantPrefill := lastMessage != nil && lastMessage.Role == "assistant" && lastMessage.Content != ""
+	switch {
+	case assistantPrefill:
+		// The prompt left an open <|START_TEXT|> for continuation.
+		p.state = cohereCollectingContent
+	case thinkingEnabled:
+		p.state = cohereCollectingThinking
+	default:
+		p.state = cohereAwaitingBlock
+	}
+	return tools
+}
+
+func (p *CohereParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
+	p.buffer.WriteString(s)
+
+	var contentSb, thinkingSb strings.Builder
+	for {
+		c, t, tc, more := p.eat(done)
+		contentSb.WriteString(c)
+		thinkingSb.WriteString(t)
+		calls = append(calls, tc...)
+		if !more {
+			break
+		}
+	}
+
+	for i := range calls {
+		calls[i].Function.Index = p.callIndex
+		p.callIndex++
+	}
+
+	return contentSb.String(), thinkingSb.String(), calls, nil
+}
+
+// eat consumes what it can from the buffer for the current state. It returns
+// more=true when a state transition happened and the remaining buffer should
+// be reprocessed.
+func (p *CohereParser) eat(done bool) (content string, thinking string, calls []api.ToolCall, more bool) {
+	buf := p.buffer.String()
+	if buf == "" {
+		return "", "", nil, false
+	}
+
+	switch p.state {
+	case cohereCollectingThinking:
+		if idx := strings.Index(buf, cohereEndThinking); idx != -1 {
+			thinking = strings.TrimRightFunc(buf[:idx], unicode.IsSpace)
+			p.resetBuffer(buf[idx+len(cohereEndThinking):])
+			p.state = cohereAwaitingBlock
+			return "", thinking, nil, true
+		}
+		// Emit all but a possible partial tag / trailing whitespace.
+		keep := overlap(buf, cohereEndThinking)
+		emitEnd := len(buf) - keep
+		emitEnd -= trailingWhitespaceLen(buf[:emitEnd])
+		if emitEnd > 0 {
+			thinking = buf[:emitEnd]
+			p.resetBuffer(buf[emitEnd:])
+		}
+		return "", thinking, nil, false
+
+	case cohereAwaitingBlock:
+		// Between blocks: look for the next text or action opener, skipping
+		// whitespace and the end-of-turn token.
+		for _, open := range []string{cohereStartText, cohereStartResponse} {
+			if idx := strings.Index(buf, open); idx != -1 {
+				p.resetBuffer(buf[idx+len(open):])
+				p.state = cohereCollectingContent
+				return "", "", nil, true
+			}
+		}
+		if idx := strings.Index(buf, cohereStartAction); idx != -1 {
+			p.resetBuffer(buf[idx+len(cohereStartAction):])
+			p.state = cohereCollectingAction
+			return "", "", nil, true
+		}
+		trimmed := strings.TrimLeftFunc(buf, unicode.IsSpace)
+		if strings.HasPrefix(trimmed, cohereEndOfTurn) {
+			p.resetBuffer(trimmed[len(cohereEndOfTurn):])
+			return "", "", nil, true
+		}
+		if trimmed == "" {
+			if done {
+				p.buffer.Reset()
+			}
+			return "", "", nil, false
+		}
+		// Wait only while the buffer could still grow into one of the
+		// expected tags. Anything else — bare content or an unrecognized
+		// tag — streams out as content rather than buffering forever.
+		if !done && maybePartialTag(trimmed) {
+			return "", "", nil, false
+		}
+		p.buffer.Reset()
+		p.state = cohereCollectingContent
+		p.buffer.WriteString(trimmed)
+		return "", "", nil, true
+
+	case cohereCollectingContent:
+		// END_OF_TURN also closes content for models that skip END_TEXT.
+		for _, close := range []string{cohereEndText, cohereEndResponse, cohereEndOfTurn} {
+			if idx := strings.Index(buf, close); idx != -1 {
+				content = buf[:idx]
+				p.resetBuffer(buf[idx+len(close):])
+				p.state = cohereAwaitingBlock
+				return content, "", nil, true
+			}
+		}
+		keep := max(overlap(buf, cohereEndText), overlap(buf, cohereEndResponse), overlap(buf, cohereEndOfTurn))
+		emitEnd := len(buf) - keep
+		if emitEnd > 0 {
+			content = buf[:emitEnd]
+			p.resetBuffer(buf[emitEnd:])
+		}
+		if done && p.buffer.Len() > 0 {
+			content += p.buffer.String()
+			p.buffer.Reset()
+		}
+		return content, "", nil, false
+
+	case cohereCollectingAction:
+		if idx := strings.Index(buf, cohereEndAction); idx != -1 {
+			payload := buf[:idx]
+			p.resetBuffer(buf[idx+len(cohereEndAction):])
+			p.state = cohereAwaitingBlock
+			calls = parseCohereActions(payload)
+			return "", "", calls, true
+		}
+		if done {
+			// Best effort on truncated output.
+			calls = parseCohereActions(buf)
+			p.buffer.Reset()
+			return "", "", calls, false
+		}
+		return "", "", nil, false
+	}
+
+	return "", "", nil, false
+}
+
+func (p *CohereParser) resetBuffer(s string) {
+	p.buffer.Reset()
+	p.buffer.WriteString(s)
+}
+
+// maybePartialTag reports whether s is a proper prefix of one of the tags the
+// awaiting-block state recognizes — the only case worth waiting on for more
+// output before treating the buffer as content.
+func maybePartialTag(s string) bool {
+	for _, tag := range []string{cohereStartText, cohereStartResponse, cohereStartAction, cohereEndOfTurn} {
+		if len(s) < len(tag) && strings.HasPrefix(tag, s) {
+			return true
+		}
+	}
+	return false
+}
+
+type cohereToolCall struct {
+	ToolCallID string                        `json:"tool_call_id"`
+	ToolName   string                        `json:"tool_name"`
+	Parameters api.ToolCallFunctionArguments `json:"parameters"`
+}
+
+func (c cohereToolCall) toolCall() api.ToolCall {
+	return api.ToolCall{
+		ID: c.ToolCallID,
+		Function: api.ToolCallFunction{
+			Name:      c.ToolName,
+			Arguments: c.Parameters,
+		},
+	}
+}
+
+// parseCohereActions parses the JSON array inside an action block:
+// [{"tool_call_id": "0", "tool_name": ..., "parameters": {...}}, ...]
+//
+// Sampled output occasionally malforms the JSON (a missing comma between
+// calls, an unquoted value). When the array as a whole fails to parse, fall
+// back to scanning its balanced top-level objects and parse each
+// independently, so one bad call doesn't drop its siblings.
+func parseCohereActions(payload string) []api.ToolCall {
+	payload = strings.TrimSpace(payload)
+	if payload == "" {
+		return nil
+	}
+
+	var parsed []cohereToolCall
+	if err := json.Unmarshal([]byte(payload), &parsed); err == nil {
+		calls := make([]api.ToolCall, 0, len(parsed))
+		for _, c := range parsed {
+			calls = append(calls, c.toolCall())
+		}
+		return calls
+	}
+
+	var calls []api.ToolCall
+	for _, obj := range scanJSONObjects(payload) {
+		var c cohereToolCall
+		if err := json.Unmarshal([]byte(obj), &c); err != nil {
+			if len(obj) > 200 {
+				obj = obj[:200] + "…"
+			}
+			slog.Warn("cohere action parsing failed", "error", err, "action", obj)
+			continue
+		}
+		calls = append(calls, c.toolCall())
+	}
+	return calls
+}
+
+// scanJSONObjects returns the balanced top-level {...} chunks of s, tracking
+// strings and escapes so braces inside values don't split objects.
+func scanJSONObjects(s string) []string {
+	var objects []string
+	depth, start := 0, -1
+	inString, escaped := false, false
+	for i := range len(s) {
+		switch c := s[i]; {
+		case escaped:
+			escaped = false
+		case c == '\\' && inString:
+			escaped = true
+		case c == '"':
+			inString = !inString
+		case inString:
+		case c == '{':
+			if depth == 0 {
+				start = i
+			}
+			depth++
+		case c == '}':
+			if depth > 0 {
+				depth--
+				if depth == 0 && start >= 0 {
+					objects = append(objects, s[start:i+1])
+					start = -1
+				}
+			}
+		}
+	}
+	return objects
+}

model/parsers/cohere_test.go

@@ -0,0 +1,268 @@
+package parsers
+
+import (
+	"testing"
+
+	"github.com/ollama/ollama/api"
+)
+
+// Model output begins inside <|START_THINKING|> when reasoning is on (the
+// generation prompt ends with that tag).
+
+func cohereAddAll(t *testing.T, p *CohereParser, chunks []string) (content, thinking string, calls []api.ToolCall) {
+	t.Helper()
+	for i, c := range chunks {
+		done := i == len(chunks)-1
+		ct, th, tc, err := p.Add(c, done)
+		if err != nil {
+			t.Fatal(err)
+		}
+		content += ct
+		thinking += th
+		calls = append(calls, tc...)
+	}
+	return content, thinking, calls
+}
+
+func TestCohereParseThinkingThenText(t *testing.T) {
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+
+	content, thinking, calls := cohereAddAll(t, p, []string{
+		"Let me think", " about this.<|END_THINKING|>",
+		"<|START_TEXT|>Hello", " world!<|END_TEXT|>",
+	})
+	if thinking != "Let me think about this." {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "Hello world!" {
+		t.Errorf("content = %q", content)
+	}
+	if len(calls) != 0 {
+		t.Errorf("unexpected calls: %v", calls)
+	}
+}
+
+func TestCohereParseSplitTags(t *testing.T) {
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+
+	// Tags split across chunk boundaries must not leak into output.
+	content, thinking, _ := cohereAddAll(t, p, []string{
+		"think<|END_TH", "INKING|><|STAR", "T_TEXT|>ans", "wer<|END_", "TEXT|>",
+	})
+	if thinking != "think" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "answer" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParseToolCall(t *testing.T) {
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+
+	content, thinking, calls := cohereAddAll(t, p, []string{
+		"plan<|END_THINKING|><|START_ACTION|>[\n",
+		`    {"tool_call_id": "0", "tool_name": "get_weather", "parameters": {"city": "Paris"}},`,
+		`    {"tool_call_id": "1", "tool_name": "get_time", "parameters": {}}`,
+		"\n]<|END_ACTION|>",
+	})
+	if thinking != "plan" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "" {
+		t.Errorf("content = %q", content)
+	}
+	if len(calls) != 2 {
+		t.Fatalf("calls = %d, want 2", len(calls))
+	}
+	if calls[0].Function.Name != "get_weather" || calls[1].Function.Name != "get_time" {
+		t.Errorf("call names = %q, %q", calls[0].Function.Name, calls[1].Function.Name)
+	}
+	if v, ok := calls[0].Function.Arguments.Get("city"); !ok || v != "Paris" {
+		t.Errorf("call 0 city = %v %v", v, ok)
+	}
+	if calls[0].Function.Index != 0 || calls[1].Function.Index != 1 {
+		t.Errorf("call indices = %d, %d", calls[0].Function.Index, calls[1].Function.Index)
+	}
+}
+
+func TestCohereParseReasoningOff(t *testing.T) {
+	p := &CohereParser{}
+	think := &api.ThinkValue{Value: false}
+	p.Init(nil, nil, think)
+
+	content, thinking, _ := cohereAddAll(t, p, []string{
+		"<|START_TEXT|>direct answer<|END_TEXT|>",
+	})
+	if thinking != "" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "direct answer" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParseBareContent(t *testing.T) {
+	// Models occasionally skip the START_TEXT wrapper; treat raw text after
+	// thinking as content.
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+
+	content, thinking, _ := cohereAddAll(t, p, []string{
+		"thought<|END_THINKING|>", "Just plain text", " output",
+	})
+	if thinking != "thought" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "Just plain text output" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParseEndOfTurnWithoutEndText(t *testing.T) {
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+
+	content, _, _ := cohereAddAll(t, p, []string{
+		"t<|END_THINKING|><|START_TEXT|>answer<|END_OF_TURN_TOKEN|>",
+	})
+	if content != "answer" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParsePrefillContinuation(t *testing.T) {
+	p := &CohereParser{}
+	last := &api.Message{Role: "assistant", Content: "partial"}
+	p.Init(nil, last, nil)
+
+	content, thinking, _ := cohereAddAll(t, p, []string{" continued<|END_TEXT|>"})
+	if thinking != "" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != " continued" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParserRegistered(t *testing.T) {
+	p := ParserForName("cohere")
+	if p == nil {
+		t.Fatal("cohere parser not registered")
+	}
+	if !p.HasToolSupport() || !p.HasThinkingSupport() {
+		t.Error("cohere parser should support tools and thinking")
+	}
+}
+
+func TestCohereParseMalformedActions(t *testing.T) {
+	// One malformed call (unquoted value) must not drop its well-formed
+	// sibling, and a missing comma between calls must not drop either.
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+	_, _, calls := cohereAddAll(t, p, []string{
+		"plan<|END_THINKING|><|START_ACTION|>[\n",
+		`    {"tool_call_id": "0", "tool_name": "set_alarm", "parameters": {"time": 15:30}},`,
+		`    {"tool_call_id": "1", "tool_name": "get_weather", "parameters": {"city": "Oslo"}}`,
+		"\n]<|END_ACTION|>",
+	})
+	if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
+		t.Fatalf("calls = %v, want the well-formed get_weather call", calls)
+	}
+
+	p = &CohereParser{}
+	p.Init(nil, nil, nil)
+	_, _, calls = cohereAddAll(t, p, []string{
+		`<|END_THINKING|><|START_ACTION|>[`,
+		`{"tool_call_id": "0", "tool_name": "a", "parameters": {}}`,
+		`{"tool_call_id": "1", "tool_name": "b", "parameters": {"x": "{not json}"}}`,
+		`]<|END_ACTION|>`,
+	})
+	if len(calls) != 2 || calls[0].Function.Name != "a" || calls[1].Function.Name != "b" {
+		t.Fatalf("calls = %v, want both calls despite missing comma", calls)
+	}
+	if v, ok := calls[1].Function.Arguments.Get("x"); !ok || v != "{not json}" {
+		t.Fatalf("braces inside string values must not split objects, got %v", v)
+	}
+
+	// Unparseable garbage yields no calls and no panic.
+	p = &CohereParser{}
+	p.Init(nil, nil, nil)
+	_, _, calls = cohereAddAll(t, p, []string{"x<|END_THINKING|><|START_ACTION|>[!!!]<|END_ACTION|>"})
+	if len(calls) != 0 {
+		t.Fatalf("calls = %v, want none", calls)
+	}
+}
+
+func TestCohereParseLegacyResponseMarkers(t *testing.T) {
+	// Models trained on the older Command A template sometimes emit
+	// <|START_RESPONSE|>/<|END_RESPONSE|> instead of START_TEXT/END_TEXT.
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+	content, thinking, _ := cohereAddAll(t, p, []string{
+		"plan<|END_THINKING|>", "<|START_RESPONSE|>Hello", " there<|END_RESPONSE|>",
+	})
+	if thinking != "plan" {
+		t.Errorf("thinking = %q", thinking)
+	}
+	if content != "Hello there" {
+		t.Errorf("content = %q", content)
+	}
+}
+
+func TestCohereParseStreamsBeforeDone(t *testing.T) {
+	// Regression test: output after END_THINKING that opens with an
+	// unrecognized tag must stream as it arrives, not buffer until the
+	// generation finishes (it previously buffered forever, presenting as a
+	// hung response).
+	for _, tc := range []struct {
+		name  string
+		chunk string
+	}{
+		{"legacy response marker", "<|START_RESPONSE|>The answer is 42."},
+		{"unrecognized tag", "<|TOOL_PLAN|>I should call the tool."},
+		{"bare content", "Just plain text."},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			p := &CohereParser{}
+			p.Init(nil, nil, nil)
+			if _, _, _, err := p.Add("x<|END_THINKING|>", false); err != nil {
+				t.Fatal(err)
+			}
+			content, _, _, err := p.Add(tc.chunk, false) // done=false: still streaming
+			if err != nil {
+				t.Fatal(err)
+			}
+			if content == "" {
+				t.Fatalf("content did not stream before done for %q", tc.chunk)
+			}
+		})
+	}
+}
+
+func TestCohereParseEndOfTurnBetweenBlocks(t *testing.T) {
+	// A literal end-of-turn marker between blocks is consumed, not shown.
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+	content, thinking, _ := cohereAddAll(t, p, []string{
+		"t<|END_THINKING|><|START_TEXT|>hi<|END_TEXT|><|END_OF_TURN_TOKEN|>",
+	})
+	if thinking != "t" || content != "hi" {
+		t.Errorf("thinking = %q, content = %q", thinking, content)
+	}
+}
+
+func TestCohereParseBareContentBeforeEndOfTurn(t *testing.T) {
+	// Bare content followed by an end-of-turn marker keeps the content.
+	p := &CohereParser{}
+	p.Init(nil, nil, nil)
+	content, _, _ := cohereAddAll(t, p, []string{
+		"t<|END_THINKING|>plain answer<|END_OF_TURN_TOKEN|>",
+	})
+	if content != "plain answer" {
+		t.Errorf("content = %q, want %q", content, "plain answer")
+	}
+}

model/parsers/lfm2.go

@@ -14,7 +14,12 @@ import (
 type LFM2ParserState int
 
 const (
-	LFM2CollectingThinking LFM2ParserState = iota
+	// LFM2LookingForThinking is the initial state when thinking is enabled.
+	// LFM2 models emit an explicit <think> tag only when they reason; a direct
+	// answer has no tag at all. This state waits to see which one the output
+	// begins with before committing to thinking or content.
+	LFM2LookingForThinking LFM2ParserState = iota
+	LFM2CollectingThinking
 	LFM2CollectingContent
 	LFM2CollectingToolCalls
 )
@@ -70,8 +75,10 @@ func (p *LFM2Parser) setInitialState(lastMessage *api.Message, thinkValue *api.T
 		return
 	}
 
-	p.state = LFM2CollectingThinking
-	p.needsThinkingLeadingTrim = true
+	// Thinking is enabled, but the model decides per-turn whether to reason. Wait
+	// for a leading <think> tag before treating output as thinking; otherwise it's
+	// a direct answer.
+	p.state = LFM2LookingForThinking
 }
 
 func (p *LFM2Parser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
@@ -109,6 +116,18 @@ func (lfm2EventToolCall) isLFM2Event()        {}
 
 func (p *LFM2Parser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
 	p.buffer.WriteString(s)
+
+	// On the final chunk a partial "<think>" prefix can never complete, so commit
+	// the buffered output as a direct answer rather than withholding it.
+	if done && p.state == LFM2LookingForThinking {
+		trimmed := strings.TrimLeftFunc(p.buffer.String(), unicode.IsSpace)
+		if !strings.HasPrefix(trimmed, lfm2ThinkingOpenTag) {
+			p.buffer.Reset()
+			p.buffer.WriteString(trimmed)
+			p.state = LFM2CollectingContent
+		}
+	}
+
 	events := p.parseEvents()
 
 	var toolCalls []api.ToolCall
@@ -180,6 +199,28 @@ func (p *LFM2Parser) eat() ([]lfm2Event, bool) {
 	}
 
 	switch p.state {
+	case LFM2LookingForThinking:
+		// Decide whether this turn is a reasoning turn (begins with <think>) or a
+		// direct answer (no tag). Leading whitespace is ignored either way.
+		trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
+		if strings.HasPrefix(trimmed, lfm2ThinkingOpenTag) {
+			after := trimmed[len(lfm2ThinkingOpenTag):]
+			p.buffer.Reset()
+			p.buffer.WriteString(after)
+			p.state = LFM2CollectingThinking
+			p.needsThinkingLeadingTrim = true
+			return events, true
+		}
+		if trimmed == "" || strings.HasPrefix(lfm2ThinkingOpenTag, trimmed) {
+			// Only whitespace so far, or a partial "<think>" prefix; wait for more.
+			return events, false
+		}
+		// Direct answer: no thinking block.
+		p.buffer.Reset()
+		p.buffer.WriteString(trimmed)
+		p.state = LFM2CollectingContent
+		return events, true
+
 	case LFM2CollectingThinking:
 		// Strip opening <think> tag if present
 		if strings.HasPrefix(bufStr, lfm2ThinkingOpenTag) {

model/parsers/lfm2_test.go

@@ -25,14 +25,36 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_content",
-			input:            "I need to think about this...</think>The answer is 42.",
+			input:            "<think>I need to think about this...</think>The answer is 42.",
 			expectedThinking: "I need to think about this...",
 			expectedContent:  "The answer is 42.",
 			hasThinking:      true,
 		},
+		{
+			name:            "direct_answer_with_thinking_enabled",
+			input:           "The answer is 42.",
+			expectedContent: "The answer is 42.",
+			hasThinking:     true,
+		},
+		{
+			name:            "direct_answer_with_tool_call",
+			input:           "I'll check the weather.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
+			expectedContent: "I'll check the weather.",
+			expectedCalls: []api.ToolCall{
+				{
+					Function: api.ToolCallFunction{
+						Name: "get_weather",
+						Arguments: testArgs(map[string]any{
+							"location": "Paris",
+						}),
+					},
+				},
+			},
+			hasThinking: true,
+		},
 		{
 			name:             "thinking_with_newlines",
-			input:            "Let me think:\n- Point 1\n- Point 2</think>\n\nHere's my answer.",
+			input:            "<think>Let me think:\n- Point 1\n- Point 2</think>\n\nHere's my answer.",
 			expectedThinking: "Let me think:\n- Point 1\n- Point 2",
 			expectedContent:  "Here's my answer.",
 			hasThinking:      true,
@@ -98,7 +120,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_with_tool_call",
-			input:            "Let me check the weather...</think>I'll get that for you.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
+			input:            "<think>Let me check the weather...</think>I'll get that for you.<|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|>",
 			expectedThinking: "Let me check the weather...",
 			expectedContent:  "I'll get that for you.",
 			expectedCalls: []api.ToolCall{
@@ -121,7 +143,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "only_thinking",
-			input:            "Just thinking content</think>",
+			input:            "<think>Just thinking content</think>",
 			expectedThinking: "Just thinking content",
 			expectedContent:  "",
 			hasThinking:      true,
@@ -140,7 +162,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_with_unicode",
-			input:            "我在思考这个问题...</think>答案是42。",
+			input:            "<think>我在思考这个问题...</think>答案是42。",
 			expectedThinking: "我在思考这个问题...",
 			expectedContent:  "答案是42。",
 			hasThinking:      true,
@@ -164,7 +186,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_with_special_chars",
-			input:            "Let me calculate: 2+2=4 & 3*3=9...</think>The results are correct!",
+			input:            "<think>Let me calculate: 2+2=4 & 3*3=9...</think>The results are correct!",
 			expectedThinking: "Let me calculate: 2+2=4 & 3*3=9...",
 			expectedContent:  "The results are correct!",
 			hasThinking:      true,
@@ -228,7 +250,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_then_python_tool_call",
-			input:            "I should check the status...</think>Let me look that up.<|tool_call_start|>[get_status(id=\"123\")]<|tool_call_end|>",
+			input:            "<think>I should check the status...</think>Let me look that up.<|tool_call_start|>[get_status(id=\"123\")]<|tool_call_end|>",
 			expectedThinking: "I should check the status...",
 			expectedContent:  "Let me look that up.",
 			expectedCalls: []api.ToolCall{
@@ -291,7 +313,7 @@ func TestLFM2Parser(t *testing.T) {
 		},
 		{
 			name:             "thinking_directly_to_tool_call",
-			input:            "Let me run this command...</think><|tool_call_start|>[bash(command='ls')]<|tool_call_end|>",
+			input:            "<think>Let me run this command...</think><|tool_call_start|>[bash(command='ls')]<|tool_call_end|>",
 			expectedThinking: "Let me run this command...",
 			expectedContent:  "",
 			expectedCalls: []api.ToolCall{
@@ -350,11 +372,17 @@ func TestLFM2Parser_Streaming(t *testing.T) {
 		},
 		{
 			name:             "streaming_thinking",
-			chunks:           []string{"I need to ", "think about this", "...</think>", "The answer is 42."},
+			chunks:           []string{"<think>", "I need to ", "think about this", "...</think>", "The answer is 42."},
 			expectedThinking: "I need to think about this...",
 			expectedContent:  "The answer is 42.",
 			hasThinking:      true,
 		},
+		{
+			name:            "streaming_direct_answer",
+			chunks:          []string{"The answer ", "is ", "42."},
+			expectedContent: "The answer is 42.",
+			hasThinking:     true,
+		},
 		{
 			name:            "streaming_tool_call",
 			chunks:          []string{"I'll check weather.", "<|tool_call_start|>", "[get_weather(", "location=\"Paris\")]", "<|tool_call_end|>"},
@@ -373,7 +401,7 @@ func TestLFM2Parser_Streaming(t *testing.T) {
 		},
 		{
 			name:             "streaming_thinking_with_partial_tag",
-			chunks:           []string{"Thinking about this", "...</", "think>", "Done thinking."},
+			chunks:           []string{"<think>", "Thinking about this", "...</", "think>", "Done thinking."},
 			expectedThinking: "Thinking about this...",
 			expectedContent:  "Done thinking.",
 			hasThinking:      true,
@@ -401,6 +429,22 @@ func TestLFM2Parser_Streaming(t *testing.T) {
 			},
 			hasThinking: false,
 		},
+		{
+			// Opening <think> tag split across chunks must still be recognized.
+			name:             "streaming_thinking_split_open_tag",
+			chunks:           []string{"<th", "ink>", "reasoning", "</think>", "answer"},
+			expectedThinking: "reasoning",
+			expectedContent:  "answer",
+			hasThinking:      true,
+		},
+		{
+			// A direct answer starting with '<' is ambiguous until enough arrives
+			// to rule out an opening <think> tag.
+			name:            "streaming_direct_answer_starting_with_angle",
+			chunks:          []string{"<", "html> is a tag"},
+			expectedContent: "<html> is a tag",
+			hasThinking:     true,
+		},
 		{
 			// Test that leading whitespace after <think> is trimmed even when in separate chunks
 			name:             "streaming_thinking_whitespace_after_tag",
@@ -506,9 +550,9 @@ func TestLFM2Parser_Init(t *testing.T) {
 		t.Errorf("Init() returned tools mismatch (-want +got):\n%s", diff)
 	}
 
-	// Test initial state is set to thinking when enabled
-	if parser.state != LFM2CollectingThinking {
-		t.Errorf("Expected initial state to be LFM2CollectingThinking, got %v", parser.state)
+	// Test initial state looks for a leading <think> tag when thinking is enabled
+	if parser.state != LFM2LookingForThinking {
+		t.Errorf("Expected initial state to be LFM2LookingForThinking, got %v", parser.state)
 	}
 }
 
@@ -1071,18 +1115,24 @@ func TestLFM2Parser_EdgeCases(t *testing.T) {
 	}{
 		{
 			name:             "multiple_think_close_tags",
-			input:            "First thought</think>Second thought</think>Final content",
+			input:            "<think>First thought</think>Second thought</think>Final content",
 			expectedThinking: "First thought",
 			expectedContent:  "Second thought</think>Final content",
 			hasThinking:      true,
 		},
 		{
-			name:             "empty_thinking_content",
-			input:            "</think>Just content",
+			name:             "empty_thinking_block",
+			input:            "<think></think>Just content",
 			expectedThinking: "",
 			expectedContent:  "Just content",
 			hasThinking:      true,
 		},
+		{
+			name:            "direct_answer_with_leading_whitespace",
+			input:           "  \n  Hello there",
+			expectedContent: "Hello there",
+			hasThinking:     true,
+		},
 		{
 			name:            "thinking_disabled_with_think_tags",
 			input:           "Some content</think>More content",

model/parsers/parsers.go

@@ -92,6 +92,8 @@ func ParserForName(name string) Parser {
 		return &LFM2Parser{hasThinkingSupport: true}
 	case "laguna":
 		return &LagunaParser{}
+	case "cohere":
+		return &CohereParser{}
 	default:
 		return nil
 	}

model/renderers/cohere.go

@@ -0,0 +1,233 @@
+package renderers
+
+import (
+	"encoding/json"
+	"strconv"
+	"strings"
+
+	"github.com/ollama/ollama/api"
+)
+
+// CohereRenderer renders the Cohere North / Command A 2026 chat template
+// (Cohere2 MoE models such as CohereLabs/North-Mini-Code-1.0): a platform
+// system turn with an Available Tools section, <|START_TEXT|>-wrapped message
+// bodies, <|START_THINKING|> reasoning, <|START_ACTION|> tool calls, and
+// <|START_TOOL_RESULT|> tool results.
+//
+// The chat template's model-specific platform instructions (identity, default
+// policies) belong in the model's Modelfile SYSTEM prompt, which arrives here
+// as the first system message.
+type CohereRenderer struct{}
+
+func (r *CohereRenderer) LeadingBOS() string {
+	return "<BOS_TOKEN>"
+}
+
+// cohereToolJSON renders one tool entry exactly as the template's tojson
+// filter does: {"name": ..., "description": ..., "parameters": {...},
+// "responses": null} with ", " / ": " separators.
+func cohereToolJSON(tool api.Tool) (string, error) {
+	params, err := marshalWithSpaces(tool.Function.Parameters)
+	if err != nil {
+		return "", err
+	}
+	name, err := json.Marshal(tool.Function.Name)
+	if err != nil {
+		return "", err
+	}
+	desc, err := json.Marshal(tool.Function.Description)
+	if err != nil {
+		return "", err
+	}
+	var sb strings.Builder
+	sb.WriteString(`{"name": `)
+	sb.Write(name)
+	sb.WriteString(`, "description": `)
+	sb.Write(desc)
+	sb.WriteString(`, "parameters": `)
+	sb.WriteString(string(params))
+	sb.WriteString(`, "responses": null}`)
+	return sb.String(), nil
+}
+
+// writeToolsSection writes the "# Available Tools" block, reproducing the
+// template's whitespace for the empty and populated cases.
+func writeToolsSection(sb *strings.Builder, tools []api.Tool) error {
+	sb.WriteString("# Available Tools\n```json\n[\n")
+	if len(tools) == 0 {
+		sb.WriteString("\n\n")
+	} else {
+		for i, tool := range tools {
+			entry, err := cohereToolJSON(tool)
+			if err != nil {
+				return err
+			}
+			sb.WriteString("\n    ")
+			sb.WriteString(entry)
+			if i < len(tools)-1 {
+				sb.WriteString(",")
+			}
+			sb.WriteString("\n\n")
+		}
+	}
+	sb.WriteString("\n]\n```")
+	return nil
+}
+
+// cohereToolResult is one entry of a <|START_TOOL_RESULT|> array.
+func writeToolResult(sb *strings.Builder, callID string, content string) error {
+	wrapped, err := marshalWithSpaces(map[string]string{"content": content})
+	if err != nil {
+		return err
+	}
+	sb.WriteString("\n    {\n        \"tool_call_id\": \"")
+	sb.WriteString(callID)
+	sb.WriteString("\",\n        \"results\": {\n\n            \n            \"0\": ")
+	sb.WriteString(string(wrapped))
+	sb.WriteString("\n\n        },\n        \"is_error\": null\n    }")
+	return nil
+}
+
+func (r *CohereRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) {
+	var sb strings.Builder
+
+	// The template defaults reasoning to true; an explicit think=false
+	// disables it.
+	reasoning := think == nil || think.Bool()
+
+	// The first system message — the request's system prompt, or the model's
+	// Modelfile SYSTEM when the request has none — fills the template's
+	// platform instruction slot.
+	var system string
+	rest := messages
+	if len(messages) > 0 && strings.EqualFold(messages[0].Role, "system") {
+		system = messages[0].Content
+		rest = messages[1:]
+	}
+
+	// Platform system turn: system prompt plus the Available Tools section
+	// (rendered even when no tools are defined).
+	sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>")
+	if system != "" {
+		sb.WriteString(system)
+		sb.WriteString("\n\n\n\n")
+	}
+	if err := writeToolsSection(&sb, tools); err != nil {
+		return "", err
+	}
+	sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>")
+
+	// Tool call ids regenerate as sequential indices across the whole
+	// conversation (regen_tool_call_ids default); results reference the
+	// index of their originating call.
+	callIndex := 0
+	callIDToIndex := map[string]string{}
+	nextCallID := func(id string) string {
+		idx := strconv.Itoa(callIndex)
+		callIndex++
+		if id != "" {
+			if _, seen := callIDToIndex[id]; !seen {
+				callIDToIndex[id] = idx
+			}
+		}
+		return idx
+	}
+	resolveResultID := func(m api.Message) string {
+		if idx, ok := callIDToIndex[m.ToolCallID]; ok {
+			return idx
+		}
+		// Fall back to call order when ids are absent.
+		return m.ToolCallID
+	}
+
+	prefill := false
+	for i := 0; i < len(rest); i++ {
+		message := rest[i]
+		switch strings.ToLower(message.Role) {
+		case "system":
+			sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>")
+			sb.WriteString(message.Content)
+			sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>")
+		case "user":
+			sb.WriteString("<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>")
+			sb.WriteString(message.Content)
+			sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>")
+		case "assistant", "chatbot":
+			sb.WriteString("<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>")
+			if len(message.ToolCalls) > 0 {
+				// Whitespace before THINKING/ACTION matches the template's
+				// (untrimmed jinja blocks).
+				sb.WriteString("\n            \n                ")
+				if message.Thinking != "" {
+					sb.WriteString("<|START_THINKING|>")
+					sb.WriteString(message.Thinking)
+					sb.WriteString("<|END_THINKING|>")
+				}
+				sb.WriteString("<|START_ACTION|>[")
+				for j, tc := range message.ToolCalls {
+					args, err := marshalWithSpaces(tc.Function.Arguments)
+					if err != nil {
+						return "", err
+					}
+					sb.WriteString("\n\n    {\"tool_call_id\": \"")
+					sb.WriteString(nextCallID(toolCallID(tc)))
+					sb.WriteString("\", \"tool_name\": \"")
+					sb.WriteString(tc.Function.Name)
+					sb.WriteString("\", \"parameters\": ")
+					sb.WriteString(string(args))
+					sb.WriteString("}")
+					if j < len(message.ToolCalls)-1 {
+						sb.WriteString(",")
+					}
+				}
+				sb.WriteString("\n\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>")
+			} else {
+				if message.Thinking != "" {
+					sb.WriteString("<|START_THINKING|>")
+					sb.WriteString(message.Thinking)
+					sb.WriteString("<|END_THINKING|>")
+				}
+				sb.WriteString("<|START_TEXT|>")
+				sb.WriteString(message.Content)
+				if i == len(rest)-1 {
+					// Assistant prefill: leave the text open for
+					// continuation.
+					prefill = true
+				} else {
+					sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>")
+				}
+			}
+		case "tool":
+			// Consecutive tool messages merge into one TOOL_RESULT array.
+			sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[")
+			if err := writeToolResult(&sb, resolveResultID(message), message.Content); err != nil {
+				return "", err
+			}
+			for i+1 < len(rest) && strings.EqualFold(rest[i+1].Role, "tool") {
+				i++
+				sb.WriteString(",")
+				if err := writeToolResult(&sb, resolveResultID(rest[i]), rest[i].Content); err != nil {
+					return "", err
+				}
+			}
+			sb.WriteString("\n\n]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>")
+		}
+	}
+
+	if !prefill {
+		sb.WriteString("<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>")
+		if reasoning {
+			sb.WriteString("<|START_THINKING|>")
+		} else {
+			sb.WriteString("<|START_THINKING|><|END_THINKING|>")
+		}
+	}
+
+	return sb.String(), nil
+}
+
+// toolCallID returns the tool call's id when the client supplied one. The
+// api.ToolCall ID field may be empty for calls synthesized by ollama.
+func toolCallID(tc api.ToolCall) string {
+	return tc.ID
+}

model/renderers/cohere_test.go

@@ -0,0 +1,189 @@
+package renderers
+
+import (
+	"testing"
+
+	"github.com/ollama/ollama/api"
+)
+
+// Ground truth in these tests comes from rendering North-Mini-Code-1.0's
+// chat_template.jinja with HF jinja semantics (add_generation_prompt=true),
+// minus the leading "<BOS>" from {{ bos_token }} (the tokenizer adds BOS as a
+// token at encode time).
+
+const cohereSystemTurnNoTools = "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>" +
+	"# Available Tools\n```json\n[\n\n\n\n]\n```" +
+	"<|END_TEXT|><|END_OF_TURN_TOKEN|>"
+
+func TestCohereRenderUserOnly(t *testing.T) {
+	r := &CohereRenderer{}
+	got, err := r.Render([]api.Message{{Role: "user", Content: "USERMSG"}}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := cohereSystemTurnNoTools +
+		"<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>USERMSG<|END_TEXT|><|END_OF_TURN_TOKEN|>" +
+		"<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>"
+	if got != want {
+		t.Errorf("render mismatch:\ngot:  %q\nwant: %q", got, want)
+	}
+}
+
+func TestCohereRenderSystemHistoryAndThinking(t *testing.T) {
+	r := &CohereRenderer{}
+	got, err := r.Render([]api.Message{
+		{Role: "system", Content: "DEVPREAMBLE"},
+		{Role: "user", Content: "Q1"},
+		{Role: "assistant", Content: "A1", Thinking: "THINK1"},
+		{Role: "user", Content: "Q2"},
+	}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>" +
+		"DEVPREAMBLE\n\n\n\n# Available Tools\n```json\n[\n\n\n\n]\n```" +
+		"<|END_TEXT|><|END_OF_TURN_TOKEN|>" +
+		"<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>Q1<|END_TEXT|><|END_OF_TURN_TOKEN|>" +
+		"<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>THINK1<|END_THINKING|><|START_TEXT|>A1<|END_TEXT|><|END_OF_TURN_TOKEN|>" +
+		"<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>Q2<|END_TEXT|><|END_OF_TURN_TOKEN|>" +
+		"<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>"
+	if got != want {
+		t.Errorf("render mismatch:\ngot:  %q\nwant: %q", got, want)
+	}
+}
+
+func TestCohereRenderReasoningOff(t *testing.T) {
+	r := &CohereRenderer{}
+	think := &api.ThinkValue{Value: false}
+	got, err := r.Render([]api.Message{{Role: "user", Content: "Q"}}, nil, think)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wantSuffix := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|><|END_THINKING|>"
+	if got[len(got)-len(wantSuffix):] != wantSuffix {
+		t.Errorf("reasoning-off generation prompt mismatch, got tail %q", got[len(got)-len(wantSuffix):])
+	}
+}
+
+func TestCohereRenderToolFlow(t *testing.T) {
+	r := &CohereRenderer{}
+	tools := []api.Tool{{
+		Type: "function",
+		Function: api.ToolFunction{
+			Name:        "get_weather",
+			Description: "Get weather",
+			Parameters: api.ToolFunctionParameters{
+				Type:     "object",
+				Required: []string{"city"},
+				Properties: testPropsOrdered([]orderedProp{
+					{Key: "city", Value: api.ToolProperty{Type: api.PropertyType{"string"}}},
+				}),
+			},
+		},
+	}}
+	args := api.ToolCallFunctionArguments{}
+	args.Set("city", "Paris")
+	got, err := r.Render([]api.Message{
+		{Role: "user", Content: "weather in Paris?"},
+		{Role: "assistant", Thinking: "I should call the tool", ToolCalls: []api.ToolCall{
+			{ID: "call_x", Function: api.ToolCallFunction{Name: "get_weather", Arguments: args}},
+		}},
+		{Role: "tool", ToolCallID: "call_x", Content: "15C sunny"},
+		{Role: "user", Content: "thanks"},
+	}, tools, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Tools section (from jinja): one entry, surrounded by the template's
+	// whitespace. Key order inside "parameters" follows Go's struct order
+	// (jinja preserves whatever order the client sent; both are valid JSON
+	// schema).
+	wantTools := "# Available Tools\n```json\n[\n\n    {\"name\": \"get_weather\", \"description\": \"Get weather\", \"parameters\": {\"type\": \"object\", \"required\": [\"city\"], \"properties\": {\"city\": {\"type\": \"string\"}}}, \"responses\": null}\n\n\n]\n```"
+	if !contains(t, got, wantTools, "tools section") {
+		return
+	}
+
+	// Assistant tool call turn.
+	wantAction := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>\n            \n                <|START_THINKING|>I should call the tool<|END_THINKING|><|START_ACTION|>[\n\n    {\"tool_call_id\": \"0\", \"tool_name\": \"get_weather\", \"parameters\": {\"city\": \"Paris\"}}\n\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>"
+	if !contains(t, got, wantAction, "action turn") {
+		return
+	}
+
+	// Tool result turn.
+	wantResult := "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[\n    {\n        \"tool_call_id\": \"0\",\n        \"results\": {\n\n            \n            \"0\": {\"content\": \"15C sunny\"}\n\n        },\n        \"is_error\": null\n    }\n\n]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>"
+	contains(t, got, wantResult, "tool result turn")
+}
+
+func TestCohereRenderMultiToolCallsAndResults(t *testing.T) {
+	r := &CohereRenderer{}
+	emptyArgs := api.ToolCallFunctionArguments{}
+	xArgs := api.ToolCallFunctionArguments{}
+	xArgs.Set("x", 1)
+	got, err := r.Render([]api.Message{
+		{Role: "user", Content: "go"},
+		{Role: "assistant", ToolCalls: []api.ToolCall{
+			{ID: "a", Function: api.ToolCallFunction{Name: "t1", Arguments: emptyArgs}},
+			{ID: "b", Function: api.ToolCallFunction{Name: "t2", Arguments: xArgs}},
+		}},
+		{Role: "tool", ToolCallID: "a", Content: "r1"},
+		{Role: "tool", ToolCallID: "b", Content: "r2"},
+	}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wantAction := "<|START_ACTION|>[\n\n    {\"tool_call_id\": \"0\", \"tool_name\": \"t1\", \"parameters\": {}},\n\n    {\"tool_call_id\": \"1\", \"tool_name\": \"t2\", \"parameters\": {\"x\": 1}}\n\n]<|END_ACTION|>"
+	if !contains(t, got, wantAction, "multi action") {
+		return
+	}
+
+	// Consecutive tool messages merge into one result block with sequential
+	// regenerated ids.
+	wantResults := "<|START_TOOL_RESULT|>[\n    {\n        \"tool_call_id\": \"0\",\n        \"results\": {\n\n            \n            \"0\": {\"content\": \"r1\"}\n\n        },\n        \"is_error\": null\n    },\n    {\n        \"tool_call_id\": \"1\",\n        \"results\": {\n\n            \n            \"0\": {\"content\": \"r2\"}\n\n        },\n        \"is_error\": null\n    }\n\n]<|END_TOOL_RESULT|>"
+	contains(t, got, wantResults, "merged tool results")
+}
+
+func TestCohereRenderAssistantPrefill(t *testing.T) {
+	r := &CohereRenderer{}
+	got, err := r.Render([]api.Message{
+		{Role: "user", Content: "Q"},
+		{Role: "assistant", Content: "partial"},
+	}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wantSuffix := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_TEXT|>partial"
+	if got[len(got)-len(wantSuffix):] != wantSuffix {
+		t.Errorf("prefill should leave the text open, got tail %q", got[max(0, len(got)-90):])
+	}
+}
+
+func TestCohereRendererRegistered(t *testing.T) {
+	if rendererForName("cohere") == nil {
+		t.Fatal("cohere renderer not registered")
+	}
+	if got := LeadingBOSForRenderer("cohere"); got != "<BOS_TOKEN>" {
+		t.Errorf("LeadingBOS = %q, want <BOS_TOKEN>", got)
+	}
+}
+
+func contains(t *testing.T, haystack, needle, what string) bool {
+	t.Helper()
+	if idx := indexOf(haystack, needle); idx == -1 {
+		t.Errorf("missing %s:\nwant substring: %q\nin: %q", what, needle, haystack)
+		return false
+	}
+	return true
+}
+
+func indexOf(s, sub string) int {
+	for i := 0; i+len(sub) <= len(s); i++ {
+		if s[i:i+len(sub)] == sub {
+			return i
+		}
+	}
+	return -1
+}

model/renderers/lfm2.go

@@ -17,6 +17,8 @@ type LFM2Renderer struct {
 const lfm2BOSToken = "<|startoftext|>"
 
 const (
+	lfm2ThinkingOpenTag      = "<think>"
+	lfm2ThinkingCloseTag     = "</think>"
 	lfm2ToolListStartTag     = "<|tool_list_start|>"
 	lfm2ToolListEndTag       = "<|tool_list_end|>"
 	lfm2ToolCallStartTag     = "<|tool_call_start|>"
@@ -285,11 +287,6 @@ func (r *LFM2Renderer) Render(messages []api.Message, tools []api.Tool, thinkVal
 
 		content := r.renderMessageContent(message, imageOffset)
 		imageOffset += len(message.Images)
-		if message.Role == "assistant" && !keepPastThinking && i != lastAssistantIndex {
-			if idx := strings.LastIndex(content, "</think>"); idx >= 0 {
-				content = strings.TrimSpace(content[idx+len("</think>"):])
-			}
-		}
 		if message.Role == "assistant" && len(message.ToolCalls) > 0 && !strings.Contains(content, lfm2ToolCallStartTag) {
 			if strings.TrimSpace(content) == "" {
 				content = lfm2RenderToolCalls(message.ToolCalls) + content
@@ -297,6 +294,23 @@ func (r *LFM2Renderer) Render(messages []api.Message, tools []api.Tool, thinkVal
 				content = lfm2RenderToolCalls(message.ToolCalls) + "\n" + content
 			}
 		}
+		// Reconstruct the inline <think>...</think> block from the separate
+		// Thinking field so reasoning turns round-trip in the model's own format:
+		// thinking precedes any tool calls and content. A direct answer carries no
+		// Thinking, so nothing is added. Only the thinking variant emits these tags;
+		// the non-thinking renderer must never send them, including for a trailing
+		// assistant prefill (which is exempt from the stripping below).
+		if r.IsThinking && message.Role == "assistant" && message.Thinking != "" && !strings.Contains(content, lfm2ThinkingCloseTag) {
+			content = lfm2ThinkingOpenTag + message.Thinking + lfm2ThinkingCloseTag + content
+		}
+		// Drop reasoning from earlier assistant turns unless thinking is kept; the
+		// <think>...</think> block is a clean prefix, so everything after the close
+		// tag (tool calls and content) is preserved.
+		if message.Role == "assistant" && !keepPastThinking && i != lastAssistantIndex {
+			if idx := strings.LastIndex(content, lfm2ThinkingCloseTag); idx >= 0 {
+				content = strings.TrimSpace(content[idx+len(lfm2ThinkingCloseTag):])
+			}
+		}
 		if message.Role == "tool" && !strings.Contains(content, lfm2ToolResponseStartTag) {
 			content = lfm2ToolResponseStartTag + content + lfm2ToolResponseEndTag
 		}

model/renderers/lfm2_test.go

@@ -159,6 +159,78 @@ func TestLFM2Renderer_ChatTemplateParity(t *testing.T) {
 			thinkValue: &api.ThinkValue{Value: true},
 			expected:   "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\n<think>reason1</think>A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
 		},
+		{
+			name:     "thinking_field_reconstructed_when_enabled",
+			renderer: &LFM2Renderer{IsThinking: true},
+			messages: []api.Message{
+				{Role: "user", Content: "Q1"},
+				{Role: "assistant", Content: "A1", Thinking: "reason1"},
+				{Role: "user", Content: "Q2"},
+				{Role: "assistant", Content: "A2", Thinking: "reason2"},
+			},
+			thinkValue: &api.ThinkValue{Value: true},
+			expected:   "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\n<think>reason1</think>A1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
+		},
+		{
+			name:     "thinking_field_stripped_for_non_last_when_disabled",
+			renderer: &LFM2Renderer{IsThinking: true},
+			messages: []api.Message{
+				{Role: "user", Content: "Q1"},
+				{Role: "assistant", Content: "A1", Thinking: "reason1"},
+				{Role: "user", Content: "Q2"},
+				{Role: "assistant", Content: "A2", Thinking: "reason2"},
+			},
+			thinkValue: &api.ThinkValue{Value: false},
+			expected:   "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\nA1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n<think>reason2</think>A2",
+		},
+		{
+			name:     "thinking_precedes_tool_calls",
+			renderer: &LFM2Renderer{IsThinking: true},
+			messages: []api.Message{
+				{Role: "user", Content: "Weather?"},
+				{
+					Role:     "assistant",
+					Content:  "",
+					Thinking: "Let me check the weather.",
+					ToolCalls: []api.ToolCall{
+						{
+							Function: api.ToolCallFunction{
+								Name: "get_weather",
+								Arguments: testArgs(map[string]any{
+									"location": "Paris",
+								}),
+							},
+						},
+					},
+				},
+				{Role: "tool", Content: "22C"},
+			},
+			thinkValue: &api.ThinkValue{Value: true},
+			expected:   "<|startoftext|><|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think>Let me check the weather.</think><|tool_call_start|>[get_weather(location=\"Paris\")]<|tool_call_end|><|im_end|>\n<|im_start|>tool\n<|tool_response_start|>22C<|tool_response_end|><|im_end|>\n<|im_start|>assistant\n",
+		},
+		{
+			name:     "direct_answer_history_has_no_think_tags",
+			renderer: &LFM2Renderer{IsThinking: true},
+			messages: []api.Message{
+				{Role: "user", Content: "Q1"},
+				{Role: "assistant", Content: "A1"},
+				{Role: "user", Content: "Q2"},
+			},
+			thinkValue: &api.ThinkValue{Value: true},
+			expected:   "<|startoftext|><|im_start|>user\nQ1<|im_end|>\n<|im_start|>assistant\nA1<|im_end|>\n<|im_start|>user\nQ2<|im_end|>\n<|im_start|>assistant\n",
+		},
+		{
+			// The non-thinking variant must never emit <think> tags, even for a
+			// trailing assistant prefill (which is exempt from thinking stripping).
+			name:     "non_thinking_renderer_drops_thinking_metadata_on_prefill",
+			renderer: &LFM2Renderer{IsThinking: false},
+			messages: []api.Message{
+				{Role: "user", Content: "Hi"},
+				{Role: "assistant", Content: "Hello", Thinking: "some reasoning"},
+			},
+			thinkValue: &api.ThinkValue{Value: false},
+			expected:   "<|startoftext|><|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello",
+		},
 		{
 			name:     "arbitrary_roles_are_rendered_verbatim",
 			renderer: &LFM2Renderer{IsThinking: false},

model/renderers/renderer.go

@@ -107,6 +107,8 @@ func rendererForName(name string) Renderer {
 		return &LFM2Renderer{IsThinking: true, useImgTags: RenderImgTags}
 	case "laguna":
 		return &LagunaRenderer{}
+	case "cohere":
+		return &CohereRenderer{}
 	default:
 		return nil
 	}

server/sched.go

@@ -132,14 +132,24 @@ func (s *Scheduler) GetRunner(c context.Context, m *Model, opts api.Options, ses
 	return s.getRunner(c, m, opts, sessionDuration, false, false, nil)
 }
 
-const contextShiftSmallContextLimit = 8192
-
-func resolveContextShift(shift *bool, numCtx int) bool {
+func resolveContextShift(shift *bool, m *Model) bool {
 	if shift != nil {
 		return *shift
 	}
 
-	return numCtx > 0 && numCtx < contextShiftSmallContextLimit
+	return supportsContextShift(m)
+}
+
+func supportsContextShift(m *Model) bool {
+	if m == nil {
+		return true
+	}
+
+	if m.Config.ModelFamily == "deepseek2" || slices.Contains(m.Config.ModelFamilies, "deepseek2") {
+		return false
+	}
+
+	return true
 }
 
 func effectiveModelContext(numCtx int, f *ggml.GGML) int {
@@ -174,7 +184,7 @@ func (s *Scheduler) getRunner(c context.Context, m *Model, opts api.Options, ses
 
 	contextShift := false
 	if m.ModelPath != "" {
-		contextShift = resolveContextShift(shift, opts.NumCtx)
+		contextShift = resolveContextShift(shift, m)
 	}
 
 	req := &LlmRequest{
@@ -566,7 +576,7 @@ func (s *Scheduler) load(req *LlmRequest, systemInfo ml.SystemInfo, gpus []ml.De
 			}
 
 			launchOpts = s.applyLlamaServerMmapDefaults(req, launchOpts, systemInfo, loadGpus, f, numParallel)
-			req.contextShift = resolveContextShift(req.shift, effectiveModelContext(launchOpts.NumCtx, f))
+			req.contextShift = resolveContextShift(req.shift, req.model)
 
 			config := llamaServerConfigForModel(req.model)
 			config.ContextShift = req.contextShift
@@ -691,7 +701,7 @@ iGPUScan:
 	trainContext := modelTrainContext(f)
 	if effectiveNumCtx := llama.ContextLength(); req.model.ModelPath != "" && effectiveNumCtx > 0 {
 		req.opts.NumCtx = effectiveNumCtx
-		req.contextShift = resolveContextShift(req.shift, effectiveNumCtx)
+		req.contextShift = resolveContextShift(req.shift, req.model)
 	}
 	runner := &runnerRef{
 		model:           req.model,
@@ -1420,7 +1430,7 @@ func (runner *runnerRef) needsReload(ctx context.Context, req *LlmRequest) bool
 
 	contextShift := req.contextShift
 	if req.model.ModelPath != "" {
-		contextShift = resolveContextShift(req.shift, optsNew.NumCtx)
+		contextShift = resolveContextShift(req.shift, req.model)
 	}
 	if runner.contextShift != contextShift {
 		return true

server/sched_test.go

@@ -907,19 +907,21 @@ func TestResolveContextShift(t *testing.T) {
 	tests := []struct {
 		name  string
 		shift *bool
-		ctx   int
+		model *Model
 		want  bool
 	}{
-		{name: "unset small context keeps legacy shift", ctx: 128, want: true},
-		{name: "unset large context disables shift", ctx: contextShiftSmallContextLimit, want: false},
-		{name: "unset invalid context disables shift", ctx: 0, want: false},
-		{name: "explicit false wins for small context", shift: &falseValue, ctx: 128, want: false},
-		{name: "explicit true wins for large context", shift: &trueValue, ctx: 32768, want: true},
+		{name: "unset defaults to shift", want: true},
+		{name: "unset deepseek2 disables shift", model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false},
+		{name: "unset deepseek2 family disables shift", model: &Model{Config: model.ConfigV2{ModelFamilies: []string{"llama", "deepseek2"}}}, want: false},
+		{name: "explicit false disables shift", shift: &falseValue, want: false},
+		{name: "explicit false disables shift for deepseek2", shift: &falseValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false},
+		{name: "explicit true enables shift", shift: &trueValue, want: true},
+		{name: "explicit true enables shift for deepseek2", shift: &trueValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: true},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.ctx))
+			require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.model))
 		})
 	}
 }

x/create/client/create.go

@@ -745,6 +745,9 @@ func getParserName(modelDir string) string {
 		if strings.Contains(archLower, "laguna") {
 			return "laguna"
 		}
+		if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") {
+			return "cohere"
+		}
 		if strings.Contains(archLower, "glm4") || strings.Contains(archLower, "glm-4") {
 			return "glm-4.7"
 		}
@@ -765,6 +768,9 @@ func getParserName(modelDir string) string {
 		if strings.Contains(typeLower, "laguna") {
 			return "laguna"
 		}
+		if strings.Contains(typeLower, "cohere2_moe") {
+			return "cohere"
+		}
 		if strings.Contains(typeLower, "glm4") || strings.Contains(typeLower, "glm-4") {
 			return "glm-4.7"
 		}
@@ -805,6 +811,9 @@ func getRendererName(modelDir string) string {
 		if strings.Contains(archLower, "laguna") {
 			return "laguna"
 		}
+		if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") {
+			return "cohere"
+		}
 		if strings.Contains(archLower, "gemma4") {
 			return "gemma4"
 		}
@@ -825,6 +834,9 @@ func getRendererName(modelDir string) string {
 		if strings.Contains(typeLower, "laguna") {
 			return "laguna"
 		}
+		if strings.Contains(typeLower, "cohere2_moe") {
+			return "cohere"
+		}
 		if strings.Contains(typeLower, "gemma4") {
 			return "gemma4"
 		}

x/create/cohere2moe.go

@@ -0,0 +1,106 @@
+package create
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/ollama/ollama/x/safetensors"
+)
+
+// cohere2MoeImportTransform adjusts quantization for Cohere2 MoE imports
+// (Command A family / North models).
+type cohere2MoeImportTransform struct {
+	numLayers int
+}
+
+func newCohere2MoeImportTransform(modelDir string, _ sourceModelConfig) (tensorImportTransform, error) {
+	data, err := os.ReadFile(filepath.Join(modelDir, "config.json"))
+	if err != nil {
+		return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic
+	}
+	var cfg struct {
+		NumHiddenLayers int `json:"num_hidden_layers"`
+	}
+	if err := json.Unmarshal(data, &cfg); err != nil {
+		return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic
+	}
+	return cohere2MoeImportTransform{numLayers: cfg.NumHiddenLayers}, nil
+}
+
+func (cohere2MoeImportTransform) skipTensor(string) bool { return false }
+
+func (cohere2MoeImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) {
+	return []*safetensors.TensorData{td}, nil
+}
+
+var cohere2MoeLayerIndexRe = regexp.MustCompile(`\.layers\.(\d+)\.`)
+
+func (t cohere2MoeImportTransform) quantizationType(name string, shape []int32, quantize string) string {
+	quantNorm := normalizeQuantType(quantize)
+
+	// The embedding serves double duty: lookup (via QuantizedEmbedding) and the
+	// tied lm_head projection (via AsLinear). With a 262k vocab the bf16
+	// embedding dominates decode bandwidth through the lm_head matmul, so
+	// quantize it to the 8-bit variant of the requested mode.
+	if strings.HasSuffix(name, "embed_tokens.weight") && len(shape) == 2 {
+		switch quantNorm {
+		case "int4", "int8":
+			if isAligned(shape, "int8") {
+				return "int8"
+			}
+		case "mxfp4", "nvfp4", "mxfp8":
+			if isAligned(shape, "mxfp8") {
+				return "mxfp8"
+			}
+		}
+		return ""
+	}
+
+	// The MoE router picks the top-k expert set; quantization noise there can
+	// flip expert selection and compound downstream. It is tiny, so keep it in
+	// source precision. (GetTensorQuantization already skips "mlp.gate.weight";
+	// kept explicit here so renames in the default policy cannot regress this.)
+	if strings.HasSuffix(name, ".mlp.gate.weight") {
+		return ""
+	}
+
+	// Sensitive tensors (v_proj, k_proj, down_proj) get higher precision only
+	// at quantization-sensitive layer positions (gemma4's useMoreBits
+	// heuristic) instead of the default policy's blanket promotion. The
+	// blanket int8 down_proj costs ~25% of decode bandwidth on a top-8 MoE;
+	// the layer-position heuristic keeps the early/late layers (and every
+	// third in between) at 8 bits where residual-stream error matters most.
+	promote := ""
+	switch quantNorm {
+	case "int4":
+		promote = "int8"
+	case "mxfp4", "nvfp4":
+		promote = "mxfp8"
+	}
+	isSensitive := strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj") || strings.Contains(name, "down_proj")
+	if promote != "" && isSensitive && t.numLayers > 0 {
+		layerIdx := -1
+		if m := cohere2MoeLayerIndexRe.FindStringSubmatch(name); m != nil {
+			if idx, err := strconv.Atoi(m[1]); err == nil {
+				layerIdx = idx
+			}
+		}
+		if layerIdx >= 0 {
+			if useMoreBits(layerIdx, t.numLayers) && isAligned(shape, promote) {
+				return promote
+			}
+			if !isAligned(shape, quantNorm) {
+				return ""
+			}
+			// Bypass GetTensorQuantization's blanket promotion — the
+			// layer-position heuristic is authoritative here.
+			return quantNorm
+		}
+	}
+
+	return GetTensorQuantization(name, shape, quantize)
+}

x/create/create.go

@@ -821,6 +821,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{
 	"gemma4_unified":                        newGemma4ImportTransform,
 	"gemma4_unified_text":                   newGemma4ImportTransform,
 	"LagunaForCausalLM":                     newLagunaImportTransform,
+	"Cohere2MoeForCausalLM":                 newCohere2MoeImportTransform,
 	"Gemma4AssistantForCausalLM":            newGemma4ImportTransform,
 	"Gemma4UnifiedAssistantForCausalLM":     newGemma4ImportTransform,
 	"gemma4_unified_assistant":              newGemma4ImportTransform,

x/mlxrunner/imports.go

@@ -1,6 +1,7 @@
 package mlxrunner
 
 import (
+	_ "github.com/ollama/ollama/x/models/cohere2_moe"
 	_ "github.com/ollama/ollama/x/models/gemma3"
 	_ "github.com/ollama/ollama/x/models/gemma4"
 	_ "github.com/ollama/ollama/x/models/glm4_moe_lite"

x/models/cohere2_moe/cohere2_moe.go

@@ -0,0 +1,770 @@
+// Package cohere2_moe provides the Cohere2 MoE (Command A family, North) text
+// model implementation for MLX.
+//
+// Architecture notes (matches transformers' Cohere2MoeForCausalLM):
+//   - Parallel residual blocks: a single input layernorm feeds both attention
+//     and the MLP, and their outputs are summed onto the residual.
+//   - Interleaved sliding-window and full attention layers. Sliding layers use
+//     interleaved ("traditional") RoPE; full-attention layers use no positional
+//     encoding (NoPE), except prefix dense layers when
+//     prefix_dense_sliding_window_pattern == 1, which force RoPE.
+//   - The first first_k_dense_replace layers use a dense SwiGLU MLP with
+//     prefix_dense_intermediate_size; the rest are sparse MoE layers routed by
+//     a linear gate with sigmoid or softmax selection over the top-k logits.
+//   - Logits are scaled by logit_scale. Embeddings are tied by default.
+package cohere2_moe
+
+import (
+	"encoding/json"
+	"fmt"
+	"math"
+
+	"github.com/ollama/ollama/x/mlxrunner/batch"
+	"github.com/ollama/ollama/x/mlxrunner/cache"
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+	"github.com/ollama/ollama/x/mlxrunner/model"
+	"github.com/ollama/ollama/x/mlxrunner/model/base"
+	"github.com/ollama/ollama/x/models/nn"
+	"github.com/ollama/ollama/x/tokenizer"
+)
+
+func init() {
+	base.Register("Cohere2MoeForCausalLM", NewModel)
+}
+
+// Config holds the Cohere2 MoE configuration (HuggingFace config.json).
+type Config struct {
+	HiddenSize            int32    `json:"hidden_size"`
+	NumHiddenLayers       int32    `json:"num_hidden_layers"`
+	IntermediateSize      int32    `json:"intermediate_size"`
+	NumAttentionHeads     int32    `json:"num_attention_heads"`
+	NumKeyValueHeads      int32    `json:"num_key_value_heads"`
+	HeadDim               int32    `json:"head_dim"`
+	VocabSize             int32    `json:"vocab_size"`
+	MaxPositionEmbeddings int32    `json:"max_position_embeddings"`
+	LayerNormEps          float32  `json:"layer_norm_eps"`
+	RMSNormEps            *float32 `json:"rms_norm_eps"`
+	RopeTheta             float32  `json:"rope_theta"`
+	LogitScale            float32  `json:"logit_scale"`
+	AttentionBias         bool     `json:"attention_bias"`
+	TieWordEmbeddings     *bool    `json:"tie_word_embeddings"`
+
+	SlidingWindow                   int32    `json:"sliding_window"`
+	SlidingWindowPattern            int32    `json:"sliding_window_pattern"`
+	PrefixDenseSlidingWindowPattern int32    `json:"prefix_dense_sliding_window_pattern"`
+	LayerTypes                      []string `json:"layer_types"`
+	MLPLayerTypes                   []string `json:"mlp_layer_types"`
+	FirstKDenseReplace              int32    `json:"first_k_dense_replace"`
+	PrefixDenseIntermediateSize     int32    `json:"prefix_dense_intermediate_size"`
+
+	NumExperts                      int32  `json:"num_experts"`
+	NumExpertsPerTok                int32  `json:"num_experts_per_tok"`
+	NumSharedExperts                int32  `json:"num_shared_experts"`
+	SharedExpertCombinationStrategy string `json:"shared_expert_combination_strategy"`
+	ExpertSelectionFn               string `json:"expert_selection_fn"`
+	NormTopKProb                    bool   `json:"norm_topk_prob"`
+
+	// Quantization metadata (set at load, not from config.json).
+	QuantGroupSize int                               `json:"-"`
+	QuantBits      int                               `json:"-"`
+	QuantMode      string                            `json:"-"`
+	TensorQuant    map[string]*model.TensorQuantInfo `json:"-"`
+
+	// Computed fields.
+	Scale float32 `json:"-"`
+}
+
+// normLayer abstracts the per-config choice between RMSNorm (rms_norm_eps set)
+// and Cohere-style bias-free LayerNorm.
+type normLayer interface {
+	Forward(x *mlx.Array) *mlx.Array
+}
+
+type rmsNorm struct {
+	Weight *mlx.Array
+	Eps    float32
+}
+
+func (n *rmsNorm) Forward(x *mlx.Array) *mlx.Array { return mlx.RMSNormFn(x, n.Weight, n.Eps) }
+
+type layerNorm struct {
+	Weight *mlx.Array
+	Eps    float32
+}
+
+func (n *layerNorm) Forward(x *mlx.Array) *mlx.Array {
+	return mlx.LayerNormFn(x, n.Weight, nil, n.Eps)
+}
+
+// Model is the Cohere2 MoE model.
+type Model struct {
+	EmbedTokens nn.EmbeddingLayer
+	Layers      []*Layer
+	Norm        normLayer
+	LMHead      nn.LinearLayer
+
+	tok *tokenizer.Tokenizer
+	*Config
+}
+
+// Layer is a parallel-residual transformer block.
+type Layer struct {
+	InputNorm normLayer
+	Attention *Attention
+	MLP       MLPBlock
+
+	IsSliding bool
+	UseRope   bool
+}
+
+// Attention implements Cohere2 attention (no q/k norm).
+type Attention struct {
+	QProj nn.LinearLayer
+	KProj nn.LinearLayer
+	VProj nn.LinearLayer
+	OProj nn.LinearLayer
+}
+
+// MLPBlock is the feed-forward interface for dense and MoE blocks.
+type MLPBlock interface {
+	Forward(x *mlx.Array, cfg *Config) *mlx.Array
+}
+
+// DenseMLP is a SwiGLU feed-forward block.
+type DenseMLP struct {
+	GateProj nn.LinearLayer
+	UpProj   nn.LinearLayer
+	DownProj nn.LinearLayer
+}
+
+// SparseMoE routes each token to the top-k of NumExperts expert MLPs.
+type SparseMoE struct {
+	Router       nn.LinearLayer
+	SwitchMLP    *SwitchMLP
+	SharedExpert *DenseMLP
+}
+
+// SwitchMLP executes the selected expert MLPs with stacked expert weights.
+type SwitchMLP struct {
+	GateWeight *mlx.Array
+	UpWeight   *mlx.Array
+	DownWeight *mlx.Array
+
+	GateWeightQ, GateScales, GateBiases *mlx.Array
+	UpWeightQ, UpScales, UpBiases       *mlx.Array
+	DownWeightQ, DownScales, DownBiases *mlx.Array
+
+	GateBits, UpBits, DownBits                int
+	GateGroupSize, UpGroupSize, DownGroupSize int
+	GateMode, UpMode, DownMode                string
+
+	UseQuantized bool
+}
+
+type stackedExpertWeights struct {
+	Weight    *mlx.Array
+	Scales    *mlx.Array
+	Biases    *mlx.Array
+	Bits      int
+	GroupSize int
+	Mode      string
+}
+
+func parseConfig(configData []byte) (Config, error) {
+	var raw map[string]json.RawMessage
+	if err := json.Unmarshal(configData, &raw); err != nil {
+		return Config{}, fmt.Errorf("parse config envelope: %w", err)
+	}
+
+	var cfg Config
+	if err := json.Unmarshal(configData, &cfg); err != nil {
+		return Config{}, fmt.Errorf("parse config: %w", err)
+	}
+
+	if cfg.HiddenSize <= 0 {
+		return Config{}, fmt.Errorf("invalid hidden_size: %d", cfg.HiddenSize)
+	}
+	if cfg.NumHiddenLayers <= 0 {
+		return Config{}, fmt.Errorf("invalid num_hidden_layers: %d", cfg.NumHiddenLayers)
+	}
+	if cfg.NumAttentionHeads <= 0 {
+		return Config{}, fmt.Errorf("invalid num_attention_heads: %d", cfg.NumAttentionHeads)
+	}
+	if cfg.NumKeyValueHeads <= 0 {
+		cfg.NumKeyValueHeads = cfg.NumAttentionHeads
+	}
+	if cfg.HeadDim <= 0 {
+		if cfg.HiddenSize%cfg.NumAttentionHeads != 0 {
+			return Config{}, fmt.Errorf("hidden_size (%d) must be divisible by num_attention_heads (%d)", cfg.HiddenSize, cfg.NumAttentionHeads)
+		}
+		cfg.HeadDim = cfg.HiddenSize / cfg.NumAttentionHeads
+	}
+
+	// Defaults follow transformers' Cohere2MoeConfig.
+	if cfg.LayerNormEps == 0 {
+		cfg.LayerNormEps = 1e-5
+	}
+	if cfg.RopeTheta == 0 {
+		cfg.RopeTheta = 10000
+	}
+	if cfg.LogitScale == 0 {
+		cfg.LogitScale = 0.0625
+	}
+	if _, ok := raw["sliding_window"]; !ok {
+		cfg.SlidingWindow = 4096
+	}
+	if cfg.SlidingWindowPattern <= 0 {
+		cfg.SlidingWindowPattern = 4
+	}
+	if cfg.PrefixDenseSlidingWindowPattern <= 0 {
+		cfg.PrefixDenseSlidingWindowPattern = 1
+	}
+	if cfg.MaxPositionEmbeddings <= 0 {
+		cfg.MaxPositionEmbeddings = 8192
+	}
+
+	if cfg.NumExperts <= 0 {
+		cfg.NumExperts = 8
+	}
+	if cfg.NumExpertsPerTok <= 0 {
+		cfg.NumExpertsPerTok = 2
+	}
+	if cfg.NumExpertsPerTok > cfg.NumExperts {
+		return Config{}, fmt.Errorf("num_experts_per_tok (%d) exceeds num_experts (%d)", cfg.NumExpertsPerTok, cfg.NumExperts)
+	}
+	if cfg.ExpertSelectionFn == "" {
+		cfg.ExpertSelectionFn = "softmax"
+	}
+	if cfg.ExpertSelectionFn != "softmax" && cfg.ExpertSelectionFn != "sigmoid" {
+		return Config{}, fmt.Errorf("unsupported expert_selection_fn: %q", cfg.ExpertSelectionFn)
+	}
+	if cfg.SharedExpertCombinationStrategy == "" {
+		cfg.SharedExpertCombinationStrategy = "average"
+	}
+	if cfg.SharedExpertCombinationStrategy != "average" && cfg.SharedExpertCombinationStrategy != "sum" {
+		return Config{}, fmt.Errorf("unsupported shared_expert_combination_strategy: %q", cfg.SharedExpertCombinationStrategy)
+	}
+	if _, ok := raw["norm_topk_prob"]; !ok {
+		cfg.NormTopKProb = true
+	}
+	if cfg.PrefixDenseIntermediateSize <= 0 {
+		cfg.PrefixDenseIntermediateSize = cfg.IntermediateSize
+	}
+
+	// Derive per-layer attention types when absent: the first
+	// first_k_dense_replace layers follow prefix_dense_sliding_window_pattern,
+	// the rest follow sliding_window_pattern (full attention every Nth layer).
+	if len(cfg.LayerTypes) == 0 {
+		cfg.LayerTypes = make([]string, cfg.NumHiddenLayers)
+		for i := range cfg.NumHiddenLayers {
+			if i < cfg.FirstKDenseReplace {
+				cfg.LayerTypes[i] = patternLayerType(i, cfg.PrefixDenseSlidingWindowPattern)
+			} else {
+				cfg.LayerTypes[i] = patternLayerType(i-cfg.FirstKDenseReplace, cfg.SlidingWindowPattern)
+			}
+		}
+	}
+	if len(cfg.LayerTypes) != int(cfg.NumHiddenLayers) {
+		return Config{}, fmt.Errorf("layer_types has %d entries, want %d", len(cfg.LayerTypes), cfg.NumHiddenLayers)
+	}
+
+	// Derive per-layer MLP types when absent: the first first_k_dense_replace
+	// layers are dense, the rest sparse.
+	if len(cfg.MLPLayerTypes) == 0 {
+		cfg.MLPLayerTypes = make([]string, cfg.NumHiddenLayers)
+		for i := range cfg.NumHiddenLayers {
+			if i < cfg.FirstKDenseReplace {
+				cfg.MLPLayerTypes[i] = "dense"
+			} else {
+				cfg.MLPLayerTypes[i] = "sparse"
+			}
+		}
+	}
+	if len(cfg.MLPLayerTypes) != int(cfg.NumHiddenLayers) {
+		return Config{}, fmt.Errorf("mlp_layer_types has %d entries, want %d", len(cfg.MLPLayerTypes), cfg.NumHiddenLayers)
+	}
+
+	cfg.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim)))
+	return cfg, nil
+}
+
+func patternLayerType(i, pattern int32) string {
+	if pattern > 0 && (i+1)%pattern == 0 {
+		return "full_attention"
+	}
+	return "sliding_attention"
+}
+
+func (cfg *Config) layerIsSliding(i int32) bool {
+	return cfg.LayerTypes[i] == "sliding_attention"
+}
+
+func (cfg *Config) layerIsDense(i int32) bool {
+	return cfg.MLPLayerTypes[i] == "dense"
+}
+
+// layerUsesRope reports whether layer i applies rotary embeddings: all sliding
+// layers do, and prefix dense layers force RoPE even with full attention when
+// prefix_dense_sliding_window_pattern == 1 (matching Cohere2MoeAttention's
+// force_rope). Other full-attention layers use no positional encoding.
+func (cfg *Config) layerUsesRope(i int32) bool {
+	if cfg.layerIsSliding(i) {
+		return true
+	}
+	return cfg.layerIsDense(i) && cfg.PrefixDenseSlidingWindowPattern == 1
+}
+
+func (cfg *Config) newNorm(weight *mlx.Array) normLayer {
+	if cfg.RMSNormEps != nil {
+		return &rmsNorm{Weight: weight, Eps: *cfg.RMSNormEps}
+	}
+	return &layerNorm{Weight: weight, Eps: cfg.LayerNormEps}
+}
+
+func (cfg *Config) tieEmbeddings() bool {
+	return cfg.TieWordEmbeddings == nil || *cfg.TieWordEmbeddings
+}
+
+// NewModel creates a Cohere2 MoE model from a manifest root.
+func NewModel(root *model.Root) (base.Model, error) {
+	configData, err := root.Manifest.ReadConfig("config.json")
+	if err != nil {
+		return nil, fmt.Errorf("load config: %w", err)
+	}
+
+	cfg, err := parseConfig(configData)
+	if err != nil {
+		return nil, err
+	}
+
+	if qt := root.QuantType(); qt != "" {
+		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt)
+		if gs := root.GroupSize(); gs > 0 {
+			cfg.QuantGroupSize = gs
+		}
+	} else {
+		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams("")
+	}
+	cfg.TensorQuant = root.AllTensorQuant()
+
+	tokData, err := root.Manifest.ReadConfig("tokenizer.json")
+	if err != nil {
+		return nil, fmt.Errorf("load tokenizer config: %w", err)
+	}
+
+	tokConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData}
+	if genConfigData, err := root.Manifest.ReadConfig("generation_config.json"); err == nil {
+		tokConfig.GenerationConfigJSON = genConfigData
+	}
+	if tokConfigData, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil {
+		tokConfig.TokenizerConfigJSON = tokConfigData
+	}
+
+	tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig)
+	if err != nil {
+		return nil, fmt.Errorf("parse tokenizer: %w", err)
+	}
+
+	m := &Model{
+		Layers: make([]*Layer, cfg.NumHiddenLayers),
+		Config: &cfg,
+		tok:    tok,
+	}
+
+	for i := range cfg.NumHiddenLayers {
+		m.Layers[i] = &Layer{
+			IsSliding: cfg.layerIsSliding(i),
+			UseRope:   cfg.layerUsesRope(i),
+		}
+	}
+
+	return m, nil
+}
+
+func supportsGatherQMM(mode string, bits int) bool {
+	switch mode {
+	case "affine":
+		return bits == 4 || bits == 8
+	case "mxfp8":
+		return bits == 8
+	case "nvfp4", "mxfp4":
+		return bits == 4
+	default:
+		return false
+	}
+}
+
+// transposeExpertWeightForGatherMM converts stacked [E, out, in] expert
+// weights to the [E, in, out] layout GatherMM consumes, materialized once at
+// load so the forward path avoids per-call transposes.
+func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array {
+	if w == nil || !w.Valid() || w.NumDims() != 3 {
+		return w
+	}
+	t := mlx.Transpose(w, 0, 2, 1)
+	cloned := t.Clone()
+	mlx.Eval(cloned)
+	return cloned
+}
+
+// loadStackedProjection returns expert weights already stacked as a single 3D
+// tensor (layers.N.mlp.switch_mlp.<proj>.weight) — the layout `ollama create`
+// writes when it packs per-expert tensors at import.
+func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, base string) *stackedExpertWeights {
+	key := base + ".weight"
+	w := tensors[key]
+	if w == nil {
+		return nil
+	}
+
+	scales := tensors[key+"_scale"]
+	if scales == nil {
+		return &stackedExpertWeights{Weight: w}
+	}
+
+	qbiases := tensors[key+"_qbias"]
+	groupSize, bits, mode := model.ResolveLinearQuantParams(
+		cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant,
+		key, w, scales,
+	)
+	if useQuantized && supportsGatherQMM(mode, bits) {
+		return &stackedExpertWeights{
+			Weight:    w,
+			Scales:    scales,
+			Biases:    qbiases,
+			Bits:      bits,
+			GroupSize: groupSize,
+			Mode:      mode,
+		}
+	}
+
+	return &stackedExpertWeights{
+		Weight:    mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode),
+		Bits:      bits,
+		GroupSize: groupSize,
+		Mode:      mode,
+	}
+}
+
+// LoadWeights assigns tensors to model fields.
+func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
+	cfg := m.Config
+
+	linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
+
+	embedTokens := model.MakeEmbeddingLayer(tensors, "model.embed_tokens", cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
+	if embedTokens == nil {
+		return fmt.Errorf("missing embedding weight: model.embed_tokens.weight")
+	}
+	m.EmbedTokens = embedTokens
+
+	normWeight := tensors["model.norm.weight"]
+	if normWeight == nil {
+		return fmt.Errorf("missing final norm weight: model.norm.weight")
+	}
+	m.Norm = cfg.newNorm(normWeight)
+
+	if cfg.tieEmbeddings() {
+		m.LMHead = m.EmbedTokens.AsLinear()
+	} else if lmHead := linears.Make("lm_head"); lmHead != nil {
+		m.LMHead = lmHead
+	} else {
+		m.LMHead = m.EmbedTokens.AsLinear()
+	}
+
+	useQuantizedExperts := supportsGatherQMM(cfg.QuantMode, cfg.QuantBits)
+	if !useQuantizedExperts && cfg.TensorQuant != nil {
+		for _, tq := range cfg.TensorQuant {
+			if tq == nil {
+				continue
+			}
+			_, bits, mode := model.QuantizationParams(tq.QuantType)
+			if supportsGatherQMM(mode, bits) {
+				useQuantizedExperts = true
+				break
+			}
+		}
+	}
+
+	for i := range cfg.NumHiddenLayers {
+		layerPrefix := fmt.Sprintf("model.layers.%d", i)
+		layer := &Layer{
+			IsSliding: cfg.layerIsSliding(i),
+			UseRope:   cfg.layerUsesRope(i),
+		}
+
+		normWeight := tensors[layerPrefix+".input_layernorm.weight"]
+		if normWeight == nil {
+			return fmt.Errorf("layer %d: missing input_layernorm", i)
+		}
+		layer.InputNorm = cfg.newNorm(normWeight)
+
+		attn := &Attention{
+			QProj: linears.Make(layerPrefix + ".self_attn.q_proj"),
+			KProj: linears.Make(layerPrefix + ".self_attn.k_proj"),
+			VProj: linears.Make(layerPrefix + ".self_attn.v_proj"),
+			OProj: linears.Make(layerPrefix + ".self_attn.o_proj"),
+		}
+		if attn.QProj == nil || attn.KProj == nil || attn.VProj == nil || attn.OProj == nil {
+			return fmt.Errorf("layer %d: missing attention projections", i)
+		}
+		layer.Attention = attn
+
+		if cfg.layerIsDense(i) {
+			mlp := &DenseMLP{
+				GateProj: linears.Make(layerPrefix + ".mlp.gate_proj"),
+				UpProj:   linears.Make(layerPrefix + ".mlp.up_proj"),
+				DownProj: linears.Make(layerPrefix + ".mlp.down_proj"),
+			}
+			if mlp.GateProj == nil || mlp.UpProj == nil || mlp.DownProj == nil {
+				return fmt.Errorf("layer %d: missing dense mlp projections", i)
+			}
+			layer.MLP = mlp
+		} else {
+			moe := &SparseMoE{}
+			moe.Router = linears.Make(layerPrefix + ".mlp.gate")
+			if moe.Router == nil {
+				return fmt.Errorf("layer %d: missing moe router gate", i)
+			}
+
+			gateW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.gate_proj")
+			upW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.up_proj")
+			downW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.down_proj")
+			if gateW == nil || upW == nil || downW == nil {
+				return fmt.Errorf("layer %d: missing stacked switch_mlp expert weights (import the model with `ollama create`)", i)
+			}
+
+			switchMLP := &SwitchMLP{}
+			if gateW.Scales != nil && upW.Scales != nil && downW.Scales != nil {
+				switchMLP.UseQuantized = true
+				switchMLP.GateWeightQ = gateW.Weight
+				switchMLP.GateScales = gateW.Scales
+				switchMLP.GateBiases = gateW.Biases
+				switchMLP.GateBits = gateW.Bits
+				switchMLP.GateGroupSize = gateW.GroupSize
+				switchMLP.GateMode = gateW.Mode
+				switchMLP.UpWeightQ = upW.Weight
+				switchMLP.UpScales = upW.Scales
+				switchMLP.UpBiases = upW.Biases
+				switchMLP.UpBits = upW.Bits
+				switchMLP.UpGroupSize = upW.GroupSize
+				switchMLP.UpMode = upW.Mode
+				switchMLP.DownWeightQ = downW.Weight
+				switchMLP.DownScales = downW.Scales
+				switchMLP.DownBiases = downW.Biases
+				switchMLP.DownBits = downW.Bits
+				switchMLP.DownGroupSize = downW.GroupSize
+				switchMLP.DownMode = downW.Mode
+			} else {
+				switchMLP.GateWeight = transposeExpertWeightForGatherMM(gateW.Weight)
+				switchMLP.UpWeight = transposeExpertWeightForGatherMM(upW.Weight)
+				switchMLP.DownWeight = transposeExpertWeightForGatherMM(downW.Weight)
+			}
+			moe.SwitchMLP = switchMLP
+
+			if cfg.NumSharedExperts > 0 {
+				shared := &DenseMLP{
+					GateProj: linears.Make(layerPrefix + ".mlp.shared_experts.gate_proj"),
+					UpProj:   linears.Make(layerPrefix + ".mlp.shared_experts.up_proj"),
+					DownProj: linears.Make(layerPrefix + ".mlp.shared_experts.down_proj"),
+				}
+				if shared.GateProj == nil {
+					shared.GateProj = linears.Make(layerPrefix + ".mlp.shared_expert.gate_proj")
+					shared.UpProj = linears.Make(layerPrefix + ".mlp.shared_expert.up_proj")
+					shared.DownProj = linears.Make(layerPrefix + ".mlp.shared_expert.down_proj")
+				}
+				if shared.GateProj == nil || shared.UpProj == nil || shared.DownProj == nil {
+					return fmt.Errorf("layer %d: missing shared expert projections", i)
+				}
+				moe.SharedExpert = shared
+			}
+
+			layer.MLP = moe
+		}
+
+		m.Layers[i] = layer
+	}
+
+	return nil
+}
+
+func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, useRope bool, cfg *Config) *mlx.Array {
+	q := a.QProj.Forward(x)
+	k := a.KProj.Forward(x)
+	v := a.VProj.Forward(x)
+
+	q = mlx.Transpose(mlx.Reshape(q, B, L, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3)
+	k = mlx.Transpose(mlx.Reshape(k, B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3)
+	v = mlx.Transpose(mlx.Reshape(v, B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3)
+
+	// Cohere uses interleaved pairs (traditional RoPE). Full-attention layers
+	// outside the forced-RoPE prefix use no positional encoding.
+	if useRope {
+		q = mlx.RoPEWithBase(q, int(cfg.HeadDim), true, cfg.RopeTheta, 1.0, positions)
+		k = mlx.RoPEWithBase(k, int(cfg.HeadDim), true, cfg.RopeTheta, 1.0, positions)
+	}
+
+	var kv nn.SDPAOption
+	if c != nil {
+		history := c.(cache.Attention).Update(b, k, v)
+		kv = nn.WithKVHistory(history)
+	} else {
+		kv = nn.WithKV(k, v, b.SeqQueryLens)
+	}
+	out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask()))
+	out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim)
+	return a.OProj.Forward(out)
+}
+
+func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array {
+	return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x)))
+}
+
+// route selects the top-k experts. Selection happens on the raw router logits
+// and the activation (sigmoid or softmax) is applied to just the selected
+// entries, matching Cohere2MoeTopKRouter (both activations are monotonic, so
+// selection order is unchanged).
+func (moe *SparseMoE) route(x *mlx.Array, cfg *Config) (inds, scores *mlx.Array) {
+	logits := moe.Router.Forward(x)
+
+	inds = mlx.Argpartition(mlx.Neg(logits), int(cfg.NumExpertsPerTok)-1, -1)
+	dims := inds.Dims()
+	inds = mlx.SliceStartStop(inds, []int32{0, 0, 0}, []int32{int32(dims[0]), int32(dims[1]), cfg.NumExpertsPerTok})
+
+	selected := mlx.TakeAlongAxis(logits, inds, -1)
+	if cfg.ExpertSelectionFn == "sigmoid" {
+		scores = mlx.Sigmoid(selected)
+		if cfg.NormTopKProb && cfg.NumExpertsPerTok > 1 {
+			scores = mlx.Div(scores, mlx.Sum(scores, -1, true))
+		}
+	} else {
+		scores = mlx.SoftmaxAxis(selected, -1, true)
+	}
+	return inds, scores
+}
+
+func (moe *SparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array {
+	dims := x.Dims()
+	B, L := int32(dims[0]), int32(dims[1])
+
+	inds, scores := moe.route(x, cfg)
+
+	expertOut := moe.SwitchMLP.Forward(x, inds, cfg)
+	y := mlx.Sum(mlx.Mul(expertOut, mlx.ExpandDims(scores, -1)), 2, false)
+
+	if moe.SharedExpert != nil {
+		y = mlx.Add(y, moe.SharedExpert.Forward(x, cfg))
+		if cfg.SharedExpertCombinationStrategy == "average" {
+			y = mlx.MulScalar(y, 0.5)
+		}
+	}
+
+	return mlx.Reshape(y, B, L, cfg.HiddenSize)
+}
+
+func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.Array {
+	dims := x.Dims()
+	B, L := int32(dims[0]), int32(dims[1])
+	topK := cfg.NumExpertsPerTok
+
+	xFlat := mlx.Reshape(x, B*L, 1, 1, cfg.HiddenSize)
+	idxFlat := mlx.Reshape(indices, B*L, topK)
+
+	// Sorting tokens by expert improves gather matmul locality for prefill
+	// batches; the cost outweighs the benefit for small decode batches.
+	doSort := B*L >= 64
+	var invOrder *mlx.Array
+	n := B * L * topK
+
+	if doSort {
+		idxAll := mlx.Flatten(idxFlat)
+		order := mlx.Argsort(idxAll, 0)
+		invOrder = mlx.Argsort(order, 0)
+		xFlat = mlx.ExpandDims(mlx.Take(mlx.Squeeze(xFlat, 1), mlx.FloorDivideScalar(order, topK), 0), 1)
+		idxFlat = mlx.Reshape(mlx.Take(idxAll, order, 0), n, 1)
+	}
+
+	var gate, up, hidden, down *mlx.Array
+	if s.UseQuantized {
+		gate = mlx.GatherQMM(xFlat, s.GateWeightQ, s.GateScales, s.GateBiases,
+			nil, idxFlat, true, s.GateGroupSize, s.GateBits, s.GateMode, doSort)
+		up = mlx.GatherQMM(xFlat, s.UpWeightQ, s.UpScales, s.UpBiases,
+			nil, idxFlat, true, s.UpGroupSize, s.UpBits, s.UpMode, doSort)
+		hidden = mlx.SwiGLU(gate, up)
+		down = mlx.GatherQMM(hidden, s.DownWeightQ, s.DownScales, s.DownBiases,
+			nil, idxFlat, true, s.DownGroupSize, s.DownBits, s.DownMode, doSort)
+	} else {
+		gate = mlx.GatherMM(xFlat, s.GateWeight, nil, idxFlat, doSort)
+		up = mlx.GatherMM(xFlat, s.UpWeight, nil, idxFlat, doSort)
+		hidden = mlx.SwiGLU(gate, up)
+		down = mlx.GatherMM(hidden, s.DownWeight, nil, idxFlat, doSort)
+	}
+
+	if doSort {
+		down = mlx.Reshape(mlx.Take(mlx.Squeeze(mlx.Squeeze(down, 2), 1), invOrder, 0), B*L, topK, cfg.HiddenSize)
+	} else {
+		down = mlx.Squeeze(down, 2)
+	}
+
+	return mlx.Reshape(down, B, L, topK, cfg.HiddenSize)
+}
+
+// Forward runs a parallel-residual block: one shared layernorm feeds both
+// attention and the MLP, and the residual adds both outputs.
+func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array {
+	normed := l.InputNorm.Forward(x)
+	attnOut := l.Attention.Forward(normed, b, c, positions, B, L, l.UseRope, cfg)
+	mlpOut := l.MLP.Forward(normed, cfg)
+	return mlx.Add(x, mlx.Add(attnOut, mlpOut))
+}
+
+func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array {
+	dims := b.InputIDs.Dims()
+	B, L := int32(dims[0]), int32(dims[1])
+	positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))
+
+	h := m.EmbedTokens.Forward(b.InputIDs)
+	for i, layer := range m.Layers {
+		var c cache.Cache
+		if caches != nil && i < len(caches) {
+			c = caches[i]
+		}
+		h = layer.Forward(h, b, c, positions, B, L, m.Config)
+	}
+
+	return m.Norm.Forward(h)
+}
+
+func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
+	logits := m.LMHead.Forward(x)
+	if m.LogitScale != 1.0 {
+		logits = mlx.MulScalar(logits, m.LogitScale)
+	}
+	return logits
+}
+
+func (m *Model) NumLayers() int {
+	return len(m.Layers)
+}
+
+func (m *Model) MaxContextLength() int {
+	return int(m.MaxPositionEmbeddings)
+}
+
+func (m *Model) Tokenizer() *tokenizer.Tokenizer {
+	return m.tok
+}
+
+// NewCaches creates per-layer caches: rotating (bounded) caches for sliding
+// window layers and standard KV caches for full attention layers.
+func (m *Model) NewCaches() []cache.Cache {
+	caches := make([]cache.Cache, len(m.Layers))
+	for i, layer := range m.Layers {
+		if m.SlidingWindow > 0 && layer.IsSliding {
+			caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow))
+		} else {
+			caches[i] = cache.NewKVCache()
+		}
+	}
+	return caches
+}