models: add cohere2_moe (Command A / North) to the MLX engine (#16670)

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

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/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/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
 	}

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
+}