From b82e30d5b59a704b9292dcb8a29d721df7473dd6 Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 06:36:20 +0000 Subject: [PATCH 01/58] cmake: reconfigure MLX external builds Co-authored-by: Codex --- cmake/local.cmake | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cmake/local.cmake b/cmake/local.cmake index e136d9e0b79..7844f629cf2 100644 --- a/cmake/local.cmake +++ b/cmake/local.cmake @@ -546,7 +546,8 @@ function(ollama_add_mlx_build name) SOURCE_DIR ${CMAKE_SOURCE_DIR}/cmake/mlx BINARY_DIR ${_build_dir} CONFIGURE_COMMAND ${_configure_command} - BUILD_COMMAND ${OLLAMA_NATIVE_BUILD_TOOL_COMMAND} + BUILD_COMMAND ${_configure_command} + COMMAND ${OLLAMA_NATIVE_BUILD_TOOL_COMMAND} ${OLLAMA_NATIVE_CONFIG_ARG} ${OLLAMA_NATIVE_BUILD_TARGET_ARG} mlx ${OLLAMA_NATIVE_BUILD_TARGET_ARG} mlxc From dfd005039bb77499e0fbb8f27b9585de35a4c26e Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 12:10:56 +0000 Subject: [PATCH 02/58] responses: make namespace tool identity injective Preserve namespace wire identities across declarations, history, and output while rejecting ambiguous model-facing names before dispatch. Co-authored-by: Codex --- middleware/openai.go | 16 +- middleware/responses_web_search_test.go | 149 +++++--- openai/responses.go | 385 +++++++++++++++++---- openai/responses_test.go | 430 +++++++++++++++++++++++- 4 files changed, 844 insertions(+), 136 deletions(-) diff --git a/middleware/openai.go b/middleware/openai.go index 98e96151369..b6fbf0d4f25 100644 --- a/middleware/openai.go +++ b/middleware/openai.go @@ -1018,7 +1018,7 @@ func (w *WebSearchResponsesWriter) writeWebSearchResponse(final api.ChatResponse response := openai.ToResponse(w.inner.model, w.inner.responseID, w.inner.itemID, final, w.req) completedAt := time.Now().Unix() response.CompletedAt = &completedAt - response.Output = buildResponsesWebSearchOutput(response.Output, w.preSearchThinking, w.preSearchContent, calls, w.otherToolCalls) + response.Output = buildResponsesWebSearchOutput(w.req, response.Output, w.preSearchThinking, w.preSearchContent, calls, w.otherToolCalls) if response.Usage != nil { response.Usage.InputTokens = usage.PromptEvalCount response.Usage.OutputTokens = usage.EvalCount @@ -1032,7 +1032,7 @@ func (w *WebSearchResponsesWriter) writeWebSearchResponse(final api.ChatResponse // buildResponsesWebSearchOutput assembles the final non-streaming output in // model-leg order: pre-search reasoning/text, server and mixed tool calls, then // the final model output. -func buildResponsesWebSearchOutput(output []openai.ResponsesOutputItem, preSearchThinking, preSearchContent string, searchCalls []openai.ResponsesWebSearchCall, otherToolCalls []api.ToolCall) []openai.ResponsesOutputItem { +func buildResponsesWebSearchOutput(request openai.ResponsesRequest, output []openai.ResponsesOutputItem, preSearchThinking, preSearchContent string, searchCalls []openai.ResponsesWebSearchCall, otherToolCalls []api.ToolCall) []openai.ResponsesOutputItem { items := make([]openai.ResponsesOutputItem, 0, len(output)+len(searchCalls)+len(otherToolCalls)+2) if preSearchThinking != "" { items = append(items, openai.ResponsesOutputItem{ @@ -1061,17 +1061,7 @@ func buildResponsesWebSearchOutput(output []openai.ResponsesOutputItem, preSearc items = append(items, openai.WebSearchCallOutputItem(call)) } // function_call items from mixed responses - convertedCalls := openai.ToToolCalls(otherToolCalls) - for i, tc := range convertedCalls { - items = append(items, openai.ResponsesOutputItem{ - ID: fmt.Sprintf("fc_mixed_%d", i), - Type: "function_call", - Status: "completed", - CallID: tc.ID, - Name: tc.Function.Name, - Arguments: tc.Function.Arguments, - }) - } + items = append(items, openai.ResponsesFunctionCallOutputItems(request, "fc_mixed_", otherToolCalls)...) // remaining items (final reasoning, message, or function calls) items = append(items, output...) return items diff --git a/middleware/responses_web_search_test.go b/middleware/responses_web_search_test.go index daba465a744..f3822038a2e 100644 --- a/middleware/responses_web_search_test.go +++ b/middleware/responses_web_search_test.go @@ -764,55 +764,118 @@ func TestWebSearchResponsesWriterNonStreamingPreservesContentBeforeToolCall(t *t } func TestWebSearchResponsesWriterNonStreamingSurfacesMixedToolCalls(t *testing.T) { - gin.SetMode(gin.TestMode) - recorder := httptest.NewRecorder() - ctx, _ := gin.CreateTestContext(recorder) - request := openai.ResponsesRequest{Model: "test-model", Tools: []openai.ResponsesTool{{Type: "web_search"}, {Type: "function", Name: "get_weather", Description: ptr("weather"), Parameters: map[string]any{"type": "object"}}}} - inner := &ResponsesWriter{BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, model: request.Model, responseID: "resp_test", itemID: "msg_test", request: request} - writer := &WebSearchResponsesWriter{ - BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, inner: inner, req: request, - chat: &api.ChatRequest{Model: request.Model, Tools: api.Tools{openai.WebSearchFunctionTool()}}, - search: func(context.Context, string) (*api.WebSearchResponse, error) { - return &api.WebSearchResponse{}, nil + tests := []struct { + name string + clientTool openai.ResponsesTool + internalName string + wantName string + wantNamespace string + }{ + { + name: "namespaced", + clientTool: openai.ResponsesTool{Type: "namespace", Name: "weather", Tools: []openai.ResponsesTool{{ + Type: "function", Name: "lookup", Description: ptr("weather"), Parameters: map[string]any{"type": "object"}, + }}}, + internalName: "weather__lookup", + wantName: "lookup", + wantNamespace: "weather", }, - followUpChat: func(_ context.Context, messages []api.Message, _ api.Tools) (api.ChatResponse, error) { - // Verify the assistant message only contains the web_search tool call, - // not the get_weather tool call. - if len(messages) < 2 { - t.Fatalf("expected at least 2 messages, got %d", len(messages)) - } - assistant := messages[len(messages)-2] - if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].Function.Name != "web_search" { - t.Fatalf("assistant message should only have web_search tool call, got %#v", assistant.ToolCalls) - } - return api.ChatResponse{Done: true, Message: api.Message{Role: "assistant", Content: "done"}}, nil + { + name: "flat", + clientTool: openai.ResponsesTool{Type: "function", Name: "get_weather", Description: ptr("weather"), Parameters: map[string]any{"type": "object"}}, + internalName: "get_weather", + wantName: "get_weather", }, } - // Non-streaming response with both web_search and get_weather tool calls. - initial := api.ChatResponse{Done: true, Message: api.Message{ToolCalls: []api.ToolCall{ - {ID: "call_1", Function: api.ToolCallFunction{Name: "web_search", Arguments: testArgs(map[string]any{"query": "weather"})}}, - {ID: "call_2", Function: api.ToolCallFunction{Name: "get_weather", Arguments: testArgs(map[string]any{"city": "SF"})}}, - }}} - data, _ := json.Marshal(initial) - if _, err := writer.Write(data); err != nil { - t.Fatal(err) - } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + request := openai.ResponsesRequest{Model: "test-model", Tools: []openai.ResponsesTool{{Type: "web_search"}, test.clientTool}} + chat, err := openai.FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + inner := &ResponsesWriter{BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, model: request.Model, responseID: "resp_test", itemID: "msg_test", request: request} + writer := &WebSearchResponsesWriter{ + BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, inner: inner, req: request, + chat: chat, + search: func(context.Context, string) (*api.WebSearchResponse, error) { + return &api.WebSearchResponse{}, nil + }, + followUpChat: func(_ context.Context, messages []api.Message, _ api.Tools) (api.ChatResponse, error) { + if len(messages) < 2 { + t.Fatalf("expected at least 2 messages, got %d", len(messages)) + } + assistant := messages[len(messages)-2] + if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].Function.Name != "web_search" { + t.Fatalf("assistant message should only have web_search tool call, got %#v", assistant.ToolCalls) + } + return api.ChatResponse{Done: true, Message: api.Message{Role: "assistant", Content: "done"}}, nil + }, + } - var response openai.ResponsesResponse - if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { - t.Fatalf("decode response: %v: %s", err, recorder.Body.String()) - } + initial := api.ChatResponse{Done: true, Message: api.Message{ + Thinking: "I should search.", + Content: "Searching first.", + ToolCalls: []api.ToolCall{ + {ID: "call_search", Function: api.ToolCallFunction{Name: "web_search", Arguments: testArgs(map[string]any{"query": "weather"})}}, + {ID: "call_client", Function: api.ToolCallFunction{Name: test.internalName, Arguments: testArgs(map[string]any{"city": "SF"})}}, + }, + }} + data, _ := json.Marshal(initial) + if _, err := writer.Write(data); err != nil { + t.Fatal(err) + } - // Output should include a function_call item for get_weather. - var hasFunctionCall bool - for _, item := range response.Output { - if item.Type == "function_call" && item.Name == "get_weather" { - hasFunctionCall = true - } - } - if !hasFunctionCall { - t.Fatalf("mixed tool call (get_weather) was not surfaced: %#v", response.Output) + var response openai.ResponsesResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v: %s", err, recorder.Body.String()) + } + wantTypes := []string{"reasoning", "message", "web_search_call", "function_call", "message"} + if len(response.Output) != len(wantTypes) { + t.Fatalf("output = %#v", response.Output) + } + for i, wantType := range wantTypes { + if response.Output[i].Type != wantType { + t.Fatalf("output[%d].type = %q, want %q: %#v", i, response.Output[i].Type, wantType, response.Output) + } + } + if response.Output[0].EncryptedContent != "I should search." || response.Output[1].Content[0].Text != "Searching first." || response.Output[4].Content[0].Text != "done" { + t.Fatalf("reasoning/messages were reordered: %#v", response.Output) + } + if response.Output[2].Action == nil || response.Output[2].Action.Query != "weather" { + t.Fatalf("search item = %#v", response.Output[2]) + } + function := response.Output[3] + if function.ID != "fc_mixed_0" || function.CallID != "call_client" || function.Name != test.wantName || function.Namespace != test.wantNamespace { + t.Fatalf("function item = %#v", function) + } + var rawResponse struct { + Output []map[string]json.RawMessage `json:"output"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &rawResponse); err != nil { + t.Fatal(err) + } + _, hasNamespace := rawResponse.Output[3]["namespace"] + if hasNamespace != (test.wantNamespace != "") { + t.Fatalf("function namespace presence = %v, want %v: %s", hasNamespace, test.wantNamespace != "", recorder.Body.String()) + } + if function.Arguments != `{"city":"SF"}` { + t.Fatalf("function arguments changed: %q", function.Arguments) + } + var arguments map[string]any + if err := json.Unmarshal([]byte(function.Arguments), &arguments); err != nil || arguments["city"] != "SF" { + t.Fatalf("function arguments = %q (%v)", function.Arguments, err) + } + for _, item := range response.Output { + if item.Type == "function_call" && item.Name == "web_search" { + t.Fatalf("private web_search function leaked: %#v", response.Output) + } + } + }) } } diff --git a/openai/responses.go b/openai/responses.go index 4a399ad6ecb..7723d0799f2 100644 --- a/openai/responses.go +++ b/openai/responses.go @@ -159,7 +159,8 @@ type ResponsesFunctionCall struct { Type string `json:"type"` // always "function_call" CallID string `json:"call_id"` // the tool call ID Name string `json:"name"` // function name - Arguments string `json:"arguments"` // JSON arguments string + Namespace string `json:"namespace,omitempty"` + Arguments string `json:"arguments"` // JSON arguments string } func (ResponsesFunctionCall) responsesInputItem() {} @@ -414,6 +415,11 @@ type ResponsesRequest struct { // FromResponsesRequest converts a ResponsesRequest to api.ChatRequest func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { + resolver, err := newResponsesToolResolver(r) + if err != nil { + return nil, err + } + var messages []api.Message // Add instructions as system message if present @@ -496,7 +502,7 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { toolCall := api.ToolCall{ ID: v.CallID, Function: api.ToolCallFunction{ - Name: v.Name, + Name: resolver.internalName(v.Namespace, v.Name), Arguments: args, }, } @@ -577,28 +583,6 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { return nil, err } - // Convert tools from Responses API format to api.Tool format - var tools []api.Tool - hasWebSearch := HasWebSearchTool(r.Tools) - for _, t := range r.Tools { - if isWebSearchTool(t) { - tools = append(tools, WebSearchFunctionTool()) - continue - } - expanded, err := convertTools(t) - if err != nil { - return nil, err - } - for _, tool := range expanded { - // The built-in tool owns this name. Keeping a user-declared function - // with the same name makes a model call ambiguous. - if hasWebSearch && tool.Function.Name == "web_search" { - continue - } - tools = append(tools, tool) - } - } - // Handle text format (e.g. json_schema) var format json.RawMessage if r.Text != nil && r.Text.Format != nil { @@ -614,7 +598,7 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { Model: r.Model, Messages: messages, Options: options, - Tools: tools, + Tools: resolver.tools, Format: format, Think: think, }, nil @@ -653,32 +637,248 @@ func WebSearchFunctionTool() api.Tool { } } -// convertTools converts one Responses-API tool declaration to api.Tools. A -// "namespace" declaration groups member functions under a common name; it -// expands to those members with namespace-qualified names, since api.Tool -// carries only a flat function name. Dropping the members instead would -// leave the model with one schema-less pseudo-function and make every -// namespaced call undeclarable. -func convertTools(t ResponsesTool) ([]api.Tool, error) { +type responsesToolExternalName struct { + namespace string + name string +} + +type responsesToolName struct { + external responsesToolExternalName + internal string + declarationNode *responsesNamespaceNode + declared bool +} + +type responsesNamespaceNode struct { + id uint64 + name string + parent *responsesNamespaceNode +} + +type responsesNamespaceChild struct { + parentID uint64 + name string +} + +type responsesToolDeclaration struct { + namespaceID uint64 + name string +} + +type responsesToolResolver struct { + byExternal map[responsesToolExternalName]responsesToolName + byInternal map[string]responsesToolName + declarations map[responsesToolDeclaration]struct{} + namespaceChildren map[responsesNamespaceChild]struct{} + nextNamespaceID uint64 + tools []api.Tool + hasWebSearchTool bool +} + +func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error) { + resolver := &responsesToolResolver{ + byExternal: make(map[responsesToolExternalName]responsesToolName), + byInternal: make(map[string]responsesToolName), + declarations: make(map[responsesToolDeclaration]struct{}), + namespaceChildren: make(map[responsesNamespaceChild]struct{}), + hasWebSearchTool: HasWebSearchTool(r.Tools), + } + + for _, tool := range r.Tools { + expanded, err := convertTools(tool, nil, resolver) + if err != nil { + return nil, err + } + resolver.tools = append(resolver.tools, expanded...) + } + + for _, item := range r.Input.Items { + if call, ok := item.(ResponsesFunctionCall); ok { + if _, err := resolver.register(call.Namespace, call.Name, nil, false); err != nil { + return nil, err + } + } + } + + return resolver, nil +} + +func joinNamespaceToolName(namespace, name string) string { + if strings.HasSuffix(namespace, "__") { + return namespace + name + } + return namespace + "__" + name +} + +func namespaceParts(node *responsesNamespaceNode, leaf string) []string { + depth := 0 + for current := node; current != nil; current = current.parent { + depth++ + } + leafCount := 0 + if leaf != "" { + leafCount = 1 + } + parts := make([]string, depth+leafCount) + if leafCount != 0 { + parts[depth] = leaf + } + for current, i := node, depth-1; current != nil; current, i = current.parent, i-1 { + parts[i] = current.name + } + return parts +} + +func joinNamespaceParts(parts []string) string { + if len(parts) == 0 { + return "" + } + size := 0 + for i, part := range parts { + size += len(part) + if i > 0 && !strings.HasSuffix(parts[i-1], "__") { + size += 2 + } + } + var qualified strings.Builder + qualified.Grow(size) + qualified.WriteString(parts[0]) + for i, part := range parts[1:] { + if !strings.HasSuffix(parts[i], "__") { + qualified.WriteString("__") + } + qualified.WriteString(part) + } + return qualified.String() +} + +func qualifiedNamespace(node *responsesNamespaceNode) string { + return joinNamespaceParts(namespaceParts(node, "")) +} + +func describeResponsesToolName(name responsesToolName) string { + source := "" + if name.declared { + source = fmt.Sprintf(" declared at path %q", namespaceParts(name.declarationNode, name.external.name)) + } + if name.external.namespace == "" { + return fmt.Sprintf("flat function %q%s", name.external.name, source) + } + return fmt.Sprintf("namespace %q function %q%s", name.external.namespace, name.external.name, source) +} + +func (r *responsesToolResolver) register(namespace, name string, declarationNode *responsesNamespaceNode, declaration bool) (responsesToolName, error) { + if name == "" { + return responsesToolName{}, fmt.Errorf("responses function name must not be empty") + } + if namespace != "" && strings.HasPrefix(name, joinNamespaceToolName(namespace, "")) { + return responsesToolName{}, fmt.Errorf("responses namespace %q function %q is already namespace-qualified", namespace, name) + } + + external := responsesToolExternalName{namespace: namespace, name: name} + internal := name + if namespace != "" { + internal = joinNamespaceToolName(namespace, name) + } + candidate := responsesToolName{external: external, internal: internal, declarationNode: declarationNode, declared: declaration} + + if declaration { + namespaceID := uint64(0) + if declarationNode != nil { + namespaceID = declarationNode.id + } + declarationID := responsesToolDeclaration{namespaceID: namespaceID, name: name} + if _, ok := r.declarations[declarationID]; ok { + return responsesToolName{}, fmt.Errorf("duplicate responses tool declaration for %s", describeResponsesToolName(candidate)) + } + r.declarations[declarationID] = struct{}{} + } + + if existing, ok := r.byInternal[internal]; ok { + sameExternal := existing.external == external + sameDeclaration := !declaration || existing.declarationNode == declarationNode + if sameExternal && sameDeclaration { + return existing, nil + } + left, right := describeResponsesToolName(existing), describeResponsesToolName(candidate) + if right < left { + left, right = right, left + } + return responsesToolName{}, fmt.Errorf("responses tool identity collision for internal name %q between %s and %s", internal, left, right) + } + + r.byExternal[external] = candidate + r.byInternal[internal] = candidate + return candidate, nil +} + +func (r *responsesToolResolver) internalName(namespace, name string) string { + if resolved, ok := r.byExternal[responsesToolExternalName{namespace: namespace, name: name}]; ok { + return resolved.internal + } + return name +} + +func (r *responsesToolResolver) externalName(name string) (string, string) { + if resolved, ok := r.byInternal[name]; ok { + return resolved.external.name, resolved.external.namespace + } + return name, "" +} + +// convertTools expands one Responses declaration while registering every +// model-visible function identity. Declaration expansion and identity +// validation intentionally share this single recursive traversal. +func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolver *responsesToolResolver) ([]api.Tool, error) { + if t.Type == "web_search" && namespaceNode == nil { + return []api.Tool{WebSearchFunctionTool()}, nil + } + if t.Type != "namespace" { tool, err := convertTool(t) if err != nil { return nil, err } + if t.Type != "function" { + return []api.Tool{tool}, nil + } + + // The built-in tool retains upstream ownership of its private model name. + if resolver.hasWebSearchTool && namespaceNode == nil && t.Name == "web_search" { + return nil, nil + } + + namespace := qualifiedNamespace(namespaceNode) + resolved, err := resolver.register(namespace, t.Name, namespaceNode, true) + if err != nil { + return nil, err + } + tool.Function.Name = resolved.internal return []api.Tool{tool}, nil } + if t.Name == "" { + return nil, fmt.Errorf("responses namespace name must not be empty") + } + parentID := uint64(0) + if namespaceNode != nil { + parentID = namespaceNode.id + } + child := responsesNamespaceChild{parentID: parentID, name: t.Name} + if _, ok := resolver.namespaceChildren[child]; ok { + duplicate := &responsesNamespaceNode{name: t.Name, parent: namespaceNode} + return nil, fmt.Errorf("duplicate responses namespace declaration %q", qualifiedNamespace(duplicate)) + } + resolver.namespaceChildren[child] = struct{}{} + resolver.nextNamespaceID++ + childNode := &responsesNamespaceNode{id: resolver.nextNamespaceID, name: t.Name, parent: namespaceNode} + var tools []api.Tool for _, member := range t.Tools { - expanded, err := convertTools(member) + expanded, err := convertTools(member, childNode, resolver) if err != nil { return nil, err } - for i := range expanded { - if prefix := t.Name + "."; t.Name != "" && !strings.HasPrefix(expanded[i].Function.Name, prefix) { - expanded[i].Function.Name = prefix + expanded[i].Function.Name - } - } tools = append(tools, expanded...) } return tools, nil @@ -821,6 +1021,7 @@ type ResponsesOutputItem struct { Content []ResponsesOutputContent `json:"content,omitempty"` // for message CallID string `json:"call_id,omitempty"` // for function_call Name string `json:"name,omitempty"` // for function_call + Namespace string `json:"namespace,omitempty"` // for function_call Arguments string `json:"arguments,omitempty"` // for function_call Action *ResponsesWebSearchAction `json:"action,omitempty"` // for web_search_call @@ -914,17 +1115,7 @@ func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse, } if len(chatResponse.Message.ToolCalls) > 0 { - toolCalls := ToToolCalls(chatResponse.Message.ToolCalls) - for i, tc := range toolCalls { - output = append(output, ResponsesOutputItem{ - ID: fmt.Sprintf("fc_%s_%d", responseID, i), - Type: "function_call", - Status: "completed", - CallID: tc.ID, - Name: tc.Function.Name, - Arguments: tc.Function.Arguments, - }) - } + output = append(output, ResponsesFunctionCallOutputItems(request, fmt.Sprintf("fc_%s_", responseID), chatResponse.Message.ToolCalls)...) } else { output = append(output, ResponsesOutputItem{ ID: itemID, @@ -1020,6 +1211,31 @@ func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse, } } +// ResponsesFunctionCallOutputItems converts model-facing function calls back +// to Responses wire items using the request's validated namespace identities. +// idPrefix lets each Responses output owner preserve its established item IDs. +func ResponsesFunctionCallOutputItems(request ResponsesRequest, idPrefix string, toolCalls []api.ToolCall) []ResponsesOutputItem { + resolver, _ := newResponsesToolResolver(request) + converted := ToToolCalls(toolCalls) + items := make([]ResponsesOutputItem, 0, len(converted)) + for i, tc := range converted { + name, namespace := tc.Function.Name, "" + if resolver != nil { + name, namespace = resolver.externalName(tc.Function.Name) + } + items = append(items, ResponsesOutputItem{ + ID: fmt.Sprintf("%s%d", idPrefix, i), + Type: "function_call", + Status: "completed", + CallID: tc.ID, + Name: name, + Namespace: namespace, + Arguments: tc.Function.Arguments, + }) + } + return items +} + // Streaming events: // ResponsesStreamEvent represents a single Server-Sent Event for the Responses API. @@ -1037,6 +1253,7 @@ type ResponsesStreamConverter struct { itemID string model string request ResponsesRequest + resolver *responsesToolResolver // State tracking (mutated across Process calls) firstWrite bool @@ -1069,11 +1286,13 @@ func (c *ResponsesStreamConverter) newEvent(eventType string, data map[string]an // NewResponsesStreamConverter creates a new converter with the given configuration. func NewResponsesStreamConverter(responseID, itemID, model string, request ResponsesRequest) *ResponsesStreamConverter { + resolver, _ := newResponsesToolResolver(request) return &ResponsesStreamConverter{ responseID: responseID, itemID: itemID, model: model, request: request, + resolver: resolver, firstWrite: true, } } @@ -1130,21 +1349,7 @@ func (c *ResponsesStreamConverter) buildResponseObject(status string, output []a truncation = *c.request.Truncation } - var tools []any - if c.request.Tools != nil { - for _, t := range c.request.Tools { - tools = append(tools, map[string]any{ - "type": t.Type, - "name": t.Name, - "description": t.Description, - "strict": t.Strict, - "parameters": t.Parameters, - }) - } - } - if tools == nil { - tools = []any{} - } + tools := responsesToolsForStream(c.request.Tools) textFormat := map[string]any{"type": "text"} if c.request.Text != nil && c.request.Text.Format != nil { @@ -1223,6 +1428,28 @@ func (c *ResponsesStreamConverter) buildResponseObject(status string, output []a } } +func responsesToolsForStream(tools []ResponsesTool) []any { + if tools == nil { + return []any{} + } + + converted := make([]any, 0, len(tools)) + for _, tool := range tools { + value := map[string]any{ + "type": tool.Type, + "name": tool.Name, + "description": tool.Description, + "strict": tool.Strict, + "parameters": tool.Parameters, + } + if len(tool.Tools) > 0 { + value["tools"] = responsesToolsForStream(tool.Tools) + } + converted = append(converted, value) + } + return converted +} + func (c *ResponsesStreamConverter) createResponseCreatedEvent() ResponsesStreamEvent { return c.newEvent("response.created", map[string]any{ "response": c.buildResponseObject("in_progress", []any{}, nil), @@ -1320,28 +1547,40 @@ func (c *ResponsesStreamConverter) emitFunctionCallEvents(toolCalls []api.ToolCa for i, tc := range converted { outputIndex := c.outputIndex + i fcItemID := fmt.Sprintf("fc_%d_%d", rand.Intn(999999), i) + name, namespace := tc.Function.Name, "" + if c.resolver != nil { + name, namespace = c.resolver.externalName(tc.Function.Name) + } toolCallItem := map[string]any{ "id": fcItemID, "type": "function_call", "status": "completed", "call_id": tc.ID, - "name": tc.Function.Name, + "name": name, "arguments": tc.Function.Arguments, } + if namespace != "" { + toolCallItem["namespace"] = namespace + } c.completedItems = append(c.completedItems, toolCallItem) + addedItem := map[string]any{ + "id": fcItemID, + "type": "function_call", + "status": "in_progress", + "call_id": tc.ID, + "name": name, + "arguments": "", + } + if namespace != "" { + addedItem["namespace"] = namespace + } + events = append(events, c.newEvent("response.output_item.added", map[string]any{ "output_index": outputIndex, - "item": map[string]any{ - "id": fcItemID, - "type": "function_call", - "status": "in_progress", - "call_id": tc.ID, - "name": tc.Function.Name, - "arguments": "", - }, + "item": addedItem, }), c.newEvent("response.function_call_arguments.delta", map[string]any{ "item_id": fcItemID, diff --git a/openai/responses_test.go b/openai/responses_test.go index 98c20c5818b..0d3372965d6 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -2,12 +2,21 @@ package openai import ( "encoding/json" + "strings" "testing" "time" "github.com/ollama/ollama/api" ) +func responsesNamespace(name string, tools ...ResponsesTool) ResponsesTool { + return ResponsesTool{Type: "namespace", Name: name, Tools: tools} +} + +func responsesFunction(name string) ResponsesTool { + return ResponsesTool{Type: "function", Name: name, Parameters: map[string]any{"type": "object"}} +} + func TestResponsesInputMessage_UnmarshalJSON(t *testing.T) { tests := []struct { name string @@ -204,6 +213,21 @@ func TestUnmarshalResponsesInputItem(t *testing.T) { } }) + t.Run("namespaced function_call item", func(t *testing.T) { + got, err := unmarshalResponsesInputItem([]byte(`{"type":"function_call","call_id":"call_abc123","namespace":"mcp__marker__","name":"marker","arguments":"{}"}`)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + fc, ok := got.(ResponsesFunctionCall) + if !ok { + t.Fatalf("got type %T, want ResponsesFunctionCall", got) + } + if fc.Namespace != "mcp__marker__" || fc.Name != "marker" { + t.Fatalf("function identity = (%q, %q), want (%q, %q)", fc.Namespace, fc.Name, "mcp__marker__", "marker") + } + }) + t.Run("function_call_output item", func(t *testing.T) { got, err := unmarshalResponsesInputItem([]byte(`{"type": "function_call_output", "call_id": "call_abc123", "output": "the result"}`)) if err != nil { @@ -576,7 +600,7 @@ func TestFromResponsesRequest_NamespaceTools(t *testing.T) { }, { "type": "function", - "name": "muse.read_file", + "name": "read_file", "description": "Reads a file", "strict": true, "parameters": { @@ -613,13 +637,13 @@ func TestFromResponsesRequest_NamespaceTools(t *testing.T) { t.Fatalf("expected 3 converted tools, got %d: %v", len(chatReq.Tools), chatReq.Tools) } - // Member functions carry the namespace-qualified name; an already - // qualified member is not double-prefixed. - if got := chatReq.Tools[0].Function.Name; got != "muse.bash" { - t.Errorf("expected function name 'muse.bash', got %q", got) + // Member functions carry Codex-compatible double-underscore namespace + // qualification. + if got := chatReq.Tools[0].Function.Name; got != "muse__bash" { + t.Errorf("expected function name 'muse__bash', got %q", got) } - if got := chatReq.Tools[1].Function.Name; got != "muse.read_file" { - t.Errorf("expected function name 'muse.read_file', got %q", got) + if got := chatReq.Tools[1].Function.Name; got != "muse__read_file" { + t.Errorf("expected function name 'muse__read_file', got %q", got) } if got := chatReq.Tools[2].Function.Name; got != "plain" { t.Errorf("expected function name 'plain', got %q", got) @@ -637,6 +661,398 @@ func TestFromResponsesRequest_NamespaceTools(t *testing.T) { } } +func TestResponsesToolIdentityRejectsAmbiguity(t *testing.T) { + nestedABC := responsesNamespace("a", responsesNamespace("b", responsesFunction("c"))) + flatNamespaceABC := responsesNamespace("a__b", responsesFunction("c")) + + tests := []struct { + name string + tools []ResponsesTool + }{ + {"separator collision forward", []ResponsesTool{responsesNamespace("a", responsesFunction("b__c")), responsesNamespace("a__b", responsesFunction("c"))}}, + {"separator collision reverse", []ResponsesTool{responsesNamespace("a__b", responsesFunction("c")), responsesNamespace("a", responsesFunction("b__c"))}}, + {"trailing separator collision forward", []ResponsesTool{responsesNamespace("a", responsesFunction("b")), responsesNamespace("a__", responsesFunction("b"))}}, + {"trailing separator collision reverse", []ResponsesTool{responsesNamespace("a__", responsesFunction("b")), responsesNamespace("a", responsesFunction("b"))}}, + {"nested path collision forward", []ResponsesTool{nestedABC, flatNamespaceABC}}, + {"nested path collision reverse", []ResponsesTool{flatNamespaceABC, nestedABC}}, + {"prequalified leaf", []ResponsesTool{responsesNamespace("muse", responsesFunction("muse__read_file"))}}, + {"duplicate flat declarations", []ResponsesTool{responsesFunction("same"), responsesFunction("same")}}, + {"duplicate namespace leaves", []ResponsesTool{responsesNamespace("a", responsesFunction("b"), responsesFunction("b"))}}, + {"repeated namespace path", []ResponsesTool{responsesNamespace("a", responsesFunction("b")), responsesNamespace("a", responsesFunction("c"))}}, + {"flat namespace collision forward", []ResponsesTool{responsesFunction("a__b__c"), nestedABC}}, + {"flat namespace collision reverse", []ResponsesTool{nestedABC, responsesFunction("a__b__c")}}, + {"empty namespace", []ResponsesTool{responsesNamespace("", responsesFunction("b"))}}, + {"empty flat function", []ResponsesTool{responsesFunction("")}}, + {"empty namespace function", []ResponsesTool{responsesNamespace("a", responsesFunction(""))}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := FromResponsesRequest(ResponsesRequest{Model: "test", Input: ResponsesInput{Text: "hi"}, Tools: tt.tools}) + if err == nil { + t.Fatal("expected identity validation error") + } + if !strings.Contains(err.Error(), "responses") { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +func TestResponsesToolIdentityErrorsAreOrderIndependent(t *testing.T) { + pairs := []struct { + name string + left ResponsesTool + right ResponsesTool + }{ + {"separator", responsesNamespace("a", responsesFunction("b__c")), responsesNamespace("a__b", responsesFunction("c"))}, + {"trailing separator", responsesNamespace("a", responsesFunction("b")), responsesNamespace("a__", responsesFunction("b"))}, + {"nested path", responsesNamespace("a", responsesNamespace("b", responsesFunction("c"))), responsesNamespace("a__b", responsesFunction("c"))}, + {"flat namespace", responsesFunction("a__b__c"), responsesNamespace("a", responsesNamespace("b", responsesFunction("c")))}, + } + for _, tt := range pairs { + t.Run(tt.name, func(t *testing.T) { + getError := func(tools []ResponsesTool) string { + _, err := FromResponsesRequest(ResponsesRequest{Model: "test", Input: ResponsesInput{Text: "hi"}, Tools: tools}) + if err == nil { + t.Fatal("expected identity validation error") + } + return err.Error() + } + forward := getError([]ResponsesTool{tt.left, tt.right}) + reverse := getError([]ResponsesTool{tt.right, tt.left}) + if forward != reverse { + t.Fatalf("order-dependent errors:\nforward: %s\nreverse: %s", forward, reverse) + } + }) + } +} + +func TestResponsesNamespaceBookkeepingIsLinear(t *testing.T) { + deepNamespace := func(depth int, leaf *ResponsesTool) ResponsesTool { + var tool ResponsesTool + if leaf != nil { + tool = *leaf + } else { + tool = responsesNamespace("n") + depth-- + } + for range depth { + tool = responsesNamespace("n", tool) + } + return tool + } + + t.Run("deep namespace-only chain retains one segment per node", func(t *testing.T) { + const depth = 2048 + resolver, err := newResponsesToolResolver(ResponsesRequest{Tools: []ResponsesTool{deepNamespace(depth, nil)}}) + if err != nil { + t.Fatal(err) + } + if got := len(resolver.namespaceChildren); got != depth { + t.Fatalf("namespace node count = %d, want %d", got, depth) + } + retainedSegmentBytes := 0 + for child := range resolver.namespaceChildren { + retainedSegmentBytes += len(child.name) + } + if retainedSegmentBytes != depth { + t.Fatalf("retained namespace key bytes = %d, want %d", retainedSegmentBytes, depth) + } + if len(resolver.tools) != 0 { + t.Fatalf("namespace-only chain dispatched tools: %#v", resolver.tools) + } + }) + + t.Run("near-limit leaf builds only its qualified model name", func(t *testing.T) { + const depth = 1024 + leaf := responsesFunction("leaf") + resolver, err := newResponsesToolResolver(ResponsesRequest{Tools: []ResponsesTool{deepNamespace(depth, &leaf)}}) + if err != nil { + t.Fatal(err) + } + if got := len(resolver.namespaceChildren); got != depth { + t.Fatalf("namespace node count = %d, want %d", got, depth) + } + if len(resolver.tools) != 1 { + t.Fatalf("dispatched tool count = %d, want 1", len(resolver.tools)) + } + wantLength := depth*3 + len("leaf") + if got := len(resolver.tools[0].Function.Name); got != wantLength { + t.Fatalf("qualified leaf length = %d, want %d", got, wantLength) + } + }) +} + +func TestResponsesHistoryToolIdentity(t *testing.T) { + call := func(id, namespace, name string) ResponsesFunctionCall { + return ResponsesFunctionCall{Type: "function_call", CallID: id, Namespace: namespace, Name: name, Arguments: "{}"} + } + + rejections := []struct { + name string + items []ResponsesInputItem + }{ + {"flat namespace forward", []ResponsesInputItem{call("1", "", "a__b__c"), call("2", "a", "b__c")}}, + {"flat namespace reverse", []ResponsesInputItem{call("2", "a", "b__c"), call("1", "", "a__b__c")}}, + {"namespace collision forward", []ResponsesInputItem{call("1", "a", "b__c"), call("2", "a__b", "c")}}, + {"namespace collision reverse", []ResponsesInputItem{call("2", "a__b", "c"), call("1", "a", "b__c")}}, + {"prequalified history", []ResponsesInputItem{call("1", "muse", "muse__read_file")}}, + } + for _, tt := range rejections { + t.Run(tt.name, func(t *testing.T) { + _, err := FromResponsesRequest(ResponsesRequest{Model: "test", Input: ResponsesInput{Items: tt.items}}) + if err == nil { + t.Fatal("expected identity validation error") + } + }) + } + + t.Run("flat declaration collides with namespaced history", func(t *testing.T) { + _, err := FromResponsesRequest(ResponsesRequest{ + Model: "test", + Tools: []ResponsesTool{responsesFunction("a__b__c")}, + Input: ResponsesInput{Items: []ResponsesInputItem{call("1", "a", "b__c")}}, + }) + if err == nil { + t.Fatal("expected declaration/history identity collision") + } + }) + + t.Run("namespaced declaration collides with flat history", func(t *testing.T) { + _, err := FromResponsesRequest(ResponsesRequest{ + Model: "test", + Tools: []ResponsesTool{responsesNamespace("a", responsesFunction("b__c"))}, + Input: ResponsesInput{Items: []ResponsesInputItem{call("1", "", "a__b__c")}}, + }) + if err == nil { + t.Fatal("expected declaration/history identity collision") + } + }) + + t.Run("repeated exact history is accepted", func(t *testing.T) { + chat, err := FromResponsesRequest(ResponsesRequest{Model: "test", Input: ResponsesInput{Items: []ResponsesInputItem{ + call("1", "a", "b"), call("2", "a", "b"), + }}}) + if err != nil { + t.Fatal(err) + } + if got := chat.Messages[0].ToolCalls[0].Function.Name; got != "a__b" { + t.Fatalf("first history name = %q", got) + } + if got := chat.Messages[0].ToolCalls[1].Function.Name; got != "a__b" { + t.Fatalf("second history name = %q", got) + } + }) + + t.Run("declaration-free validator history is accepted", func(t *testing.T) { + chat, err := FromResponsesRequest(ResponsesRequest{Model: "test", Input: ResponsesInput{Items: []ResponsesInputItem{ + call("marker", "mcp__marker__", "marker"), + }}}) + if err != nil { + t.Fatal(err) + } + if got := chat.Messages[0].ToolCalls[0].Function.Name; got != "mcp__marker__marker" { + t.Fatalf("history name = %q", got) + } + }) +} + +func TestResponsesNamespaceRoundTrip(t *testing.T) { + declaration := responsesNamespace("a", responsesNamespace("b", responsesFunction("c"))) + request := ResponsesRequest{ + Model: "test", + Tools: []ResponsesTool{declaration}, + Input: ResponsesInput{Items: []ResponsesInputItem{ResponsesFunctionCall{ + Type: "function_call", CallID: "call_1", Namespace: "a__b", Name: "c", Arguments: "{}", + }}}, + } + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if got := chat.Tools[0].Function.Name; got != "a__b__c" { + t.Fatalf("declaration name = %q", got) + } + if got := chat.Messages[0].ToolCalls[0].Function.Name; got != "a__b__c" { + t.Fatalf("history name = %q", got) + } + + response := ToResponse("test", "resp", "item", api.ChatResponse{ + CreatedAt: time.Now(), + Message: api.Message{ToolCalls: []api.ToolCall{{ + ID: "call_2", Function: api.ToolCallFunction{Name: "a__b__c", Arguments: testArgs(map[string]any{})}, + }}}, + }, request) + if got := response.Output[0]; got.Name != "c" || got.Namespace != "a__b" { + t.Fatalf("non-stream identity = (%q, %q)", got.Namespace, got.Name) + } + + converter := NewResponsesStreamConverter("resp", "item", "test", request) + events := converter.Process(api.ChatResponse{ + Message: api.Message{ToolCalls: []api.ToolCall{{ + ID: "call_3", Function: api.ToolCallFunction{Name: "a__b__c", Arguments: testArgs(map[string]any{})}, + }}}, + Done: true, + }) + var added, done, terminal map[string]any + for _, event := range events { + data := event.Data.(map[string]any) + switch event.Event { + case "response.output_item.added": + item := data["item"].(map[string]any) + if item["type"] == "function_call" { + added = item + } + case "response.output_item.done": + item := data["item"].(map[string]any) + if item["type"] == "function_call" { + done = item + } + case "response.completed": + output := data["response"].(map[string]any)["output"].([]any) + terminal = output[0].(map[string]any) + } + } + for label, item := range map[string]map[string]any{"added": added, "done": done, "terminal": terminal} { + if item == nil || item["name"] != "c" || item["namespace"] != "a__b" { + t.Fatalf("%s identity = %#v", label, item) + } + } + + created := events[0].Data.(map[string]any)["response"].(map[string]any) + encoded, err := json.Marshal(created["tools"]) + if err != nil { + t.Fatal(err) + } + var tools []ResponsesTool + if err := json.Unmarshal(encoded, &tools); err != nil { + t.Fatal(err) + } + if len(tools) != 1 || len(tools[0].Tools) != 1 || len(tools[0].Tools[0].Tools) != 1 || tools[0].Tools[0].Tools[0].Name != "c" { + t.Fatalf("recursive stream metadata = %#v", tools) + } +} + +func TestResponsesToolIdentityPreservesFlatWebSearchAndNamespaceOrder(t *testing.T) { + request := ResponsesRequest{ + Model: "test", + Input: ResponsesInput{Text: "hi"}, + Tools: []ResponsesTool{ + responsesFunction("flat"), + {Type: "web_search"}, + responsesNamespace("a", responsesFunction("b")), + }, + } + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + want := []string{"flat", "web_search", "a__b"} + if len(chat.Tools) != len(want) { + t.Fatalf("tools = %#v", chat.Tools) + } + for i, tool := range chat.Tools { + if tool.Function.Name != want[i] { + t.Fatalf("tool[%d] = %q, want %q", i, tool.Function.Name, want[i]) + } + } +} + +func TestResponsesOutputInvalidResolverKeepsFlatName(t *testing.T) { + invalid := ResponsesRequest{Tools: []ResponsesTool{ + responsesFunction("a__b"), responsesNamespace("a", responsesFunction("b")), + }} + response := ToResponse("test", "resp", "item", api.ChatResponse{ + CreatedAt: time.Now(), + Message: api.Message{ToolCalls: []api.ToolCall{{ + ID: "call", Function: api.ToolCallFunction{Name: "a__b", Arguments: testArgs(map[string]any{})}, + }}}, + }, invalid) + if got := response.Output[0]; got.Name != "a__b" || got.Namespace != "" { + t.Fatalf("invalid resolver attributed output as (%q, %q)", got.Namespace, got.Name) + } + + converter := NewResponsesStreamConverter("resp", "item", "test", invalid) + events := converter.Process(api.ChatResponse{Message: api.Message{ToolCalls: []api.ToolCall{{ + ID: "call", Function: api.ToolCallFunction{Name: "a__b", Arguments: testArgs(map[string]any{})}, + }}}}) + for _, event := range events { + if event.Event != "response.output_item.done" { + continue + } + item := event.Data.(map[string]any)["item"].(map[string]any) + if item["name"] != "a__b" { + t.Fatalf("invalid stream resolver changed flat name: %#v", item) + } + if _, ok := item["namespace"]; ok { + t.Fatalf("invalid stream resolver attributed namespace: %#v", item) + } + return + } + t.Fatal("missing streamed function_call item") +} + +func TestResponsesStreamConverter_MixedNamespaceWebSearchReasoningOrder(t *testing.T) { + request := ResponsesRequest{Tools: []ResponsesTool{ + responsesNamespace("a", responsesFunction("b")), + {Type: "web_search"}, + }} + converter := NewResponsesStreamConverter("resp_mixed", "msg_mixed", "test", request) + var events []ResponsesStreamEvent + + events = append(events, converter.Process(api.ChatResponse{Message: api.Message{Thinking: "search first"}})...) + search := ResponsesWebSearchCall{ + ID: "ws_mixed", + Type: "web_search_call", + Status: "completed", + Action: &ResponsesWebSearchAction{Type: "search", Query: "Ollama"}, + } + searchIndex, searchEvents := converter.StartWebSearchCall(search) + events = append(events, searchEvents...) + events = append(events, converter.FinishWebSearchCall(search, searchIndex)...) + events = append(events, converter.EmitFunctionCallItems([]api.ToolCall{{ + ID: "call_mixed", Function: api.ToolCallFunction{Name: "a__b", Arguments: testArgs(map[string]any{})}, + }})...) + events = append(events, converter.Process(api.ChatResponse{Done: true})...) + + doneByIndex := map[int]map[string]any{} + var addedFunction map[string]any + var terminal []any + for _, event := range events { + data := event.Data.(map[string]any) + switch event.Event { + case "response.output_item.added": + item := data["item"].(map[string]any) + if item["type"] == "function_call" { + addedFunction = item + } + case "response.output_item.done": + doneByIndex[data["output_index"].(int)] = data["item"].(map[string]any) + case "response.completed": + terminal = data["response"].(map[string]any)["output"].([]any) + } + } + if addedFunction == nil || addedFunction["name"] != "b" || addedFunction["namespace"] != "a" { + t.Fatalf("namespaced added item = %#v", addedFunction) + } + wantTypes := []string{"reasoning", "web_search_call", "function_call"} + if len(doneByIndex) != len(wantTypes) || len(terminal) != len(wantTypes) { + t.Fatalf("done=%#v terminal=%#v", doneByIndex, terminal) + } + for i, wantType := range wantTypes { + done := doneByIndex[i] + final := terminal[i].(map[string]any) + if done["type"] != wantType || final["type"] != wantType || done["id"] != final["id"] { + t.Fatalf("output[%d] done=%#v terminal=%#v", i, done, final) + } + } + function := doneByIndex[2] + if function["name"] != "b" || function["namespace"] != "a" { + t.Fatalf("namespaced done item = %#v", function) + } +} + func TestFromResponsesRequest_ReasoningEffort(t *testing.T) { tests := []struct { name string From 15b572cb39026c7c39115e15322c382a115dc9ba Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 06:37:55 +0000 Subject: [PATCH 03/58] test(responses): cover prior namespace separators Co-authored-by: Codex --- openai/responses_test.go | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/openai/responses_test.go b/openai/responses_test.go index 0d3372965d6..eea5107c3c2 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -661,6 +661,36 @@ func TestFromResponsesRequest_NamespaceTools(t *testing.T) { } } +func TestJoinNamespacePartsUsesPriorSegmentSeparator(t *testing.T) { + tests := []struct { + name string + parts []string + want string + }{ + {"prior segment ends in separator", []string{"a__", "b"}, "a__b"}, + {"leaf separator is preserved", []string{"a", "b__"}, "a__b__"}, + {"nested path", []string{"a", "b__", "c", "d__", "e"}, "a__b__c__d__e"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := joinNamespaceParts(tt.parts); got != tt.want { + t.Fatalf("joinNamespaceParts(%q) = %q, want %q", tt.parts, got, tt.want) + } + }) + } + + t.Run("resolver rejects identical joined names", func(t *testing.T) { + _, err := newResponsesToolResolver(ResponsesRequest{Tools: []ResponsesTool{ + responsesNamespace("a__", responsesFunction("b")), + responsesNamespace("a", responsesFunction("b")), + }}) + if err == nil || !strings.Contains(err.Error(), "identity collision") { + t.Fatalf("resolver error = %v, want identity collision", err) + } + }) +} + func TestResponsesToolIdentityRejectsAmbiguity(t *testing.T) { nestedABC := responsesNamespace("a", responsesNamespace("b", responsesFunction("c"))) flatNamespaceABC := responsesNamespace("a__b", responsesFunction("c")) From 81ee373f38b5f165fffa13df95a042485814c49f Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 13:49:10 +0000 Subject: [PATCH 04/58] openai: adapt custom apply_patch responses tool Bridge the narrow Responses custom apply_patch grammar through the request-wide injective resolver, preserving history, mixed output, streaming identity, and safe fallback behavior. Co-authored-by: Codex --- middleware/responses_web_search_test.go | 181 ++++++++++ openai/responses.go | 231 ++++++++++++- openai/responses_test.go | 430 ++++++++++++++++++++++++ 3 files changed, 830 insertions(+), 12 deletions(-) diff --git a/middleware/responses_web_search_test.go b/middleware/responses_web_search_test.go index f3822038a2e..ccd54173191 100644 --- a/middleware/responses_web_search_test.go +++ b/middleware/responses_web_search_test.go @@ -939,6 +939,187 @@ func TestWebSearchResponsesWriterStreamingSurfacesMixedToolCalls(t *testing.T) { } } +func TestWebSearchResponsesWriterNonStreamingSurfacesCustomApplyPatch(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + format := json.RawMessage(`{"type":"grammar","syntax":"lark","definition":"start: /.+/"}`) + request := openai.ResponsesRequest{Model: "test-model", Tools: []openai.ResponsesTool{ + {Type: "web_search"}, + {Type: "custom", Name: "apply_patch", Format: format}, + }} + chat, err := openai.FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + inner := &ResponsesWriter{BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, model: request.Model, responseID: "resp_test", itemID: "msg_test", request: request} + writer := &WebSearchResponsesWriter{ + BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, inner: inner, req: request, chat: chat, + search: func(context.Context, string) (*api.WebSearchResponse, error) { return &api.WebSearchResponse{}, nil }, + followUpChat: func(_ context.Context, messages []api.Message, _ api.Tools) (api.ChatResponse, error) { + assistant := messages[len(messages)-2] + if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].Function.Name != "web_search" { + t.Fatalf("assistant search call = %#v", assistant.ToolCalls) + } + return api.ChatResponse{Done: true, Message: api.Message{Role: "assistant", Content: "done"}}, nil + }, + } + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + initial := api.ChatResponse{Done: true, Message: api.Message{ + Thinking: "I should search.", Content: "Searching first.", + ToolCalls: []api.ToolCall{ + {ID: "call_search", Function: api.ToolCallFunction{Name: "web_search", Arguments: testArgs(map[string]any{"query": "weather"})}}, + {ID: "call_patch", Function: api.ToolCallFunction{Name: "apply_patch", Arguments: testArgs(map[string]any{"input": patch})}}, + }, + }} + data, _ := json.Marshal(initial) + if _, err := writer.Write(data); err != nil { + t.Fatal(err) + } + + var response openai.ResponsesResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v: %s", err, recorder.Body.String()) + } + wantTypes := []string{"reasoning", "message", "web_search_call", "custom_tool_call", "message"} + if len(response.Output) != len(wantTypes) { + t.Fatalf("output = %#v", response.Output) + } + for i, want := range wantTypes { + if response.Output[i].Type != want { + t.Fatalf("output[%d].type = %q, want %q: %#v", i, response.Output[i].Type, want, response.Output) + } + } + custom := response.Output[3] + if custom.ID != "ctc_mixed_0" || custom.CallID != "call_patch" || custom.Name != "apply_patch" || custom.Input != patch || custom.Arguments != "" { + t.Fatalf("custom item = %#v", custom) + } + var raw struct { + Output []map[string]json.RawMessage `json:"output"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &raw); err != nil { + t.Fatal(err) + } + if _, ok := raw.Output[3]["arguments"]; ok { + t.Fatalf("custom item leaked function arguments: %s", recorder.Body.String()) + } + for _, item := range response.Output { + if item.Type == "function_call" && item.Name == "web_search" { + t.Fatalf("private web_search function leaked: %#v", response.Output) + } + } +} + +func TestWebSearchResponsesWriterStreamingSurfacesCustomApplyPatch(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + stream := true + request := openai.ResponsesRequest{Model: "test-model", Stream: &stream, Tools: []openai.ResponsesTool{ + {Type: "web_search"}, {Type: "custom", Name: "apply_patch"}, + }} + inner := &ResponsesWriter{BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, converter: openai.NewResponsesStreamConverter("resp_test", "msg_test", request.Model, request), model: request.Model, stream: true, responseID: "resp_test", itemID: "msg_test", request: request} + writer := &WebSearchResponsesWriter{ + BaseWriter: BaseWriter{ResponseWriter: ctx.Writer}, inner: inner, req: request, + chat: &api.ChatRequest{Model: request.Model, Tools: api.Tools{openai.WebSearchFunctionTool()}}, + search: func(context.Context, string) (*api.WebSearchResponse, error) { return &api.WebSearchResponse{}, nil }, + followUpChat: func(_ context.Context, messages []api.Message, _ api.Tools) (api.ChatResponse, error) { + return api.ChatResponse{Done: true, Message: api.Message{Role: "assistant", Content: "done"}}, nil + }, + } + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + initial := api.ChatResponse{Done: true, Message: api.Message{ + Thinking: "think", Content: "searching", + ToolCalls: []api.ToolCall{ + {ID: "call_search", Function: api.ToolCallFunction{Name: "web_search", Arguments: testArgs(map[string]any{"query": "weather"})}}, + {ID: "call_patch", Function: api.ToolCallFunction{Name: "apply_patch", Arguments: testArgs(map[string]any{"input": patch})}}, + }, + }} + data, _ := json.Marshal(initial) + if _, err := writer.Write(data); err != nil { + t.Fatal(err) + } + body := recorder.Body.String() + for _, want := range []string{"response.web_search_call.completed", "response.custom_tool_call_input.delta", "response.custom_tool_call_input.done"} { + if !strings.Contains(body, want) { + t.Fatalf("missing %s: %s", want, body) + } + } + output := completedResponseOutput(t, body) + wantTypes := []string{"reasoning", "message", "custom_tool_call", "web_search_call", "message"} + if len(output) != len(wantTypes) { + t.Fatalf("terminal output = %#v", output) + } + for i, want := range wantTypes { + if output[i]["type"] != want { + t.Fatalf("terminal[%d] = %#v, want %q", i, output[i], want) + } + } + custom := output[2] + if custom["call_id"] != "call_patch" || custom["name"] != "apply_patch" || custom["input"] != patch { + t.Fatalf("terminal custom item = %#v", custom) + } + if _, ok := custom["arguments"]; ok { + t.Fatalf("terminal custom item leaked arguments: %#v", custom) + } + + doneItems := map[int]map[string]any{} + lastDoneIndex := -1 + var customDeltaID, customDoneID string + var customDeltaIndex, customDoneIndex int + for _, block := range strings.Split(body, "\n\n") { + if !strings.HasPrefix(block, "event: ") { + continue + } + lineEnd := strings.IndexByte(block, '\n') + dataAt := strings.Index(block, "\ndata: ") + if lineEnd < 0 || dataAt < 0 { + continue + } + eventType := strings.TrimPrefix(block[:lineEnd], "event: ") + var event map[string]any + if err := json.Unmarshal([]byte(block[dataAt+7:]), &event); err != nil { + t.Fatalf("decode %s event: %v: %s", eventType, err, block) + } + switch eventType { + case "response.output_item.done": + index := int(event["output_index"].(float64)) + if index <= lastDoneIndex { + t.Fatalf("non-monotonic done index %d after %d", index, lastDoneIndex) + } + lastDoneIndex = index + doneItems[index] = event["item"].(map[string]any) + case "response.custom_tool_call_input.delta": + customDeltaID = event["item_id"].(string) + customDeltaIndex = int(event["output_index"].(float64)) + case "response.custom_tool_call_input.done": + customDoneID = event["item_id"].(string) + customDoneIndex = int(event["output_index"].(float64)) + } + } + if len(doneItems) != len(output) { + t.Fatalf("done items = %#v, terminal = %#v", doneItems, output) + } + for index, terminal := range output { + done := doneItems[index] + if done == nil || done["id"] != terminal["id"] || done["type"] != terminal["type"] { + t.Fatalf("done/terminal mismatch at %d: done=%#v terminal=%#v", index, done, terminal) + } + } + if customDeltaID != custom["id"] || customDoneID != custom["id"] || customDeltaIndex != 2 || customDoneIndex != 2 { + t.Fatalf("custom event identity/index mismatch: delta=(%q,%d) done=(%q,%d) terminal=%#v", customDeltaID, customDeltaIndex, customDoneID, customDoneIndex, custom) + } + search := output[3] + if search["type"] != "web_search_call" || doneItems[3]["id"] != search["id"] { + t.Fatalf("search done/terminal identity mismatch: done=%#v terminal=%#v", doneItems[3], search) + } + for _, item := range output { + if item["type"] == "function_call" && item["name"] == "web_search" { + t.Fatalf("private web_search function leaked: %#v", output) + } + } +} + func completedResponseOutput(t *testing.T, body string) []map[string]any { t.Helper() for _, block := range strings.Split(body, "\n\n") { diff --git a/openai/responses.go b/openai/responses.go index 7723d0799f2..cbd18c4eea0 100644 --- a/openai/responses.go +++ b/openai/responses.go @@ -228,6 +228,28 @@ func (o *ResponsesFunctionCallOutput) UnmarshalJSON(data []byte) error { func (ResponsesFunctionCallOutput) responsesInputItem() {} +// ResponsesCustomToolCall represents an assistant custom/freeform tool call +// in conversation history. +type ResponsesCustomToolCall struct { + ID string `json:"id,omitempty"` + Type string `json:"type"` + CallID string `json:"call_id"` + Name string `json:"name"` + Input string `json:"input"` +} + +func (ResponsesCustomToolCall) responsesInputItem() {} + +// ResponsesCustomToolCallOutput represents a custom/freeform tool result from +// the client. +type ResponsesCustomToolCallOutput struct { + Type string `json:"type"` + CallID string `json:"call_id"` + Output string `json:"output"` +} + +func (ResponsesCustomToolCallOutput) responsesInputItem() {} + // ResponsesReasoningInput represents a reasoning item passed back as input. // This is used when the client sends previous reasoning back for context. type ResponsesReasoningInput struct { @@ -275,6 +297,18 @@ func unmarshalResponsesInputItem(data []byte) (ResponsesInputItem, error) { return nil, err } return output, nil + case "custom_tool_call": + var call ResponsesCustomToolCall + if err := json.Unmarshal(data, &call); err != nil { + return nil, err + } + return call, nil + case "custom_tool_call_output": + var output ResponsesCustomToolCallOutput + if err := json.Unmarshal(data, &output); err != nil { + return nil, err + } + return output, nil case "reasoning": var reasoning ResponsesReasoningInput if err := json.Unmarshal(data, &reasoning); err != nil { @@ -354,7 +388,7 @@ type ResponsesText struct { // ResponsesTool represents a tool in the Responses API format. // Note: This differs from api.Tool which nests fields under "function". type ResponsesTool struct { - Type string `json:"type"` // "function", "namespace", or "web_search" + Type string `json:"type"` // "function", "namespace", "custom", or "web_search" Name string `json:"name"` Description *string `json:"description"` // nullable but required Strict *bool `json:"strict"` // nullable but required @@ -363,7 +397,8 @@ type ResponsesTool struct { // Tools carries a "namespace" declaration's member functions. The // Responses API groups related tools by domain under a namespace tool // whose nested tools array holds the real function definitions. - Tools []ResponsesTool `json:"tools,omitempty"` + Tools []ResponsesTool `json:"tools,omitempty"` + Format json.RawMessage `json:"format,omitempty"` } type ResponsesRequest struct { @@ -462,7 +497,10 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { // so the tool result immediately follows the call it answers. var outputCallID string if i+1 < len(r.Input.Items) { - if output, ok := r.Input.Items[i+1].(ResponsesFunctionCallOutput); ok { + switch output := r.Input.Items[i+1].(type) { + case ResponsesFunctionCallOutput: + outputCallID = output.CallID + case ResponsesCustomToolCallOutput: outputCallID = output.CallID } } @@ -502,7 +540,7 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { toolCall := api.ToolCall{ ID: v.CallID, Function: api.ToolCallFunction{ - Name: resolver.internalName(v.Namespace, v.Name), + Name: resolver.internalName(responsesToolKindFunction, v.Namespace, v.Name), Arguments: args, }, } @@ -542,6 +580,37 @@ func FromResponsesRequest(r ResponsesRequest) (*api.ChatRequest, error) { Images: images, ToolCallID: v.CallID, }) + case ResponsesCustomToolCall: + args := api.NewToolCallFunctionArguments() + args.Set("input", v.Input) + toolCall := api.ToolCall{ + ID: v.CallID, + Function: api.ToolCallFunction{ + Name: resolver.internalName(responsesToolKindCustom, "", v.Name), + Arguments: args, + }, + } + if len(messages) > 0 && messages[len(messages)-1].Role == "assistant" { + lastMsg := &messages[len(messages)-1] + lastMsg.ToolCalls = append(lastMsg.ToolCalls, toolCall) + if pendingThinking != "" { + lastMsg.Thinking = pendingThinking + pendingThinking = "" + } + } else { + msg := api.Message{Role: "assistant", ToolCalls: []api.ToolCall{toolCall}} + if pendingThinking != "" { + msg.Thinking = pendingThinking + pendingThinking = "" + } + messages = append(messages, msg) + } + case ResponsesCustomToolCallOutput: + messages = append(messages, api.Message{ + Role: "tool", + Content: v.Output, + ToolCallID: v.CallID, + }) case ResponsesWebSearchCall: // Built-in tool calls are history metadata. The assistant message // that follows carries the model-visible result of the prior search. @@ -637,7 +706,15 @@ func WebSearchFunctionTool() api.Tool { } } +type responsesToolKind string + +const ( + responsesToolKindFunction responsesToolKind = "function" + responsesToolKindCustom responsesToolKind = "custom" +) + type responsesToolExternalName struct { + kind responsesToolKind namespace string name string } @@ -661,6 +738,7 @@ type responsesNamespaceChild struct { } type responsesToolDeclaration struct { + kind responsesToolKind namespaceID uint64 name string } @@ -693,8 +771,16 @@ func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error } for _, item := range r.Input.Items { - if call, ok := item.(ResponsesFunctionCall); ok { - if _, err := resolver.register(call.Namespace, call.Name, nil, false); err != nil { + switch call := item.(type) { + case ResponsesFunctionCall: + if _, err := resolver.register(responsesToolKindFunction, call.Namespace, call.Name, nil, false); err != nil { + return nil, err + } + case ResponsesCustomToolCall: + if call.Name != "apply_patch" { + return nil, fmt.Errorf("unsupported responses custom tool call %q", call.Name) + } + if _, err := resolver.register(responsesToolKindCustom, "", call.Name, nil, false); err != nil { return nil, err } } @@ -761,13 +847,16 @@ func describeResponsesToolName(name responsesToolName) string { if name.declared { source = fmt.Sprintf(" declared at path %q", namespaceParts(name.declarationNode, name.external.name)) } + if name.external.kind == responsesToolKindCustom { + return fmt.Sprintf("custom tool %q%s", name.external.name, source) + } if name.external.namespace == "" { return fmt.Sprintf("flat function %q%s", name.external.name, source) } return fmt.Sprintf("namespace %q function %q%s", name.external.namespace, name.external.name, source) } -func (r *responsesToolResolver) register(namespace, name string, declarationNode *responsesNamespaceNode, declaration bool) (responsesToolName, error) { +func (r *responsesToolResolver) register(kind responsesToolKind, namespace, name string, declarationNode *responsesNamespaceNode, declaration bool) (responsesToolName, error) { if name == "" { return responsesToolName{}, fmt.Errorf("responses function name must not be empty") } @@ -775,7 +864,7 @@ func (r *responsesToolResolver) register(namespace, name string, declarationNode return responsesToolName{}, fmt.Errorf("responses namespace %q function %q is already namespace-qualified", namespace, name) } - external := responsesToolExternalName{namespace: namespace, name: name} + external := responsesToolExternalName{kind: kind, namespace: namespace, name: name} internal := name if namespace != "" { internal = joinNamespaceToolName(namespace, name) @@ -787,7 +876,7 @@ func (r *responsesToolResolver) register(namespace, name string, declarationNode if declarationNode != nil { namespaceID = declarationNode.id } - declarationID := responsesToolDeclaration{namespaceID: namespaceID, name: name} + declarationID := responsesToolDeclaration{kind: kind, namespaceID: namespaceID, name: name} if _, ok := r.declarations[declarationID]; ok { return responsesToolName{}, fmt.Errorf("duplicate responses tool declaration for %s", describeResponsesToolName(candidate)) } @@ -812,8 +901,8 @@ func (r *responsesToolResolver) register(namespace, name string, declarationNode return candidate, nil } -func (r *responsesToolResolver) internalName(namespace, name string) string { - if resolved, ok := r.byExternal[responsesToolExternalName{namespace: namespace, name: name}]; ok { +func (r *responsesToolResolver) internalName(kind responsesToolKind, namespace, name string) string { + if resolved, ok := r.byExternal[responsesToolExternalName{kind: kind, namespace: namespace, name: name}]; ok { return resolved.internal } return name @@ -826,6 +915,11 @@ func (r *responsesToolResolver) externalName(name string) (string, string) { return name, "" } +func (r *responsesToolResolver) customName(name string) (string, bool) { + resolved, ok := r.byInternal[name] + return resolved.external.name, ok && resolved.external.kind == responsesToolKindCustom +} + // convertTools expands one Responses declaration while registering every // model-visible function identity. Declaration expansion and identity // validation intentionally share this single recursive traversal. @@ -833,6 +927,24 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv if t.Type == "web_search" && namespaceNode == nil { return []api.Tool{WebSearchFunctionTool()}, nil } + if t.Type == "custom" { + if namespaceNode != nil { + return nil, fmt.Errorf("responses custom tool %q must be top-level", t.Name) + } + if t.Name != "apply_patch" { + return nil, fmt.Errorf("unsupported responses custom tool %q", t.Name) + } + tool, err := convertTool(t) + if err != nil { + return nil, err + } + resolved, err := resolver.register(responsesToolKindCustom, "", t.Name, nil, true) + if err != nil { + return nil, err + } + tool.Function.Name = resolved.internal + return []api.Tool{tool}, nil + } if t.Type != "namespace" { tool, err := convertTool(t) @@ -849,7 +961,7 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv } namespace := qualifiedNamespace(namespaceNode) - resolved, err := resolver.register(namespace, t.Name, namespaceNode, true) + resolved, err := resolver.register(responsesToolKindFunction, namespace, t.Name, namespaceNode, true) if err != nil { return nil, err } @@ -885,6 +997,16 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv } func convertTool(t ResponsesTool) (api.Tool, error) { + if t.Type == "custom" && t.Name == "apply_patch" { + return api.Tool{ + Type: "function", + Function: api.ToolFunction{ + Name: "apply_patch", + Description: "Apply a patch to files. The input field must contain the complete raw patch text.", + Parameters: applyPatchFunctionParameters(), + }, + }, nil + } // Convert parameters from map[string]any to api.ToolFunctionParameters var params api.ToolFunctionParameters if t.Parameters != nil { @@ -913,6 +1035,37 @@ func convertTool(t ResponsesTool) (api.Tool, error) { }, nil } +func applyPatchFunctionParameters() api.ToolFunctionParameters { + properties := api.NewToolPropertiesMap() + properties.Set("input", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "Raw patch text beginning with *** Begin Patch and ending with *** End Patch.", + }) + return api.ToolFunctionParameters{ + Type: "object", + Required: []string{"input"}, + Properties: properties, + } +} + +func applyPatchInput(toolCall api.ToolCall) (string, bool) { + if toolCall.Function.Name != "apply_patch" { + return "", false + } + input, ok := toolCall.Function.Arguments.Get("input") + if !ok { + return "", false + } + patch, ok := input.(string) + if !ok || !strings.HasPrefix(patch, "*** Begin Patch\n") { + return "", false + } + if !strings.HasSuffix(strings.TrimRight(patch, "\n"), "*** End Patch") { + return "", false + } + return patch, true +} + func convertInputMessage(m ResponsesInputMessage) (api.Message, error) { content, images, err := convertResponsesContent(m.Content) if err != nil { @@ -1023,6 +1176,7 @@ type ResponsesOutputItem struct { Name string `json:"name,omitempty"` // for function_call Namespace string `json:"namespace,omitempty"` // for function_call Arguments string `json:"arguments,omitempty"` // for function_call + Input string `json:"input,omitempty"` // for custom_tool_call Action *ResponsesWebSearchAction `json:"action,omitempty"` // for web_search_call // Reasoning fields @@ -1219,6 +1373,21 @@ func ResponsesFunctionCallOutputItems(request ResponsesRequest, idPrefix string, converted := ToToolCalls(toolCalls) items := make([]ResponsesOutputItem, 0, len(converted)) for i, tc := range converted { + if resolver != nil { + if name, custom := resolver.customName(tc.Function.Name); custom { + if input, ok := applyPatchInput(toolCalls[i]); ok { + items = append(items, ResponsesOutputItem{ + ID: customToolCallItemID(idPrefix, i), + Type: "custom_tool_call", + Status: "completed", + CallID: tc.ID, + Name: name, + Input: input, + }) + continue + } + } + } name, namespace := tc.Function.Name, "" if resolver != nil { name, namespace = resolver.externalName(tc.Function.Name) @@ -1236,6 +1405,10 @@ func ResponsesFunctionCallOutputItems(request ResponsesRequest, idPrefix string, return items } +func customToolCallItemID(idPrefix string, index int) string { + return fmt.Sprintf("ctc_%s%d", strings.TrimPrefix(idPrefix, "fc_"), index) +} + // Streaming events: // ResponsesStreamEvent represents a single Server-Sent Event for the Responses API. @@ -1445,6 +1618,9 @@ func responsesToolsForStream(tools []ResponsesTool) []any { if len(tool.Tools) > 0 { value["tools"] = responsesToolsForStream(tool.Tools) } + if tool.Format != nil { + value["format"] = tool.Format + } converted = append(converted, value) } return converted @@ -1546,6 +1722,37 @@ func (c *ResponsesStreamConverter) emitFunctionCallEvents(toolCalls []api.ToolCa for i, tc := range converted { outputIndex := c.outputIndex + i + if c.resolver != nil { + if name, custom := c.resolver.customName(tc.Function.Name); custom { + if input, ok := applyPatchInput(toolCalls[i]); ok { + itemID := fmt.Sprintf("ctc_%d_%d", rand.Intn(999999), i) + item := map[string]any{ + "id": itemID, "type": "custom_tool_call", "status": "completed", + "call_id": tc.ID, "name": name, "input": input, + } + c.completedItems = append(c.completedItems, item) + events = append(events, + c.newEvent("response.output_item.added", map[string]any{ + "output_index": outputIndex, + "item": map[string]any{ + "id": itemID, "type": "custom_tool_call", "status": "in_progress", + "call_id": tc.ID, "name": name, "input": "", + }, + }), + c.newEvent("response.custom_tool_call_input.delta", map[string]any{ + "item_id": itemID, "output_index": outputIndex, "delta": input, + }), + c.newEvent("response.custom_tool_call_input.done", map[string]any{ + "item_id": itemID, "output_index": outputIndex, "input": input, + }), + c.newEvent("response.output_item.done", map[string]any{ + "output_index": outputIndex, "item": item, + }), + ) + continue + } + } + } fcItemID := fmt.Sprintf("fc_%d_%d", rand.Intn(999999), i) name, namespace := tc.Function.Name, "" if c.resolver != nil { diff --git a/openai/responses_test.go b/openai/responses_test.go index eea5107c3c2..e24bd4e9e47 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -2888,3 +2888,433 @@ func TestResponsesStreamConverter_FinalOutputKeepsStreamedItemOrder(t *testing.T } } } + +func customApplyPatchTool() ResponsesTool { + return ResponsesTool{ + Type: "custom", + Name: "apply_patch", + Format: json.RawMessage(`{"type":"grammar","syntax":"lark","definition":"start: /.+/"}`), + } +} + +func customApplyPatchCall(id, patch string) api.ToolCall { + return api.ToolCall{ + ID: id, + Function: api.ToolCallFunction{ + Name: "apply_patch", + Arguments: testArgs(map[string]any{"input": patch}), + }, + } +} + +func TestFromResponsesRequest_CustomApplyPatchDeclarationAndHistory(t *testing.T) { + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + request := ResponsesRequest{ + Tools: []ResponsesTool{ + customApplyPatchTool(), + responsesFunction("get_weather"), + responsesNamespace("editor", responsesFunction("apply_patch")), + {Type: "web_search"}, + }, + Input: ResponsesInput{Items: []ResponsesInputItem{ + ResponsesReasoningInput{EncryptedContent: "thinking"}, + ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "call_patch", Name: "apply_patch", Input: patch}, + ResponsesCustomToolCallOutput{Type: "custom_tool_call_output", CallID: "call_patch", Output: "Done"}, + }}, + } + + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if len(chat.Tools) != 4 { + t.Fatalf("tool count = %d, want 4: %#v", len(chat.Tools), chat.Tools) + } + wantNames := []string{"apply_patch", "get_weather", "editor__apply_patch", "web_search"} + for i, want := range wantNames { + if got := chat.Tools[i].Function.Name; got != want { + t.Fatalf("tool[%d].name = %q, want %q", i, got, want) + } + } + custom := chat.Tools[0] + input, ok := custom.Function.Parameters.Properties.Get("input") + if custom.Type != "function" || !ok || input.Type.String() != "string" || len(custom.Function.Parameters.Required) != 1 || custom.Function.Parameters.Required[0] != "input" { + t.Fatalf("custom schema = %#v", custom) + } + if len(chat.Messages) != 2 || chat.Messages[0].Thinking != "thinking" || len(chat.Messages[0].ToolCalls) != 1 { + t.Fatalf("history messages = %#v", chat.Messages) + } + call := chat.Messages[0].ToolCalls[0] + gotInput, _ := call.Function.Arguments.Get("input") + if call.ID != "call_patch" || call.Function.Name != "apply_patch" || gotInput != patch { + t.Fatalf("history custom call = %#v", call) + } + if chat.Messages[1].Role != "tool" || chat.Messages[1].ToolCallID != "call_patch" || chat.Messages[1].Content != "Done" { + t.Fatalf("history custom output = %#v", chat.Messages[1]) + } +} + +func TestFromResponsesRequest_CustomApplyPatchCollisions(t *testing.T) { + functionDeclaration := responsesFunction("apply_patch") + customDeclaration := customApplyPatchTool() + functionHistory := ResponsesFunctionCall{Type: "function_call", CallID: "fc", Name: "apply_patch", Arguments: `{}`} + customHistory := ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "ctc", Name: "apply_patch", Input: "patch"} + + tests := []struct { + name string + tools []ResponsesTool + items []ResponsesInputItem + }{ + {name: "duplicate custom", tools: []ResponsesTool{customDeclaration, customDeclaration}}, + {name: "function then custom declarations", tools: []ResponsesTool{functionDeclaration, customDeclaration}}, + {name: "custom then function declarations", tools: []ResponsesTool{customDeclaration, functionDeclaration}}, + {name: "function then custom histories", items: []ResponsesInputItem{functionHistory, customHistory}}, + {name: "custom then function histories", items: []ResponsesInputItem{customHistory, functionHistory}}, + {name: "function declaration custom history", tools: []ResponsesTool{functionDeclaration}, items: []ResponsesInputItem{customHistory}}, + {name: "custom declaration function history", tools: []ResponsesTool{customDeclaration}, items: []ResponsesInputItem{functionHistory}}, + {name: "nested custom", tools: []ResponsesTool{responsesNamespace("editor", customDeclaration)}}, + {name: "unsupported declaration", tools: []ResponsesTool{{Type: "custom", Name: "shell"}}}, + {name: "unsupported history", items: []ResponsesInputItem{ResponsesCustomToolCall{Type: "custom_tool_call", Name: "shell"}}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if chat, err := FromResponsesRequest(ResponsesRequest{Tools: test.tools, Input: ResponsesInput{Items: test.items}}); err == nil || chat != nil { + t.Fatalf("FromResponsesRequest() = (%#v, %v), want pre-dispatch rejection", chat, err) + } + }) + } + + left := ResponsesRequest{Tools: []ResponsesTool{functionDeclaration, customDeclaration}} + right := ResponsesRequest{Tools: []ResponsesTool{customDeclaration, functionDeclaration}} + _, leftErr := FromResponsesRequest(left) + _, rightErr := FromResponsesRequest(right) + if leftErr.Error() != rightErr.Error() { + t.Fatalf("declaration-order errors differ:\nleft: %v\nright: %v", leftErr, rightErr) + } + historyLeft := ResponsesRequest{Input: ResponsesInput{Items: []ResponsesInputItem{functionHistory, customHistory}}} + historyRight := ResponsesRequest{Input: ResponsesInput{Items: []ResponsesInputItem{customHistory, functionHistory}}} + _, leftErr = FromResponsesRequest(historyLeft) + _, rightErr = FromResponsesRequest(historyRight) + if leftErr.Error() != rightErr.Error() { + t.Fatalf("history-order errors differ:\nleft: %v\nright: %v", leftErr, rightErr) + } + + accepted := []ResponsesRequest{ + {Tools: []ResponsesTool{customDeclaration, responsesNamespace("editor", responsesFunction("apply_patch"))}}, + {Tools: []ResponsesTool{responsesNamespace("editor", responsesFunction("apply_patch")), customDeclaration}}, + {Tools: []ResponsesTool{customDeclaration}, Input: ResponsesInput{Items: []ResponsesInputItem{customHistory, ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "ctc2", Name: "apply_patch"}}}}, + } + for i, request := range accepted { + if _, err := FromResponsesRequest(request); err != nil { + t.Fatalf("accepted request[%d] rejected: %v", i, err) + } + } + + unrelated := responsesFunction("get_weather") + duplicateOrders := [][]ResponsesTool{ + {customDeclaration, unrelated, customDeclaration, responsesFunction("other")}, + {responsesFunction("other"), customDeclaration, unrelated, customDeclaration}, + } + var duplicateError string + for i, tools := range duplicateOrders { + chat, err := FromResponsesRequest(ResponsesRequest{Tools: tools}) + if err == nil || chat != nil { + t.Fatalf("duplicate order[%d] = (%#v, %v), want rejection", i, chat, err) + } + if i == 0 { + duplicateError = err.Error() + } else if err.Error() != duplicateError { + t.Fatalf("duplicate order errors differ:\nfirst: %s\nsecond: %s", duplicateError, err) + } + } +} + +func TestFromResponsesRequest_CustomApplyPatchAdjacentAssistantMerge(t *testing.T) { + requestJSON := `{ + "input": [ + {"type":"custom_tool_call","call_id":"call_patch","name":"apply_patch","input":"*** Begin Patch\n*** End Patch\n"}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"patch prepared"}]}, + {"type":"custom_tool_call_output","call_id":"call_patch","output":"Done"} + ] + }` + var request ResponsesRequest + if err := json.Unmarshal([]byte(requestJSON), &request); err != nil { + t.Fatal(err) + } + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 2 || chat.Messages[0].Role != "assistant" || chat.Messages[0].Content != "patch prepared" || len(chat.Messages[0].ToolCalls) != 1 { + t.Fatalf("adjacent assistant state = %#v", chat.Messages) + } + if chat.Messages[1].Role != "tool" || chat.Messages[1].ToolCallID != "call_patch" || chat.Messages[1].Content != "Done" { + t.Fatalf("custom output message = %#v", chat.Messages[1]) + } +} + +func TestFromResponsesRequest_RepeatedCustomApplyPatchHistoryPreservesCalls(t *testing.T) { + patches := []string{ + "*** Begin Patch\n*** Add File: one.txt\n+one\n*** End Patch\n", + "*** Begin Patch\n*** Add File: two.txt\n+two\n*** End Patch\n\n", + } + request := ResponsesRequest{ + Tools: []ResponsesTool{customApplyPatchTool()}, + Input: ResponsesInput{Items: []ResponsesInputItem{ + ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "call_one", Name: "apply_patch", Input: patches[0]}, + ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "call_two", Name: "apply_patch", Input: patches[1]}, + }}, + } + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 1 || len(chat.Messages[0].ToolCalls) != 2 { + t.Fatalf("repeated history messages = %#v", chat.Messages) + } + for i, call := range chat.Messages[0].ToolCalls { + input, ok := call.Function.Arguments.Get("input") + wantID := []string{"call_one", "call_two"}[i] + if call.ID != wantID || call.Function.Name != "apply_patch" || !ok || input != patches[i] { + t.Fatalf("history call[%d] = %#v input=%#v, want id=%q input=%q", i, call, input, wantID, patches[i]) + } + } +} + +func TestResponsesCustomApplyPatchNamespaceReverseOutputBothDeclarationOrders(t *testing.T) { + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + custom := customApplyPatchTool() + namespace := responsesNamespace("editor", responsesFunction("apply_patch")) + for _, test := range []struct { + name string + tools []ResponsesTool + }{ + {name: "custom then namespace", tools: []ResponsesTool{custom, namespace}}, + {name: "namespace then custom", tools: []ResponsesTool{namespace, custom}}, + } { + t.Run(test.name, func(t *testing.T) { + request := ResponsesRequest{Tools: test.tools} + if _, err := FromResponsesRequest(request); err != nil { + t.Fatal(err) + } + items := ResponsesFunctionCallOutputItems(request, "fc_reverse_", []api.ToolCall{ + customApplyPatchCall("call_custom", patch), + {ID: "call_namespaced", Function: api.ToolCallFunction{Name: "editor__apply_patch", Arguments: testArgs(map[string]any{"path": "file.txt"})}}, + }) + if len(items) != 2 { + t.Fatalf("output items = %#v", items) + } + if item := items[0]; item.ID != "ctc_reverse_0" || item.Type != "custom_tool_call" || item.CallID != "call_custom" || item.Name != "apply_patch" || item.Namespace != "" || item.Input != patch || item.Arguments != "" { + t.Fatalf("custom reverse item = %#v", item) + } + if item := items[1]; item.ID != "fc_reverse_1" || item.Type != "function_call" || item.CallID != "call_namespaced" || item.Name != "apply_patch" || item.Namespace != "editor" || item.Arguments != `{"path":"file.txt"}` || item.Input != "" { + t.Fatalf("namespace reverse item = %#v", item) + } + }) + } +} + +func TestResponsesCustomApplyPatchOutputClassification(t *testing.T) { + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n\n" + request := ResponsesRequest{Tools: []ResponsesTool{ + customApplyPatchTool(), responsesFunction("get_weather"), responsesNamespace("editor", responsesFunction("write")), + }} + items := ResponsesFunctionCallOutputItems(request, "fc_test_", []api.ToolCall{ + customApplyPatchCall("call_patch", patch), + {ID: "call_flat", Function: api.ToolCallFunction{Name: "get_weather", Arguments: testArgs(map[string]any{"city": "SF"})}}, + {ID: "call_ns", Function: api.ToolCallFunction{Name: "editor__write", Arguments: testArgs(map[string]any{"path": "a"})}}, + }) + if len(items) != 3 || items[0].Type != "custom_tool_call" || items[0].ID != "ctc_test_0" || items[0].CallID != "call_patch" || items[0].Name != "apply_patch" || items[0].Input != patch || items[0].Arguments != "" { + t.Fatalf("custom item = %#v", items) + } + if items[1].Type != "function_call" || items[1].Name != "get_weather" || items[1].Namespace != "" || items[2].Name != "write" || items[2].Namespace != "editor" { + t.Fatalf("mixed output identities = %#v", items) + } + + malformed := []struct { + call api.ToolCall + wantArguments string + }{ + {call: api.ToolCall{ID: "missing", Function: api.ToolCallFunction{Name: "apply_patch", Arguments: testArgs(map[string]any{})}}, wantArguments: `{}`}, + {call: api.ToolCall{ID: "number", Function: api.ToolCallFunction{Name: "apply_patch", Arguments: testArgs(map[string]any{"input": 1})}}, wantArguments: `{"input":1}`}, + {call: customApplyPatchCall("prefix", "garbage\n"+patch), wantArguments: `{"input":"garbage\n*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n\n"}`}, + {call: customApplyPatchCall("begin", "*** Begin Patch"), wantArguments: `{"input":"*** Begin Patch"}`}, + {call: customApplyPatchCall("end", "*** Begin Patch\nx"), wantArguments: `{"input":"*** Begin Patch\nx"}`}, + {call: customApplyPatchCall("trailing", "*** Begin Patch\nx\n*** End Patch\ngarbage"), wantArguments: `{"input":"*** Begin Patch\nx\n*** End Patch\ngarbage"}`}, + } + for _, test := range malformed { + got := ResponsesFunctionCallOutputItems(request, "fc_bad_", []api.ToolCall{test.call}) + if len(got) != 1 || got[0].Type != "function_call" || got[0].Arguments != test.wantArguments { + t.Fatalf("malformed %q output = %#v, want exact arguments %q", test.call.ID, got, test.wantArguments) + } + } + + invalidRequest := ResponsesRequest{Tools: []ResponsesTool{customApplyPatchTool(), responsesFunction("apply_patch")}} + fallback := ResponsesFunctionCallOutputItems(invalidRequest, "fc_fallback_", []api.ToolCall{customApplyPatchCall("call", patch)}) + if len(fallback) != 1 || fallback[0].Type != "function_call" || fallback[0].Name != "apply_patch" || fallback[0].Arguments == "" { + t.Fatalf("invalid-resolver fallback = %#v", fallback) + } +} + +func TestResponsesStreamConverter_CustomApplyPatchFallbacks(t *testing.T) { + validPatch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + tests := []struct { + name string + request ResponsesRequest + call api.ToolCall + want string + }{ + { + name: "invalid flat custom resolver", + request: ResponsesRequest{Tools: []ResponsesTool{ + customApplyPatchTool(), responsesFunction("apply_patch"), + }}, + call: customApplyPatchCall("call_invalid_resolver", validPatch), + want: `{"input":"*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n"}`, + }, + { + name: "malformed patch", + request: ResponsesRequest{Tools: []ResponsesTool{customApplyPatchTool()}}, + call: customApplyPatchCall("call_malformed", validPatch+"trailing garbage"), + want: `{"input":"*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\ntrailing garbage"}`, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + converter := NewResponsesStreamConverter("resp_fallback", "msg_fallback", "test-model", test.request) + events := converter.Process(api.ChatResponse{Message: api.Message{ToolCalls: []api.ToolCall{test.call}}, Done: true}) + wantArguments := test.want + + var added, argumentDelta, argumentsDone, outputDone map[string]any + var terminal []any + wantEvents := []string{ + "response.created", + "response.in_progress", + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + "response.completed", + } + if len(events) != len(wantEvents) { + t.Fatalf("fallback event count = %d, want %d: %#v", len(events), len(wantEvents), events) + } + for i, event := range events { + if event.Event != wantEvents[i] { + t.Fatalf("fallback event[%d] = %q, want %q: %#v", i, event.Event, wantEvents[i], events) + } + data := event.Data.(map[string]any) + if strings.HasPrefix(event.Event, "response.custom_tool_call_input.") { + t.Fatalf("unexpected custom event %q: %#v", event.Event, data) + } + switch event.Event { + case "response.output_item.added": + item := data["item"].(map[string]any) + if item["type"] == "function_call" { + added = data + } + case "response.function_call_arguments.delta": + argumentDelta = data + case "response.function_call_arguments.done": + argumentsDone = data + case "response.output_item.done": + item := data["item"].(map[string]any) + if item["type"] == "function_call" { + outputDone = data + } + case "response.completed": + terminal = data["response"].(map[string]any)["output"].([]any) + } + } + + if added == nil || argumentDelta == nil || argumentsDone == nil || outputDone == nil || len(terminal) != 1 { + t.Fatalf("incomplete function fallback lifecycle: added=%#v delta=%#v argsDone=%#v outputDone=%#v terminal=%#v", added, argumentDelta, argumentsDone, outputDone, terminal) + } + addedItem := added["item"].(map[string]any) + doneItem := outputDone["item"].(map[string]any) + terminalItem := terminal[0].(map[string]any) + if added["output_index"] != 0 || argumentDelta["output_index"] != 0 || argumentsDone["output_index"] != 0 || outputDone["output_index"] != 0 { + t.Fatalf("fallback indexes: added=%#v delta=%#v argsDone=%#v outputDone=%#v", added, argumentDelta, argumentsDone, outputDone) + } + if addedItem["id"] != argumentDelta["item_id"] || addedItem["id"] != argumentsDone["item_id"] || addedItem["id"] != doneItem["id"] || doneItem["id"] != terminalItem["id"] { + t.Fatalf("fallback item identity mismatch: added=%#v done=%#v terminal=%#v", addedItem, doneItem, terminalItem) + } + for label, item := range map[string]map[string]any{"added": addedItem, "output done": doneItem, "terminal": terminalItem} { + if item["type"] != "function_call" || item["call_id"] != test.call.ID || item["name"] != "apply_patch" { + t.Fatalf("fallback %s identity = %#v", label, item) + } + if _, ok := item["namespace"]; ok { + t.Fatalf("fallback %s invented namespace: %#v", label, item) + } + if _, ok := item["input"]; ok { + t.Fatalf("fallback %s invented custom input: %#v", label, item) + } + } + if addedItem["arguments"] != "" || argumentDelta["delta"] != wantArguments || argumentsDone["arguments"] != wantArguments || doneItem["arguments"] != wantArguments || terminalItem["arguments"] != wantArguments { + t.Fatalf("fallback arguments mismatch: added=%#v delta=%#v argsDone=%#v done=%#v terminal=%#v want=%q", addedItem, argumentDelta, argumentsDone, doneItem, terminalItem, wantArguments) + } + }) + } +} + +func TestResponsesStreamConverter_CustomApplyPatchMixedOrder(t *testing.T) { + patch := "*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n" + request := ResponsesRequest{Tools: []ResponsesTool{ + customApplyPatchTool(), responsesFunction("get_weather"), responsesNamespace("editor", responsesFunction("write")), {Type: "web_search"}, + }} + converter := NewResponsesStreamConverter("resp_custom", "msg_custom", "test-model", request) + var events []ResponsesStreamEvent + events = append(events, converter.Process(api.ChatResponse{Message: api.Message{Thinking: "think"}})...) + search := ResponsesWebSearchCall{ID: "ws_1", Type: "web_search_call", Status: "completed", Action: &ResponsesWebSearchAction{Type: "search", Query: "q"}} + index, searchEvents := converter.StartWebSearchCall(search) + events = append(events, searchEvents...) + events = append(events, converter.FinishWebSearchCall(search, index)...) + events = append(events, converter.EmitFunctionCallItems([]api.ToolCall{ + customApplyPatchCall("call_patch", patch), + {ID: "call_flat", Function: api.ToolCallFunction{Name: "get_weather", Arguments: testArgs(map[string]any{})}}, + {ID: "call_ns", Function: api.ToolCallFunction{Name: "editor__write", Arguments: testArgs(map[string]any{})}}, + })...) + events = append(events, converter.Process(api.ChatResponse{Done: true})...) + + var customAdded, customDelta, customInputDone bool + doneByIndex := map[int]map[string]any{} + var final []any + for _, event := range events { + data := event.Data.(map[string]any) + switch event.Event { + case "response.output_item.added": + item := data["item"].(map[string]any) + customAdded = customAdded || item["type"] == "custom_tool_call" && item["input"] == "" + case "response.custom_tool_call_input.delta": + customDelta = data["delta"] == patch + case "response.custom_tool_call_input.done": + customInputDone = data["input"] == patch + case "response.output_item.done": + doneByIndex[data["output_index"].(int)] = data["item"].(map[string]any) + case "response.completed": + response := data["response"].(map[string]any) + final = response["output"].([]any) + tools := response["tools"].([]any) + if got := tools[0].(map[string]any)["format"].(json.RawMessage); string(got) != string(customApplyPatchTool().Format) { + t.Fatalf("custom format = %s", got) + } + } + } + if !customAdded || !customDelta || !customInputDone || len(final) != 5 { + t.Fatalf("custom stream lifecycle missing: added=%v delta=%v done=%v final=%#v", customAdded, customDelta, customInputDone, final) + } + wantTypes := []string{"reasoning", "web_search_call", "custom_tool_call", "function_call", "function_call"} + for i, raw := range final { + item := raw.(map[string]any) + if item["type"] != wantTypes[i] || item["id"] != doneByIndex[i]["id"] { + t.Fatalf("terminal[%d] = %#v, done = %#v", i, item, doneByIndex[i]) + } + } + if final[2].(map[string]any)["call_id"] != "call_patch" || final[2].(map[string]any)["input"] != patch { + t.Fatalf("terminal custom item = %#v", final[2]) + } + if final[4].(map[string]any)["name"] != "write" || final[4].(map[string]any)["namespace"] != "editor" { + t.Fatalf("terminal namespace item = %#v", final[4]) + } +} From ba5d77220417fdf92356cff82d78faa715de1ce0 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 22 Aug 2026 14:17:07 +0000 Subject: [PATCH 05/58] openai: support Codex custom apply patch Co-authored-by: Codex --- openai/responses.go | 79 +++++++++++++++++++++++++++++------ openai/responses_test.go | 90 ++++++++++++++++++++++++---------------- 2 files changed, 121 insertions(+), 48 deletions(-) diff --git a/openai/responses.go b/openai/responses.go index cbd18c4eea0..5258c70dd70 100644 --- a/openai/responses.go +++ b/openai/responses.go @@ -744,22 +744,24 @@ type responsesToolDeclaration struct { } type responsesToolResolver struct { - byExternal map[responsesToolExternalName]responsesToolName - byInternal map[string]responsesToolName - declarations map[responsesToolDeclaration]struct{} - namespaceChildren map[responsesNamespaceChild]struct{} - nextNamespaceID uint64 - tools []api.Tool - hasWebSearchTool bool + byExternal map[responsesToolExternalName]responsesToolName + byInternal map[string]responsesToolName + declarations map[responsesToolDeclaration]struct{} + namespaceChildren map[responsesNamespaceChild]struct{} + nextNamespaceID uint64 + tools []api.Tool + hasWebSearchTool bool + hasCustomApplyPatch bool } func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error) { resolver := &responsesToolResolver{ - byExternal: make(map[responsesToolExternalName]responsesToolName), - byInternal: make(map[string]responsesToolName), - declarations: make(map[responsesToolDeclaration]struct{}), - namespaceChildren: make(map[responsesNamespaceChild]struct{}), - hasWebSearchTool: HasWebSearchTool(r.Tools), + byExternal: make(map[responsesToolExternalName]responsesToolName), + byInternal: make(map[string]responsesToolName), + declarations: make(map[responsesToolDeclaration]struct{}), + namespaceChildren: make(map[responsesNamespaceChild]struct{}), + hasWebSearchTool: HasWebSearchTool(r.Tools), + hasCustomApplyPatch: hasCustomApplyPatch(r), } for _, tool := range r.Tools { @@ -789,6 +791,20 @@ func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error return resolver, nil } +func hasCustomApplyPatch(r ResponsesRequest) bool { + for _, tool := range r.Tools { + if tool.Type == "custom" && tool.Name == "apply_patch" { + return true + } + } + for _, item := range r.Input.Items { + if call, ok := item.(ResponsesCustomToolCall); ok && call.Name == "apply_patch" { + return true + } + } + return false +} + func joinNamespaceToolName(namespace, name string) string { if strings.HasSuffix(namespace, "__") { return namespace + name @@ -868,6 +884,11 @@ func (r *responsesToolResolver) register(kind responsesToolKind, namespace, name internal := name if namespace != "" { internal = joinNamespaceToolName(namespace, name) + } else if kind == responsesToolKindFunction && name == "apply_patch" && r.hasCustomApplyPatch { + // Keep the custom tool's model-visible name stable for Codex. A client + // may also declare a flat function with that name; disambiguate only the + // model-facing function identity and map it back on Responses output. + internal = "apply_patch__function" } candidate := responsesToolName{external: external, internal: internal, declarationNode: declarationNode, declared: declaration} @@ -954,6 +975,12 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv if t.Type != "function" { return []api.Tool{tool}, nil } + if resolver.hasCustomApplyPatch && namespaceNode == nil && t.Name == "apply_patch" { + // Codex may redundantly send a flat function declaration alongside + // its authoritative custom apply_patch tool. Only expose the custom + // tool to the model; otherwise tool selection becomes ambiguous. + return nil, nil + } // The built-in tool retains upstream ownership of its private model name. if resolver.hasWebSearchTool && namespaceNode == nil && t.Name == "web_search" { @@ -1002,7 +1029,7 @@ func convertTool(t ResponsesTool) (api.Tool, error) { Type: "function", Function: api.ToolFunction{ Name: "apply_patch", - Description: "Apply a patch to files. The input field must contain the complete raw patch text.", + Description: applyPatchToolDescription(t), Parameters: applyPatchFunctionParameters(), }, }, nil @@ -1035,6 +1062,32 @@ func convertTool(t ResponsesTool) (api.Tool, error) { }, nil } +func applyPatchToolDescription(t ResponsesTool) string { + description := "Apply a patch to files. The input field must contain the complete raw patch text." + if t.Description != nil && strings.TrimSpace(*t.Description) != "" { + description = *t.Description + } + + if len(t.Format) == 0 || string(t.Format) == "null" { + return description + } + + var format struct { + Type string `json:"type"` + Syntax string `json:"syntax"` + Definition string `json:"definition"` + } + if err := json.Unmarshal(t.Format, &format); err == nil && + format.Type == "grammar" && format.Syntax == "lark" && + strings.Contains(format.Definition, "*** Begin Patch") && + strings.Contains(format.Definition, "*** Update File:") && + strings.Contains(format.Definition, "*** End Patch") { + return description + "\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines." + } + + return description +} + func applyPatchFunctionParameters() api.ToolFunctionParameters { properties := api.NewToolPropertiesMap() properties.Set("input", api.ToolProperty{ diff --git a/openai/responses_test.go b/openai/responses_test.go index e24bd4e9e47..0f500eacc2f 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -2897,6 +2897,47 @@ func customApplyPatchTool() ResponsesTool { } } +func TestFromResponsesRequest_CustomApplyPatchPreservesInstructions(t *testing.T) { + description := "Use apply_patch to edit files. This is a FREEFORM tool; do not wrap the patch in JSON." + request := ResponsesRequest{Tools: []ResponsesTool{{ + Type: "custom", + Name: "apply_patch", + Description: &description, + Format: json.RawMessage(`{"type":"grammar","syntax":"lark","definition":"start: begin_patch hunk+ end_patch\nbegin_patch: \"*** Begin Patch\" LF\nupdate_hunk: \"*** Update File: \" filename LF\nend_patch: \"*** End Patch\" LF?\n"}`), + }}} + + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if len(chat.Tools) != 1 { + t.Fatalf("tool count = %d, want 1: %#v", len(chat.Tools), chat.Tools) + } + if got, want := chat.Tools[0].Function.Description, description+"\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines."; got != want { + t.Fatalf("custom tool description = %q, want %q", got, want) + } + + legacy, err := FromResponsesRequest(ResponsesRequest{Tools: []ResponsesTool{{Type: "custom", Name: "apply_patch"}}}) + if err != nil { + t.Fatal(err) + } + if got, want := legacy.Tools[0].Function.Description, "Apply a patch to files. The input field must contain the complete raw patch text."; got != want { + t.Fatalf("legacy custom tool description = %q, want %q", got, want) + } + + incompatible := "A different custom grammar" + chat, err = FromResponsesRequest(ResponsesRequest{Tools: []ResponsesTool{{ + Type: "custom", Name: "apply_patch", Description: &incompatible, + Format: json.RawMessage(`{"type":"grammar","syntax":"lark","definition":"start: value\nvalue: /.+/"}`), + }}}) + if err != nil { + t.Fatal(err) + } + if got := chat.Tools[0].Function.Description; got != incompatible { + t.Fatalf("incompatible custom grammar changed description: %q", got) + } +} + func customApplyPatchCall(id, patch string) api.ToolCall { return api.ToolCall{ ID: id, @@ -2966,12 +3007,6 @@ func TestFromResponsesRequest_CustomApplyPatchCollisions(t *testing.T) { items []ResponsesInputItem }{ {name: "duplicate custom", tools: []ResponsesTool{customDeclaration, customDeclaration}}, - {name: "function then custom declarations", tools: []ResponsesTool{functionDeclaration, customDeclaration}}, - {name: "custom then function declarations", tools: []ResponsesTool{customDeclaration, functionDeclaration}}, - {name: "function then custom histories", items: []ResponsesInputItem{functionHistory, customHistory}}, - {name: "custom then function histories", items: []ResponsesInputItem{customHistory, functionHistory}}, - {name: "function declaration custom history", tools: []ResponsesTool{functionDeclaration}, items: []ResponsesInputItem{customHistory}}, - {name: "custom declaration function history", tools: []ResponsesTool{customDeclaration}, items: []ResponsesInputItem{functionHistory}}, {name: "nested custom", tools: []ResponsesTool{responsesNamespace("editor", customDeclaration)}}, {name: "unsupported declaration", tools: []ResponsesTool{{Type: "custom", Name: "shell"}}}, {name: "unsupported history", items: []ResponsesInputItem{ResponsesCustomToolCall{Type: "custom_tool_call", Name: "shell"}}}, @@ -2984,22 +3019,13 @@ func TestFromResponsesRequest_CustomApplyPatchCollisions(t *testing.T) { }) } - left := ResponsesRequest{Tools: []ResponsesTool{functionDeclaration, customDeclaration}} - right := ResponsesRequest{Tools: []ResponsesTool{customDeclaration, functionDeclaration}} - _, leftErr := FromResponsesRequest(left) - _, rightErr := FromResponsesRequest(right) - if leftErr.Error() != rightErr.Error() { - t.Fatalf("declaration-order errors differ:\nleft: %v\nright: %v", leftErr, rightErr) - } - historyLeft := ResponsesRequest{Input: ResponsesInput{Items: []ResponsesInputItem{functionHistory, customHistory}}} - historyRight := ResponsesRequest{Input: ResponsesInput{Items: []ResponsesInputItem{customHistory, functionHistory}}} - _, leftErr = FromResponsesRequest(historyLeft) - _, rightErr = FromResponsesRequest(historyRight) - if leftErr.Error() != rightErr.Error() { - t.Fatalf("history-order errors differ:\nleft: %v\nright: %v", leftErr, rightErr) - } - accepted := []ResponsesRequest{ + {Tools: []ResponsesTool{functionDeclaration, customDeclaration}}, + {Tools: []ResponsesTool{customDeclaration, functionDeclaration}}, + {Input: ResponsesInput{Items: []ResponsesInputItem{functionHistory, customHistory}}}, + {Input: ResponsesInput{Items: []ResponsesInputItem{customHistory, functionHistory}}}, + {Tools: []ResponsesTool{functionDeclaration}, Input: ResponsesInput{Items: []ResponsesInputItem{customHistory}}}, + {Tools: []ResponsesTool{customDeclaration}, Input: ResponsesInput{Items: []ResponsesInputItem{functionHistory}}}, {Tools: []ResponsesTool{customDeclaration, responsesNamespace("editor", responsesFunction("apply_patch"))}}, {Tools: []ResponsesTool{responsesNamespace("editor", responsesFunction("apply_patch")), customDeclaration}}, {Tools: []ResponsesTool{customDeclaration}, Input: ResponsesInput{Items: []ResponsesInputItem{customHistory, ResponsesCustomToolCall{Type: "custom_tool_call", CallID: "ctc2", Name: "apply_patch"}}}}, @@ -3010,6 +3036,14 @@ func TestFromResponsesRequest_CustomApplyPatchCollisions(t *testing.T) { } } + chat, err := FromResponsesRequest(ResponsesRequest{Tools: []ResponsesTool{functionDeclaration, customDeclaration}}) + if err != nil { + t.Fatal(err) + } + if len(chat.Tools) != 1 || chat.Tools[0].Function.Name != "apply_patch" { + t.Fatalf("mixed apply_patch tool identities = %#v", chat.Tools) + } + unrelated := responsesFunction("get_weather") duplicateOrders := [][]ResponsesTool{ {customDeclaration, unrelated, customDeclaration, responsesFunction("other")}, @@ -3148,12 +3182,6 @@ func TestResponsesCustomApplyPatchOutputClassification(t *testing.T) { t.Fatalf("malformed %q output = %#v, want exact arguments %q", test.call.ID, got, test.wantArguments) } } - - invalidRequest := ResponsesRequest{Tools: []ResponsesTool{customApplyPatchTool(), responsesFunction("apply_patch")}} - fallback := ResponsesFunctionCallOutputItems(invalidRequest, "fc_fallback_", []api.ToolCall{customApplyPatchCall("call", patch)}) - if len(fallback) != 1 || fallback[0].Type != "function_call" || fallback[0].Name != "apply_patch" || fallback[0].Arguments == "" { - t.Fatalf("invalid-resolver fallback = %#v", fallback) - } } func TestResponsesStreamConverter_CustomApplyPatchFallbacks(t *testing.T) { @@ -3164,14 +3192,6 @@ func TestResponsesStreamConverter_CustomApplyPatchFallbacks(t *testing.T) { call api.ToolCall want string }{ - { - name: "invalid flat custom resolver", - request: ResponsesRequest{Tools: []ResponsesTool{ - customApplyPatchTool(), responsesFunction("apply_patch"), - }}, - call: customApplyPatchCall("call_invalid_resolver", validPatch), - want: `{"input":"*** Begin Patch\n*** Add File: file.txt\n+new\n*** End Patch\n"}`, - }, { name: "malformed patch", request: ResponsesRequest{Tools: []ResponsesTool{customApplyPatchTool()}}, From 71494cb104997d0e8b69da86e7abbb2a894e41be Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 22 Aug 2026 15:26:57 +0000 Subject: [PATCH 06/58] openai: improve apply patch guidance for non-GPT-OSS models Co-authored-by: Codex --- openai/responses.go | 25 ++++++++++++++++++------- openai/responses_test.go | 22 ++++++++++++++++++++++ 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/openai/responses.go b/openai/responses.go index 5258c70dd70..7cdf75eb0d9 100644 --- a/openai/responses.go +++ b/openai/responses.go @@ -752,6 +752,7 @@ type responsesToolResolver struct { tools []api.Tool hasWebSearchTool bool hasCustomApplyPatch bool + applyPatchExample bool } func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error) { @@ -762,6 +763,7 @@ func newResponsesToolResolver(r ResponsesRequest) (*responsesToolResolver, error namespaceChildren: make(map[responsesNamespaceChild]struct{}), hasWebSearchTool: HasWebSearchTool(r.Tools), hasCustomApplyPatch: hasCustomApplyPatch(r), + applyPatchExample: applyPatchExampleForModel(r.Model), } for _, tool := range r.Tools { @@ -805,6 +807,11 @@ func hasCustomApplyPatch(r ResponsesRequest) bool { return false } +func applyPatchExampleForModel(model string) bool { + model = strings.NewReplacer("-", "", "_", "").Replace(strings.ToLower(strings.TrimSpace(model))) + return model != "" && !strings.HasPrefix(model, "gptoss") +} + func joinNamespaceToolName(namespace, name string) string { if strings.HasSuffix(namespace, "__") { return namespace + name @@ -955,7 +962,7 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv if t.Name != "apply_patch" { return nil, fmt.Errorf("unsupported responses custom tool %q", t.Name) } - tool, err := convertTool(t) + tool, err := convertTool(t, resolver.applyPatchExample) if err != nil { return nil, err } @@ -968,7 +975,7 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv } if t.Type != "namespace" { - tool, err := convertTool(t) + tool, err := convertTool(t, resolver.applyPatchExample) if err != nil { return nil, err } @@ -1023,13 +1030,13 @@ func convertTools(t ResponsesTool, namespaceNode *responsesNamespaceNode, resolv return tools, nil } -func convertTool(t ResponsesTool) (api.Tool, error) { +func convertTool(t ResponsesTool, applyPatchExample bool) (api.Tool, error) { if t.Type == "custom" && t.Name == "apply_patch" { return api.Tool{ Type: "function", Function: api.ToolFunction{ Name: "apply_patch", - Description: applyPatchToolDescription(t), + Description: applyPatchToolDescription(t, applyPatchExample), Parameters: applyPatchFunctionParameters(), }, }, nil @@ -1062,7 +1069,7 @@ func convertTool(t ResponsesTool) (api.Tool, error) { }, nil } -func applyPatchToolDescription(t ResponsesTool) string { +func applyPatchToolDescription(t ResponsesTool, includeExample bool) string { description := "Apply a patch to files. The input field must contain the complete raw patch text." if t.Description != nil && strings.TrimSpace(*t.Description) != "" { description = *t.Description @@ -1082,7 +1089,11 @@ func applyPatchToolDescription(t ResponsesTool) string { strings.Contains(format.Definition, "*** Begin Patch") && strings.Contains(format.Definition, "*** Update File:") && strings.Contains(format.Definition, "*** End Patch") { - return description + "\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines." + instructions := "\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines." + if includeExample { + instructions += "\n\nExample complete input:\n*** Begin Patch\n*** Update File: path/to/file\n@@\n-old text\n+new text\n*** End Patch" + } + return description + instructions } return description @@ -1092,7 +1103,7 @@ func applyPatchFunctionParameters() api.ToolFunctionParameters { properties := api.NewToolPropertiesMap() properties.Set("input", api.ToolProperty{ Type: api.PropertyType{"string"}, - Description: "Raw patch text beginning with *** Begin Patch and ending with *** End Patch.", + Description: "Complete raw patch text. Use the custom patch format and example in this tool's description.", }) return api.ToolFunctionParameters{ Type: "object", diff --git a/openai/responses_test.go b/openai/responses_test.go index 0f500eacc2f..21593a31514 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -2917,6 +2917,28 @@ func TestFromResponsesRequest_CustomApplyPatchPreservesInstructions(t *testing.T t.Fatalf("custom tool description = %q, want %q", got, want) } + withExample := description + "\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines.\n\nExample complete input:\n*** Begin Patch\n*** Update File: path/to/file\n@@\n-old text\n+new text\n*** End Patch" + for _, model := range []string{"gemma4:12b-mlx", "qwen3:8b"} { + request.Model = model + chat, err = FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if got := chat.Tools[0].Function.Description; got != withExample { + t.Fatalf("%s custom tool description = %q, want %q", model, got, withExample) + } + } + for _, model := range []string{"gptoss-mlx:20b-mxfp4", "gpt-oss:20b"} { + request.Model = model + chat, err = FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if got := chat.Tools[0].Function.Description; got != description+"\n\nFor the custom Lark patch format, emit only raw patch text: begin with *** Begin Patch; use *** Update File: , then a plain @@ line (never a numbered unified-diff header such as @@ -1,3 +1,3 @@ and never ---/+++ file headers), then -old and +new lines; finish with *** End Patch. Every patch control and hunk line must start in column 1; never indent it. In an update hunk, prefix unchanged context lines with one space, and make - and + the first character of removed and added lines." { + t.Fatalf("%s custom tool description = %q, want no example", model, got) + } + } + legacy, err := FromResponsesRequest(ResponsesRequest{Tools: []ResponsesTool{{Type: "custom", Name: "apply_patch"}}}) if err != nil { t.Fatal(err) From fd443f395c12931db7487b679ea90c7e13830e27 Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 06:38:45 +0000 Subject: [PATCH 07/58] fix(openai): align apply patch schema guidance Co-authored-by: Codex --- openai/responses.go | 2 +- openai/responses_test.go | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/openai/responses.go b/openai/responses.go index 7cdf75eb0d9..f057a345408 100644 --- a/openai/responses.go +++ b/openai/responses.go @@ -1103,7 +1103,7 @@ func applyPatchFunctionParameters() api.ToolFunctionParameters { properties := api.NewToolPropertiesMap() properties.Set("input", api.ToolProperty{ Type: api.PropertyType{"string"}, - Description: "Complete raw patch text. Use the custom patch format and example in this tool's description.", + Description: "Complete raw patch text. Use the custom patch format guidance in this tool's description.", }) return api.ToolFunctionParameters{ Type: "object", diff --git a/openai/responses_test.go b/openai/responses_test.go index 21593a31514..3a3e1e00539 100644 --- a/openai/responses_test.go +++ b/openai/responses_test.go @@ -2960,6 +2960,40 @@ func TestFromResponsesRequest_CustomApplyPatchPreservesInstructions(t *testing.T } } +func TestFromResponsesRequest_CustomApplyPatchSchemaGuidance(t *testing.T) { + request := ResponsesRequest{Tools: []ResponsesTool{{ + Type: "custom", + Name: "apply_patch", + Format: json.RawMessage(`{"type":"grammar","syntax":"lark","definition":"start: begin_patch hunk+ end_patch\nbegin_patch: \"*** Begin Patch\" LF\nupdate_hunk: \"*** Update File: \" filename LF\nend_patch: \"*** End Patch\" LF?\n"}`), + }}} + for _, tt := range []struct { + name string + model string + wantExample bool + }{ + {name: "non GPT OSS includes example", model: "gemma4:12b-mlx", wantExample: true}, + {name: "GPT OSS omits example", model: "gpt-oss:20b", wantExample: false}, + } { + t.Run(tt.name, func(t *testing.T) { + request.Model = tt.model + chat, err := FromResponsesRequest(request) + if err != nil { + t.Fatal(err) + } + if got := strings.Contains(chat.Tools[0].Function.Description, "Example complete input:"); got != tt.wantExample { + t.Fatalf("example presence = %v, want %v: %q", got, tt.wantExample, chat.Tools[0].Function.Description) + } + input, ok := chat.Tools[0].Function.Parameters.Properties.Get("input") + if !ok { + t.Fatal("missing apply_patch input schema") + } + if got, want := input.Description, "Complete raw patch text. Use the custom patch format guidance in this tool's description."; got != want { + t.Fatalf("input schema description = %q, want %q", got, want) + } + }) + } +} + func customApplyPatchCall(id, patch string) api.ToolCall { return api.ToolCall{ ID: id, From 9979c218d8906611019daffd1301dc9ca26eb6c8 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 14:23:00 +0000 Subject: [PATCH 08/58] fix(parser): reject unsafe Gemma4 tool calls Co-authored-by: Codex --- model/parsers/gemma4.go | 73 +++++++++--- model/parsers/gemma4_test.go | 213 +++++++++++++++++++++++++++++++++++ 2 files changed, 271 insertions(+), 15 deletions(-) diff --git a/model/parsers/gemma4.go b/model/parsers/gemma4.go index d944c4e952d..555a16d1757 100644 --- a/model/parsers/gemma4.go +++ b/model/parsers/gemma4.go @@ -3,6 +3,7 @@ package parsers import ( "encoding/json" "errors" + "fmt" "log/slog" "regexp" "strings" @@ -104,13 +105,14 @@ type gemma4EventContent struct { content string } -type gemma4EventToolCall struct { - toolCall api.ToolCall +type gemma4EventRawToolCall struct { + raw string + parseErr error } func (gemma4EventThinkingContent) isGemma4Event() {} func (gemma4EventContent) isGemma4Event() {} -func (gemma4EventToolCall) isGemma4Event() {} +func (gemma4EventRawToolCall) isGemma4Event() {} func (p *Gemma4Parser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) { p.buffer.WriteString(s) @@ -121,8 +123,17 @@ func (p *Gemma4Parser) Add(s string, done bool) (content string, thinking string var thinkingSb strings.Builder for _, event := range events { switch event := event.(type) { - case gemma4EventToolCall: - toolCalls = append(toolCalls, event.toolCall) + case gemma4EventRawToolCall: + if event.parseErr != nil { + slog.Warn("gemma4 tool call parsing failed", "error", event.parseErr, "content", event.raw) + return "", "", nil, event.parseErr + } + toolCall, err := parseGemma4ToolCall(event.raw, p.tools) + if err != nil { + slog.Warn("gemma4 tool call parsing failed", "error", err, "content", event.raw) + return "", "", nil, err + } + toolCalls = append(toolCalls, toolCall) case gemma4EventThinkingContent: if p.thinkingEnabled { thinkingSb.WriteString(event.content) @@ -172,6 +183,10 @@ func (p *Gemma4Parser) eat(done bool) ([]gemma4Event, bool) { var events []gemma4Event bufStr := p.buffer.String() if bufStr == "" { + if done && p.state == Gemma4CollectingToolCall { + p.state = Gemma4CollectingContent + return []gemma4Event{gemma4EventRawToolCall{}}, false + } return events, false } @@ -315,24 +330,23 @@ func (p *Gemma4Parser) eat(done bool) ([]gemma4Event, bool) { p.buffer.WriteString(remaining) p.state = Gemma4IgnoringPostToolCallNoise - if toolCall, err := parseGemma4ToolCall(toolCallContent, p.tools); err == nil { - events = append(events, gemma4EventToolCall{toolCall: toolCall}) - } else { - slog.Warn("gemma4 tool call parsing failed", "error", err, "content", toolCallContent) + event := gemma4EventRawToolCall{raw: toolCallContent} + switch { + case strings.Count(toolCallContent, gemma4StringDelimiter)%2 != 0: + event.parseErr = fmt.Errorf("unexpected Gemma control token %q while a tool argument string is open", gemma4ToolCallCloseTag) + case !strings.HasSuffix(strings.TrimSpace(toolCallContent), "}"): + event.parseErr = fmt.Errorf("unexpected Gemma control token %q while tool call arguments are incomplete", gemma4ToolCallCloseTag) } + events = append(events, event) return events, true } // If done, flush any accumulated tool call content even without closing tag. // The model may hit a stop token before emitting . - if done && len(bufStr) > 0 { + if done { p.buffer.Reset() p.state = Gemma4CollectingContent - if toolCall, err := parseGemma4ToolCall(bufStr, p.tools); err == nil { - events = append(events, gemma4EventToolCall{toolCall: toolCall}) - } else { - slog.Warn("gemma4 tool call flush on done failed", "error", err, "content", bufStr) - } + events = append(events, gemma4EventRawToolCall{raw: bufStr}) return events, false } @@ -391,6 +405,10 @@ func (p *Gemma4Parser) eat(done bool) ([]gemma4Event, bool) { // parseGemma4ToolCall parses a tool call in Gemma 4 format: // call:NAME{key:value,key:value} func parseGemma4ToolCall(content string, tools []api.Tool) (api.ToolCall, error) { + if err := validateGemma4ToolCallContent(content); err != nil { + return api.ToolCall{}, err + } + // Expected format: call:NAME{args} if !strings.HasPrefix(content, "call:") { return api.ToolCall{}, errors.New("expected 'call:' prefix") @@ -426,6 +444,31 @@ func parseGemma4ToolCall(content string, tools []api.Tool) (api.ToolCall, error) }, nil } +// validateGemma4ToolCallContent rejects parser grammar tokens that cannot be +// structurally valid inside a tool call. In particular, do not repair or strip +// these tokens from executable arguments: doing so could change the command +// while still leaving it valid enough for a client to execute. +func validateGemma4ToolCallContent(content string) error { + for index, r := range content { + if unicode.IsControl(r) && r != '\n' && r != '\r' && r != '\t' { + return fmt.Errorf("unexpected control character %U at byte %d inside tool call", r, index) + } + } + + for _, token := range []string{ + gemma4ThinkingOpenTag, + gemma4ThinkingCloseTag, + gemma4ToolCallOpenTag, + gemma4ToolCallCloseTag, + gemma4ToolResponseTag, + } { + if index := strings.Index(content, token); index >= 0 { + return fmt.Errorf("unexpected Gemma control token %q at byte %d inside tool call", token, index) + } + } + return nil +} + // gemma4ArgsToJSON converts Gemma 4's custom argument format to valid JSON. func gemma4ArgsToJSON(s string) string { var quotedStrings []string diff --git a/model/parsers/gemma4_test.go b/model/parsers/gemma4_test.go index 27a0f0522f9..832130b2089 100644 --- a/model/parsers/gemma4_test.go +++ b/model/parsers/gemma4_test.go @@ -589,6 +589,160 @@ func TestGemma4Parser_StreamingToolCall(t *testing.T) { } } +func TestGemma4Parser_RejectsControlTokenInsideStreamingToolCall(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init([]api.Tool{gemma4TestStringTool("exec_command", "cmd")}, nil, &api.ThinkValue{Value: true}) + + chunks := []string{ + `<|tool_call>call:exec_command{cmd:<|"|>cat << 'EOF' > project_analysis.py +`, + `print('ok') +"<|chan`, + `nel>EOF +<|"|>}`, + } + + for i, chunk := range chunks { + _, _, calls, err := parser.Add(chunk, i == len(chunks)-1) + if i < len(chunks)-1 { + if err != nil { + t.Fatalf("Add() chunk %d returned an early error: %v", i, err) + } + if len(calls) != 0 { + t.Fatalf("Add() chunk %d returned tool calls before the closing tag: %#v", i, calls) + } + continue + } + + if err == nil { + t.Fatal("Add() accepted a channel control token inside a tool call") + } + if !strings.Contains(err.Error(), gemma4ThinkingOpenTag) { + t.Fatalf("Add() error %q does not identify %q", err, gemma4ThinkingOpenTag) + } + if len(calls) != 0 { + t.Fatalf("Add() returned an invalid tool call: %#v", calls) + } + } +} + +func TestGemma4Parser_RejectsControlCharacterInsideStreamingToolCall(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init([]api.Tool{gemma4TestStringTool("exec_command", "cmd")}, nil, &api.ThinkValue{Value: true}) + + chunks := []string{ + `<|tool_call>call:exec_command{cmd:<|"|>cat << 'EOF' > project_analysis.py +print('ok') +`, + "\x0fEOF\n", + `<|"|>}`, + } + + for i, chunk := range chunks { + _, _, calls, err := parser.Add(chunk, i == len(chunks)-1) + if i < len(chunks)-1 { + if err != nil { + t.Fatalf("Add() chunk %d returned an early error: %v", i, err) + } + if len(calls) != 0 { + t.Fatalf("Add() chunk %d returned tool calls before the closing tag: %#v", i, calls) + } + continue + } + + if err == nil { + t.Fatal("Add() accepted a control character inside a tool call") + } + if !strings.Contains(err.Error(), "U+000F") { + t.Fatalf("Add() error %q does not identify U+000F", err) + } + if len(calls) != 0 { + t.Fatalf("Add() returned an invalid tool call: %#v", calls) + } + } +} + +func TestGemma4Parser_RejectsToolCallCloseInsideOpenArgumentString(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init([]api.Tool{gemma4TestStringTool("exec_command", "cmd")}, nil, &api.ThinkValue{Value: true}) + + input := `<|tool_call>call:exec_command{cmd:<|"|>printf beforeprintf after<|"|>}` + _, _, calls, err := parser.Add(input, true) + if err == nil { + t.Fatal("Add() accepted a tool-call close token inside an open argument string") + } + if !strings.Contains(err.Error(), gemma4ToolCallCloseTag) { + t.Fatalf("Add() error %q does not identify %q", err, gemma4ToolCallCloseTag) + } + if len(calls) != 0 { + t.Fatalf("Add() returned a truncated tool call: %#v", calls) + } +} + +func TestGemma4Parser_RejectsToolCallCloseInsideRawArgument(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init([]api.Tool{gemma4TestStringTool("exec_command", "cmd")}, nil, &api.ThinkValue{Value: true}) + + input := `<|tool_call>call:exec_command{cmd:printf beforeprintf after}` + _, _, calls, err := parser.Add(input, true) + if err == nil { + t.Fatal("Add() accepted a tool-call close token inside an incomplete raw argument") + } + if !strings.Contains(err.Error(), gemma4ToolCallCloseTag) { + t.Fatalf("Add() error %q does not identify %q", err, gemma4ToolCallCloseTag) + } + if len(calls) != 0 { + t.Fatalf("Add() returned a truncated executable tool call: %#v", calls) + } +} + +func TestGemma4Parser_RejectsEmptyToolCallAtEnd(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init(nil, nil, &api.ThinkValue{Value: true}) + + _, _, calls, err := parser.Add(gemma4ToolCallOpenTag, true) + if err == nil { + t.Fatal("Add() silently discarded an empty tool call") + } + if len(calls) != 0 { + t.Fatalf("Add() returned a tool call for empty input: %#v", calls) + } +} + +func TestGemma4Parser_AllowsStreamingMultilineHereDocument(t *testing.T) { + parser := &Gemma4Parser{hasThinkingSupport: true} + parser.Init([]api.Tool{gemma4TestStringTool("exec_command", "cmd")}, nil, &api.ThinkValue{Value: true}) + + command := "cat << 'EOF' > project_analysis.py\nprint('ok')\nEOF\n" + chunks := []string{ + `<|tool_call>call:exec_command{cmd:<|"|>cat << 'EOF' > project_analysis.py +`, + "print('ok')\nEOF\n", + `<|"|>}`, + } + + var calls []api.ToolCall + for i, chunk := range chunks { + _, _, chunkCalls, err := parser.Add(chunk, i == len(chunks)-1) + if err != nil { + t.Fatalf("Add() chunk %d returned error: %v", i, err) + } + calls = append(calls, chunkCalls...) + } + + want := []api.ToolCall{{ + Function: api.ToolCallFunction{ + Name: "exec_command", + Arguments: testArgs(map[string]any{ + "cmd": command, + }), + }, + }} + if diff := cmp.Diff(want, calls, argsComparer); diff != "" { + t.Fatalf("tool calls mismatch (-want +got):\n%s", diff) + } +} + func TestGemma4Parser_IgnoresExtraToolCallCloseTags(t *testing.T) { tests := []struct { name string @@ -1266,6 +1420,65 @@ func TestParseGemma4ToolCall_InvalidRawQuotedEscape(t *testing.T) { } } +func TestParseGemma4ToolCall_RejectsUnexpectedControlTokens(t *testing.T) { + tests := []struct { + name string + token string + }{ + {name: "thinking open", token: gemma4ThinkingOpenTag}, + {name: "thinking close", token: gemma4ThinkingCloseTag}, + {name: "nested tool call", token: gemma4ToolCallOpenTag}, + {name: "tool call close", token: gemma4ToolCallCloseTag}, + {name: "tool response", token: gemma4ToolResponseTag}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + content := `call:exec_command{cmd:<|"|>printf 'before` + tt.token + `after'<|"|>}` + _, err := parseGemma4ToolCall(content, []api.Tool{gemma4TestStringTool("exec_command", "cmd")}) + if err == nil { + t.Fatalf("parseGemma4ToolCall accepted %q inside a tool argument", tt.token) + } + if !strings.Contains(err.Error(), tt.token) { + t.Fatalf("parseGemma4ToolCall error %q does not identify %q", err, tt.token) + } + }) + } +} + +func TestParseGemma4ToolCall_RejectsUnexpectedControlCharacters(t *testing.T) { + content := "call:exec_command{cmd:<|\"|>printf 'before\x0fafter'<|\"|>}" + _, err := parseGemma4ToolCall(content, []api.Tool{gemma4TestStringTool("exec_command", "cmd")}) + if err == nil { + t.Fatal("parseGemma4ToolCall accepted U+000F inside a tool argument") + } + if !strings.Contains(err.Error(), "U+000F") { + t.Fatalf("parseGemma4ToolCall error %q does not identify U+000F", err) + } +} + +func TestParseGemma4ToolCall_AllowsValidMultilineHereDocument(t *testing.T) { + command := "cat << 'EOF' > project_analysis.py\nprint('ok')\nEOF\n" + content := `call:exec_command{cmd:<|"|>` + command + `<|"|>}` + + got, err := parseGemma4ToolCall(content, []api.Tool{gemma4TestStringTool("exec_command", "cmd")}) + if err != nil { + t.Fatalf("parseGemma4ToolCall returned error: %v", err) + } + + want := api.ToolCall{ + Function: api.ToolCallFunction{ + Name: "exec_command", + Arguments: testArgs(map[string]any{ + "cmd": command, + }), + }, + } + if diff := cmp.Diff(want, got, argsComparer); diff != "" { + t.Fatalf("tool call mismatch (-want +got):\n%s", diff) + } +} + func TestParseGemma4ToolCall_QuotedScalarsStayStrings(t *testing.T) { toolCall, err := parseGemma4ToolCall(`call:foo{n:<|"|>1<|"|>,b:<|"|>true<|"|>,z:<|"|>null<|"|>}`, nil) if err != nil { From 052e97b962ac9a5ad9315109741b14cc3f7469f0 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 14:46:15 +0000 Subject: [PATCH 09/58] fix(tokenizer): bound sparse tokenizer IDs Co-authored-by: Codex --- x/tokenizer/tokenizer_load.go | 123 ++++++++++++-- x/tokenizer/tokenizer_load_test.go | 248 +++++++++++++++++++++++++++++ 2 files changed, 355 insertions(+), 16 deletions(-) diff --git a/x/tokenizer/tokenizer_load.go b/x/tokenizer/tokenizer_load.go index c0b6bc25e8e..0852abd380c 100644 --- a/x/tokenizer/tokenizer_load.go +++ b/x/tokenizer/tokenizer_load.go @@ -8,6 +8,31 @@ import ( "strings" ) +const maxTokenizerVocabularySize = 1 << 20 + +type addedToken struct { + ID int32 `json:"id"` + Content string `json:"content"` + Special bool `json:"special"` +} + +func validateTokenizerRecordCount(baseCount, addedCount int) error { + if baseCount > maxTokenizerVocabularySize || addedCount > maxTokenizerVocabularySize-baseCount { + return fmt.Errorf("tokenizer has too many vocabulary records (maximum %d)", maxTokenizerVocabularySize) + } + return nil +} + +func validateTokenizerID(id int32) error { + if id < 0 { + return fmt.Errorf("tokenizer ID %d must not be negative", id) + } + if id >= maxTokenizerVocabularySize { + return fmt.Errorf("tokenizer ID %d exceeds maximum %d", id, maxTokenizerVocabularySize-1) + } + return nil +} + // TokenizerConfig holds optional configuration data that can be passed to LoadFromBytesWithConfig. type TokenizerConfig struct { TokenizerConfigJSON []byte // tokenizer_config.json content @@ -52,11 +77,7 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { } `json:"model"` PreTokenizer json.RawMessage `json:"pre_tokenizer"` Decoder json.RawMessage `json:"decoder"` - AddedTokens []struct { - ID int32 `json:"id"` - Content string `json:"content"` - Special bool `json:"special"` - } `json:"added_tokens"` + AddedTokens []addedToken `json:"added_tokens"` } if err := json.Unmarshal(data, &raw); err != nil { @@ -68,6 +89,86 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { return nil, fmt.Errorf("unsupported tokenizer type: %s", raw.Model.Type) } + if err := validateTokenizerRecordCount(len(raw.Model.Vocab), len(raw.AddedTokens)); err != nil { + return nil, fmt.Errorf("invalid tokenizer vocabulary: %w", err) + } + + baseByID := make(map[int32]string, len(raw.Model.Vocab)) + maxID := int32(-1) + // Select canonical failures instead of returning from randomized map traversal. + // Range errors take precedence, and the lowest numeric ID wins within each class. + invalidBaseID := int32(0) + hasInvalidBaseID := false + duplicateBaseID := int32(0) + hasDuplicateBaseID := false + for token, id := range raw.Model.Vocab { + if err := validateTokenizerID(id); err != nil { + if !hasInvalidBaseID || id < invalidBaseID { + invalidBaseID = id + hasInvalidBaseID = true + } + continue + } + if _, ok := baseByID[id]; ok { + if !hasDuplicateBaseID || id < duplicateBaseID { + duplicateBaseID = id + hasDuplicateBaseID = true + } + continue + } + baseByID[id] = token + if id > maxID { + maxID = id + } + } + if hasInvalidBaseID { + return nil, fmt.Errorf("invalid base token ID: %w", validateTokenizerID(invalidBaseID)) + } + if hasDuplicateBaseID { + return nil, fmt.Errorf("duplicate base token ID %d", duplicateBaseID) + } + + addedByID := make(map[int32]string, len(raw.AddedTokens)) + addedByContent := make(map[string]int32, len(raw.AddedTokens)) + for _, tok := range raw.AddedTokens { + if err := validateTokenizerID(tok.ID); err != nil { + return nil, fmt.Errorf("invalid added token %q: %w", tok.Content, err) + } + if _, ok := addedByID[tok.ID]; ok { + return nil, fmt.Errorf("duplicate added token ID %d", tok.ID) + } + if previousID, ok := addedByContent[tok.Content]; ok { + first, second := previousID, tok.ID + if first > second { + first, second = second, first + } + return nil, fmt.Errorf("duplicate added token content %q with IDs %d and %d", tok.Content, first, second) + } + addedByID[tok.ID] = tok.Content + addedByContent[tok.Content] = tok.ID + if tok.ID > maxID { + maxID = tok.ID + } + } + + for _, tok := range raw.AddedTokens { + if baseContent, ok := baseByID[tok.ID]; ok && baseContent != tok.Content { + return nil, fmt.Errorf("token ID %d has conflicting base and added content", tok.ID) + } + if baseID, ok := raw.Model.Vocab[tok.Content]; ok && baseID != tok.ID { + first, second := baseID, tok.ID + if first > second { + first, second = second, first + } + return nil, fmt.Errorf("token content %q has conflicting base and added IDs %d and %d", tok.Content, first, second) + } + } + + valuesLen := 0 + if maxID >= 0 { + valuesLen = int(maxID) + 1 + } + // Parse merges - can be []string (Llama) or [][]string (GPT-OSS). var mergesStrings []string if raw.Model.Merges != nil { @@ -91,7 +192,7 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // Build tokenizer t := &Tokenizer{ vocab: &Vocabulary{ - Values: make([]string, len(raw.Model.Vocab)), + Values: make([]string, valuesLen), Reverse: raw.Model.Vocab, Merges: make(map[string]int, len(mergesStrings)), BOS: -1, @@ -102,11 +203,6 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // Build values array for token, id := range raw.Model.Vocab { - if int(id) >= len(t.vocab.Values) { - newValues := make([]string, id+1) - copy(newValues, t.vocab.Values) - t.vocab.Values = newValues - } t.vocab.Values[id] = token } @@ -121,11 +217,6 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { // if it's a "truly special" token like BOS/EOS/PAD, but for tokenization we need // to treat all added_tokens as special to match HuggingFace behavior. for _, tok := range raw.AddedTokens { - if int(tok.ID) >= len(t.vocab.Values) { - newValues := make([]string, tok.ID+1) - copy(newValues, t.vocab.Values) - t.vocab.Values = newValues - } t.vocab.Values[tok.ID] = tok.Content t.specialTokens[tok.Content] = tok.ID // Add ALL added_tokens to special tokens } diff --git a/x/tokenizer/tokenizer_load_test.go b/x/tokenizer/tokenizer_load_test.go index c0d52e1a459..b40c2d48217 100644 --- a/x/tokenizer/tokenizer_load_test.go +++ b/x/tokenizer/tokenizer_load_test.go @@ -2,10 +2,258 @@ package tokenizer import ( "encoding/json" + "fmt" + "math" "strings" "testing" ) +func tokenizerJSON(vocab, addedTokens string) []byte { + return []byte(fmt.Sprintf(`{"model":{"type":"BPE","vocab":%s,"merges":[]},"added_tokens":%s}`, vocab, addedTokens)) +} + +func TestValidateTokenizerID(t *testing.T) { + tests := []struct { + name string + id int32 + want string + }{ + {name: "negative", id: -1, want: "must not be negative"}, + {name: "minimum int32", id: math.MinInt32, want: "must not be negative"}, + {name: "maximum int32", id: math.MaxInt32, want: "exceeds maximum"}, + {name: "exclusive cap", id: maxTokenizerVocabularySize, want: "exceeds maximum"}, + {name: "inclusive upper ID", id: maxTokenizerVocabularySize - 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateTokenizerID(tt.id) + if tt.want == "" { + if err != nil { + t.Fatalf("validateTokenizerID(%d): %v", tt.id, err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("validateTokenizerID(%d) error = %v, want containing %q", tt.id, err, tt.want) + } + }) + } +} + +func TestValidateTokenizerRecordCount(t *testing.T) { + if err := validateTokenizerRecordCount(maxTokenizerVocabularySize-1, 1); err != nil { + t.Fatalf("maximum record count rejected: %v", err) + } + if err := validateTokenizerRecordCount(maxTokenizerVocabularySize, 1); err == nil { + t.Fatal("expected aggregate record count above cap to fail") + } +} + +func TestLoadFromBytesRejectsInvalidTokenizerIDs(t *testing.T) { + tests := []struct { + name string + data []byte + wants []string + }{ + {name: "negative base", data: tokenizerJSON(`{"a":-1}`, `[]`), wants: []string{"invalid base token", "tokenizer ID -1"}}, + {name: "minimum base", data: tokenizerJSON(fmt.Sprintf(`{"a":%d}`, int64(math.MinInt32)), `[]`), wants: []string{"invalid base token", fmt.Sprint(math.MinInt32)}}, + {name: "negative added", data: tokenizerJSON(`{}`, `[{"id":-1,"content":"a"}]`), wants: []string{"invalid added token", "tokenizer ID -1"}}, + {name: "minimum added", data: tokenizerJSON(`{}`, fmt.Sprintf(`[{"id":%d,"content":"a"}]`, int64(math.MinInt32))), wants: []string{"invalid added token", fmt.Sprint(math.MinInt32)}}, + {name: "base at exclusive cap", data: tokenizerJSON(fmt.Sprintf(`{"a":%d}`, maxTokenizerVocabularySize), `[]`), wants: []string{"invalid base token", "exceeds maximum"}}, + {name: "added at exclusive cap", data: tokenizerJSON(`{}`, fmt.Sprintf(`[{"id":%d,"content":"a"}]`, maxTokenizerVocabularySize)), wants: []string{"invalid added token", "exceeds maximum"}}, + {name: "above int32", data: tokenizerJSON(`{"a":2147483648}`, `[]`), wants: []string{"failed to parse tokenizer", "cannot unmarshal number"}}, + {name: "below int32", data: tokenizerJSON(`{"a":-2147483649}`, `[]`), wants: []string{"failed to parse tokenizer", "cannot unmarshal number"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := LoadFromBytes(tt.data) + if err == nil { + t.Fatal("expected tokenizer load to fail") + } + for _, want := range tt.wants { + if !strings.Contains(err.Error(), want) { + t.Fatalf("error = %v, want containing %q", err, want) + } + } + }) + } +} + +func TestLoadFromBytesSupportsBoundedSparseIDs(t *testing.T) { + tests := []struct { + name string + data []byte + wantSize int + decodeID int32 + wantDecoded string + special string + wantSpecial int32 + }{ + { + name: "inclusive upper base ID", + data: tokenizerJSON(fmt.Sprintf(`{"edge":%d}`, maxTokenizerVocabularySize-1), `[]`), + wantSize: maxTokenizerVocabularySize, + decodeID: maxTokenizerVocabularySize - 1, + wantDecoded: "edge", + }, + {name: "sparse base", data: tokenizerJSON(`{"base":1000}`, `[]`), wantSize: 1001, decodeID: 1000, wantDecoded: "base"}, + {name: "sparse added", data: tokenizerJSON(`{}`, `[{"id":1000,"content":"added"}]`), wantSize: 1001, decodeID: 1000, wantDecoded: "added", special: "added", wantSpecial: 1000}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tok, err := LoadFromBytes(tt.data) + if err != nil { + t.Fatal(err) + } + if got := tok.VocabSize(); got != tt.wantSize { + t.Fatalf("VocabSize() = %d, want %d", got, tt.wantSize) + } + if got := tok.Decode([]int32{tt.decodeID}); got != tt.wantDecoded { + t.Fatalf("Decode(%d) = %q, want %q", tt.decodeID, got, tt.wantDecoded) + } + if tt.special != "" { + id, ok := tok.GetSpecialToken(tt.special) + if !ok || id != tt.wantSpecial { + t.Fatalf("GetSpecialToken(%q) = (%d, %v), want (%d, true)", tt.special, id, ok, tt.wantSpecial) + } + } + }) + } +} + +func TestLoadFromBytesRejectsTokenizerCollisions(t *testing.T) { + tests := []struct { + name string + inputs [][]byte + wantErr string + }{ + { + name: "duplicate base ID", + inputs: [][]byte{ + tokenizerJSON(`{"a":1,"b":1}`, `[]`), + tokenizerJSON(`{"b":1,"a":1}`, `[]`), + }, + wantErr: "duplicate base token ID 1", + }, + { + name: "duplicate added ID", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":1,"content":"b"}]`), + tokenizerJSON(`{}`, `[{"id":1,"content":"b"},{"id":1,"content":"a"}]`), + }, + wantErr: "duplicate added token ID 1", + }, + { + name: "repeated added entry", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":1,"content":"a"}]`), + }, + wantErr: "duplicate added token ID 1", + }, + { + name: "duplicate added content", + inputs: [][]byte{ + tokenizerJSON(`{}`, `[{"id":1,"content":"a"},{"id":2,"content":"a"}]`), + tokenizerJSON(`{}`, `[{"id":2,"content":"a"},{"id":1,"content":"a"}]`), + }, + wantErr: `duplicate added token content "a" with IDs 1 and 2`, + }, + { + name: "cross-source ID conflict", + inputs: [][]byte{ + tokenizerJSON(`{"a":1}`, `[{"id":1,"content":"b"}]`), + }, + wantErr: "token ID 1 has conflicting base and added content", + }, + { + name: "cross-source content conflict", + inputs: [][]byte{ + tokenizerJSON(`{"a":1}`, `[{"id":2,"content":"a"}]`), + }, + wantErr: `token content "a" has conflicting base and added IDs 1 and 2`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for i, data := range tt.inputs { + _, err := LoadFromBytes(data) + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("input %d error = %v, want %q", i, err, tt.wantErr) + } + } + }) + } +} + +func TestLoadFromBytesBaseValidationIsDeterministic(t *testing.T) { + tests := []struct { + name string + inputs [][]byte + wantErr string + }{ + { + name: "range error precedes duplicate ID", + inputs: [][]byte{ + tokenizerJSON(fmt.Sprintf(`{"duplicate-a":5,"negative":-2,"over-cap":%d,"duplicate-b":5}`, maxTokenizerVocabularySize), `[]`), + tokenizerJSON(fmt.Sprintf(`{"duplicate-b":5,"over-cap":%d,"negative":-2,"duplicate-a":5}`, maxTokenizerVocabularySize), `[]`), + }, + wantErr: "invalid base token ID: tokenizer ID -2 must not be negative", + }, + { + name: "lowest duplicate ID wins", + inputs: [][]byte{ + tokenizerJSON(`{"high-a":7,"low-a":2,"high-b":7,"low-b":2}`, `[]`), + tokenizerJSON(`{"low-b":2,"high-b":7,"low-a":2,"high-a":7}`, `[]`), + }, + wantErr: "duplicate base token ID 2", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for i, data := range tt.inputs { + for attempt := range 100 { + _, err := LoadFromBytes(data) + if err == nil || err.Error() != tt.wantErr { + t.Fatalf("input %d attempt %d error = %v, want %q", i, attempt, err, tt.wantErr) + } + } + } + }) + } +} + +func TestLoadFromBytesAcceptsExactAddedTokenPromotion(t *testing.T) { + tok, err := LoadFromBytes(tokenizerJSON(`{"a":1}`, `[{"id":1,"content":"a"}]`)) + if err != nil { + t.Fatal(err) + } + if got := tok.Decode([]int32{1}); got != "a" { + t.Fatalf("Decode(1) = %q, want a", got) + } + id, ok := tok.GetSpecialToken("a") + if !ok || id != 1 { + t.Fatalf("GetSpecialToken(a) = (%d, %v), want (1, true)", id, ok) + } +} + +func TestLoadFromBytesAcceptsEmptyAndContiguousVocabulary(t *testing.T) { + for name, data := range map[string][]byte{ + "empty": tokenizerJSON(`{}`, `[]`), + "contiguous": tokenizerJSON(`{"a":0,"b":1}`, `[]`), + } { + t.Run(name, func(t *testing.T) { + if _, err := LoadFromBytes(data); err != nil { + t.Fatal(err) + } + }) + } +} + func TestLoadFromBytesRejectsWordPiece(t *testing.T) { data := []byte(`{ "model": { From d1f3b614fc3d5689d02933dfa3a7e02875067850 Mon Sep 17 00:00:00 2001 From: Philipp Date: Wed, 26 Aug 2026 18:58:40 +0000 Subject: [PATCH 10/58] tokenizer: bound added token IDs during load Co-authored-by: Codex --- x/tokenizer/tokenizer_load.go | 27 +++++++- x/tokenizer/tokenizer_load_test.go | 105 +++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/x/tokenizer/tokenizer_load.go b/x/tokenizer/tokenizer_load.go index 0852abd380c..38a6500f9c5 100644 --- a/x/tokenizer/tokenizer_load.go +++ b/x/tokenizer/tokenizer_load.go @@ -39,6 +39,10 @@ type TokenizerConfig struct { GenerationConfigJSON []byte // generation_config.json content SpecialTokensMapJSON []byte // special_tokens_map.json content ConfigJSON []byte // config.json content + // AddedTokenIDLimit is an optional exclusive upper bound for added token IDs. + // When positive, added tokens with negative IDs or IDs at or above the limit + // are ignored. A zero value disables filtering. + AddedTokenIDLimit int32 } // LoadFromBytes loads a tokenizer from tokenizer.json bytes. @@ -46,13 +50,17 @@ type TokenizerConfig struct { // Note: This won't load special token config from companion files. Use LoadFromBytesWithConfig // to provide tokenizer_config.json data for proper PAD/EOS token loading. func LoadFromBytes(data []byte) (*Tokenizer, error) { - return loadFromTokenizerJSON(data) + return loadFromTokenizerJSON(data, 0) } // LoadFromBytesWithConfig loads a tokenizer from tokenizer.json bytes with additional config files. // This is useful when loading from blob storage where companion config files are also blobs. func LoadFromBytesWithConfig(data []byte, config *TokenizerConfig) (*Tokenizer, error) { - t, err := loadFromTokenizerJSON(data) + var addedTokenIDLimit int32 + if config != nil { + addedTokenIDLimit = config.AddedTokenIDLimit + } + t, err := loadFromTokenizerJSON(data, addedTokenIDLimit) if err != nil { return nil, err } @@ -68,7 +76,11 @@ func LoadFromBytesWithConfig(data []byte, config *TokenizerConfig) (*Tokenizer, } // loadFromTokenizerJSON parses tokenizer.json content from bytes. -func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { +func loadFromTokenizerJSON(data []byte, addedTokenIDLimit int32) (*Tokenizer, error) { + if addedTokenIDLimit < 0 { + return nil, fmt.Errorf("added token ID limit must not be negative: %d", addedTokenIDLimit) + } + var raw struct { Model struct { Type string `json:"type"` // "BPE" @@ -88,6 +100,15 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { if raw.Model.Type != "BPE" { return nil, fmt.Errorf("unsupported tokenizer type: %s", raw.Model.Type) } + if addedTokenIDLimit > 0 { + filtered := raw.AddedTokens[:0] + for _, tok := range raw.AddedTokens { + if tok.ID >= 0 && tok.ID < addedTokenIDLimit { + filtered = append(filtered, tok) + } + } + raw.AddedTokens = filtered + } if err := validateTokenizerRecordCount(len(raw.Model.Vocab), len(raw.AddedTokens)); err != nil { return nil, fmt.Errorf("invalid tokenizer vocabulary: %w", err) diff --git a/x/tokenizer/tokenizer_load_test.go b/x/tokenizer/tokenizer_load_test.go index b40c2d48217..547564a1a7b 100644 --- a/x/tokenizer/tokenizer_load_test.go +++ b/x/tokenizer/tokenizer_load_test.go @@ -81,6 +81,111 @@ func TestLoadFromBytesRejectsInvalidTokenizerIDs(t *testing.T) { } } +func TestLoadFromBytesWithConfigFiltersAddedTokenIDs(t *testing.T) { + data := tokenizerJSON(`{"base":0}`, `[ + {"id":-1,"content":"negative"}, + {"id":1,"content":"kept"}, + {"id":2,"content":"at-limit"}, + {"id":3,"content":"above-limit"} + ]`) + tok, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + + if got := tok.VocabSize(); got != 2 { + t.Fatalf("VocabSize() = %d, want 2", got) + } + if got := tok.Decode([]int32{0, 1}); got != "basekept" { + t.Fatalf("Decode([0 1]) = %q, want %q", got, "basekept") + } + if id, ok := tok.GetSpecialToken("kept"); !ok || id != 1 { + t.Fatalf("GetSpecialToken(kept) = (%d, %v), want (1, true)", id, ok) + } + for _, content := range []string{"negative", "at-limit", "above-limit"} { + if id, ok := tok.GetSpecialToken(content); ok { + t.Fatalf("GetSpecialToken(%q) = (%d, true), want absent", content, id) + } + } +} + +func TestLoadFromBytesWithConfigAddedTokenIDLimitDisabled(t *testing.T) { + data := tokenizerJSON(`{}`, `[{"id":-1,"content":"negative"}]`) + for name, config := range map[string]*TokenizerConfig{ + "nil config": nil, + "zero limit": {}, + } { + t.Run(name, func(t *testing.T) { + _, err := LoadFromBytesWithConfig(data, config) + if err == nil || !strings.Contains(err.Error(), "invalid added token") { + t.Fatalf("error = %v, want invalid added token", err) + } + }) + } +} + +func TestLoadFromBytesWithConfigRejectsNegativeAddedTokenIDLimit(t *testing.T) { + _, err := LoadFromBytesWithConfig(tokenizerJSON(`{}`, `[]`), &TokenizerConfig{AddedTokenIDLimit: -1}) + if err == nil || err.Error() != "added token ID limit must not be negative: -1" { + t.Fatalf("error = %v, want negative limit error", err) + } +} + +func TestLoadFromBytesWithConfigValidatesOnlyRetainedAddedTokens(t *testing.T) { + tests := []struct { + name string + data []byte + }{ + { + name: "duplicate IDs are filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"first"}, + {"id":2,"content":"second"}, + {"id":1,"content":"kept"} + ]`), + }, + { + name: "duplicate content is filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"duplicate"}, + {"id":3,"content":"duplicate"}, + {"id":1,"content":"kept"} + ]`), + }, + { + name: "base collision is filtered", + data: tokenizerJSON(`{"base":0}`, `[ + {"id":2,"content":"base"}, + {"id":1,"content":"kept"} + ]`), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tok, err := LoadFromBytesWithConfig(tt.data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + if got := tok.VocabSize(); got != 2 { + t.Fatalf("VocabSize() = %d, want 2", got) + } + }) + } +} + +func TestLoadFromBytesWithConfigRejectsRetainedAddedTokenCollision(t *testing.T) { + data := tokenizerJSON(`{"base":0}`, `[ + {"id":1,"content":"first"}, + {"id":1,"content":"second"}, + {"id":2,"content":"filtered"} + ]`) + _, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err == nil || err.Error() != "duplicate added token ID 1" { + t.Fatalf("error = %v, want retained-token collision", err) + } +} + func TestLoadFromBytesSupportsBoundedSparseIDs(t *testing.T) { tests := []struct { name string From 5b3d11418804ed8dcd900f9fa9365b30a41dfa98 Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 06:39:49 +0000 Subject: [PATCH 11/58] test(tokenizer): clarify added token limit scope Co-authored-by: Codex --- x/tokenizer/tokenizer_load_test.go | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/x/tokenizer/tokenizer_load_test.go b/x/tokenizer/tokenizer_load_test.go index 547564a1a7b..f1dbb3a216e 100644 --- a/x/tokenizer/tokenizer_load_test.go +++ b/x/tokenizer/tokenizer_load_test.go @@ -109,6 +109,27 @@ func TestLoadFromBytesWithConfigFiltersAddedTokenIDs(t *testing.T) { } } +func TestLoadFromBytesWithConfigAddedTokenIDLimitDoesNotFilterBaseVocabulary(t *testing.T) { + data := tokenizerJSON(`{"base":3}`, `[ + {"id":1,"content":"added"}, + {"id":3,"content":"filtered-added"} + ]`) + tok, err := LoadFromBytesWithConfig(data, &TokenizerConfig{AddedTokenIDLimit: 2}) + if err != nil { + t.Fatal(err) + } + + if got := tok.VocabSize(); got != 4 { + t.Fatalf("VocabSize() = %d, want 4", got) + } + if got := tok.Decode([]int32{3, 1}); got != "baseadded" { + t.Fatalf("Decode([3 1]) = %q, want %q", got, "baseadded") + } + if _, ok := tok.GetSpecialToken("filtered-added"); ok { + t.Fatal("added token at the limit was retained") + } +} + func TestLoadFromBytesWithConfigAddedTokenIDLimitDisabled(t *testing.T) { data := tokenizerJSON(`{}`, `[{"id":-1,"content":"negative"}]`) for name, config := range map[string]*TokenizerConfig{ From 25dfd2265796599f08aea6e3a24aa7143400b62b Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:19:13 +0000 Subject: [PATCH 12/58] create: plan GPT-OSS tensor imports Co-authored-by: Codex --- x/create/create.go | 7 +- x/create/gptoss.go | 746 +++++++++++++++++++++++++++++++++++++++++++++ x/create/plan.go | 17 ++ 3 files changed, 769 insertions(+), 1 deletion(-) create mode 100644 x/create/gptoss.go diff --git a/x/create/create.go b/x/create/create.go index 84d17dd3466..df23786e05b 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -430,7 +430,11 @@ func (cfg sourceModelConfig) QuantMetadata() map[string]string { } } - quantType := sourceQuantType(q.Mode, q.Bits) + mode := q.Mode + if mode == "" { + mode = q.QuantMethod + } + quantType := sourceQuantType(mode, q.Bits) if quantType == "" { return nil } @@ -475,6 +479,7 @@ func (cfg sourceModelConfig) HFFP8WeightBlockSize() (rows, cols int32, ok bool) type tensorImportTransformFactory func(rawConfig json.RawMessage) (quantizePolicy, error) var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ + "GptOssForCausalLM": newGPTOSSImportTransform, "Qwen3_5ForCausalLM": newQwen35ImportTransform, "Qwen3_5ForConditionalGeneration": newQwen35ImportTransform, "Qwen3NextForCausalLM": newQwen35ImportTransform, diff --git a/x/create/gptoss.go b/x/create/gptoss.go new file mode 100644 index 00000000000..71f492f01eb --- /dev/null +++ b/x/create/gptoss.go @@ -0,0 +1,746 @@ +package create + +import ( + "encoding/json" + "fmt" + "io" + "math" + "os" + "strconv" + "strings" + + "github.com/ollama/ollama/x/safetensors" +) + +type gptossPerTensorQuant struct { + mode string + bits int + groupSize int +} + +type gptossImportTransform struct { + defaultQuant gptossPerTensorQuant + perTensorQuant map[string]gptossPerTensorQuant +} + +func validateGPTOSSDequant() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("GPTOSS_VALIDATE_DEQUANT"))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func newGPTOSSImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) { + t := &gptossImportTransform{ + perTensorQuant: make(map[string]gptossPerTensorQuant), + } + + defaultQuant, perTensor := parseGPTOSSPerTensorQuant(rawConfig) + t.defaultQuant = defaultQuant + t.perTensorQuant = perTensor + + return t, nil +} + +func parseGPTOSSPerTensorQuant(rawConfig json.RawMessage) (gptossPerTensorQuant, map[string]gptossPerTensorQuant) { + defaultQ := gptossPerTensorQuant{} + perTensor := make(map[string]gptossPerTensorQuant) + + var raw struct { + Quantization json.RawMessage `json:"quantization"` + } + if err := json.Unmarshal(rawConfig, &raw); err != nil || raw.Quantization == nil { + return defaultQ, perTensor + } + + var entries map[string]json.RawMessage + if err := json.Unmarshal(raw.Quantization, &entries); err != nil { + return defaultQ, perTensor + } + + type quantEntry struct { + Bits int `json:"bits"` + GroupSize int `json:"group_size"` + Mode string `json:"mode"` + QuantMethod string `json:"quant_method"` + } + + if v, ok := entries["bits"]; ok { + json.Unmarshal(v, &defaultQ.bits) + } + if v, ok := entries["group_size"]; ok { + json.Unmarshal(v, &defaultQ.groupSize) + } + if v, ok := entries["mode"]; ok { + json.Unmarshal(v, &defaultQ.mode) + } + if defaultQ.mode == "" { + if v, ok := entries["quant_method"]; ok { + json.Unmarshal(v, &defaultQ.mode) + } + } + + for key, val := range entries { + if key == "bits" || key == "group_size" || key == "mode" || key == "quant_method" { + continue + } + var entry quantEntry + if err := json.Unmarshal(val, &entry); err != nil { + continue + } + if entry.Bits > 0 { + mode := entry.Mode + if mode == "" { + mode = entry.QuantMethod + } + // Infer mode when not specified: if the tensor has biases + // (affine pattern) or uses group_size=64 (affine default), + // it's affine quantization. The MLX checkpoint omits mode + // for some tensors like the router. + if mode == "" && entry.GroupSize == 64 { + mode = "affine" + } + q := gptossPerTensorQuant{ + mode: mode, + bits: entry.Bits, + groupSize: entry.GroupSize, + } + perTensor[key] = q + } + } + + return defaultQ, perTensor +} + +func isGptossRouterWeight(name string) bool { + return strings.HasSuffix(name, ".router.weight") +} + +func (t *gptossImportTransform) quantizationType(name string, shape []int32, quantize string) string { + // MoE router weights choose the top-k expert set. Quantization noise can + // flip expert selection, causing downstream activations to diverge sharply. + // The tensor is small, so leave it in source precision. + if isGptossRouterWeight(name) { + return "" + } + + quantNorm := normalizeQuantType(quantize) + if quantNorm == "" { + return "" + } + if strings.Contains(name, ".experts.") && strings.HasSuffix(name, ".weight") { + return "" + } + return GetTensorQuantization(name, shape, quantize) +} + +func (t *gptossImportTransform) prequantizedMetadata(sourceName string, global map[string]string) map[string]string { + prefix := strings.TrimSuffix(sourceName, ".weight") + if prefix == sourceName { + return global + } + + q, ok := t.perTensorQuant[prefix] + if !ok { + return global + } + + qt := sourceQuantType(q.mode, q.bits) + if qt == "" { + return global + } + + override := make(map[string]string, len(global)+2) + for k, v := range global { + override[k] = v + } + override["quant_type"] = qt + if q.groupSize > 0 { + override["group_size"] = strconv.Itoa(q.groupSize) + } + return override +} + +func planGPTOSS(inv Inventory, class Classification, policy quantizePolicy) ([]BlobSpec, error) { + t, ok := policy.(*gptossImportTransform) + if !ok { + return nil, fmt.Errorf("gpt-oss planner requires gptossImportTransform, got %T", policy) + } + + consumed := make(map[string]bool) + expertGroups := make(map[string][]TensorSpec) + expertMetadata := make(map[string]map[string]string) + var specs []BlobSpec + + for _, name := range sortedTensorNames(inv) { + if consumed[name] { + continue + } + if t.isNativeCompanion(inv, name) { + continue + } + if strings.HasSuffix(name, "_blocks") { + group, tensors, metadata, sources, ok, err := t.planNativeExpertTensor(inv, name) + if err != nil { + return nil, err + } + if ok { + expertGroups[group] = append(expertGroups[group], tensors...) + if len(metadata) > 0 { + if expertMetadata[group] == nil { + expertMetadata[group] = make(map[string]string) + } + for k, v := range metadata { + expertMetadata[group][k] = v + } + } + for _, source := range sources { + consumed[source] = true + } + continue + } + } + + if spec, sources, ok := t.planNativeQuantizedDense(inv, name); ok { + specs = append(specs, spec) + for _, source := range sources { + consumed[source] = true + } + continue + } + + outName := t.canonicalTensorName(name) + source := inv.Tensors[name] + q := "" + if class.Kind == SourceFloat && class.Quantize != "" { + q = policy.quantizationType(outName, source.Shape, class.Quantize) + } + specs = append(specs, BlobSpec{ + Name: outName, + Tensors: []TensorSpec{{Name: outName, Sources: []SourceTensor{source}, Quantize: q}}, + }) + } + + for _, group := range sortedKeys(expertGroups) { + specs = append(specs, BlobSpec{ + Name: group, + Tensors: expertGroups[group], + Metadata: expertMetadata[group], + }) + } + + return specs, nil +} + +func (t *gptossImportTransform) isNativeCompanion(inv Inventory, name string) bool { + switch { + case strings.HasSuffix(name, "_scales"): + return inv.Has(strings.TrimSuffix(name, "_scales") + "_blocks") + case strings.HasSuffix(name, "_bias"): + return inv.Has(strings.TrimSuffix(name, "_bias") + "_blocks") + case strings.HasSuffix(name, ".scales"): + return inv.Has(strings.TrimSuffix(name, ".scales") + ".weight") + case strings.HasSuffix(name, ".biases"): + return inv.Has(strings.TrimSuffix(name, ".biases") + ".weight") + default: + return false + } +} + +func (t *gptossImportTransform) planNativeQuantizedDense(inv Inventory, weightName string) (BlobSpec, []string, bool) { + if !strings.HasSuffix(weightName, ".weight") { + return BlobSpec{}, nil, false + } + scaleName := strings.TrimSuffix(weightName, ".weight") + ".scales" + if !inv.Has(scaleName) { + return BlobSpec{}, nil, false + } + + outWeight := t.canonicalTensorName(weightName) + if outWeight == weightName { + return BlobSpec{}, nil, false + } + + weight := inv.Tensors[weightName] + scale := inv.Tensors[scaleName] + tensors := []TensorSpec{ + {Name: outWeight, Sources: []SourceTensor{weight}}, + {Name: outWeight + ".scale", Sources: []SourceTensor{scale}}, + } + sources := []string{weightName, scaleName} + if biasName := strings.TrimSuffix(weightName, ".weight") + ".biases"; inv.Has(biasName) { + tensors = append(tensors, TensorSpec{Name: outWeight + ".bias", Sources: []SourceTensor{inv.Tensors[biasName]}}) + sources = append(sources, biasName) + } + + return BlobSpec{ + Name: outWeight, + Tensors: tensors, + Metadata: t.prequantizedMetadata(weightName, t.defaultMetadata()), + }, sources, true +} + +func (t *gptossImportTransform) planNativeExpertTensor(inv Inventory, blocksName string) (string, []TensorSpec, map[string]string, []string, bool, error) { + scalesName := strings.TrimSuffix(blocksName, "_blocks") + "_scales" + if !inv.Has(scalesName) { + return "", nil, nil, nil, false, nil + } + + outName := t.canonicalTensorName(blocksName) + group, ok := strings.CutSuffix(outName, ".gate_up_proj.weight") + if ok { + group = group + "" + } else if group, ok = strings.CutSuffix(outName, ".down_proj.weight"); ok { + group = group + "" + } + if !ok { + return "", nil, nil, nil, false, nil + } + group = strings.TrimSuffix(group, ".gate_up_proj") + group = strings.TrimSuffix(group, ".down_proj") + + blocks := inv.Tensors[blocksName] + scales := inv.Tensors[scalesName] + sources := []string{blocksName, scalesName} + metadata := t.prequantizedMetadata(strings.TrimSuffix(blocksName, "_blocks")+".weight", t.defaultMetadata()) + if strings.Contains(outName, ".gate_up_proj.weight") { + gateWeight := strings.Replace(outName, "gate_up_proj", "gate_proj", 1) + upWeight := strings.Replace(outName, "gate_up_proj", "up_proj", 1) + tensors := []TensorSpec{ + {Name: gateWeight, Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSGateUpWeight}, + {Name: gateWeight + ".scale", Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSGateUpScale}, + {Name: upWeight, Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSUpWeight}, + {Name: upWeight + ".scale", Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSUpScale}, + } + if biasName := strings.TrimSuffix(blocksName, "_blocks") + "_bias"; inv.Has(biasName) { + bias := inv.Tensors[biasName] + gateBias := strings.Replace(outName, "gate_up_proj.weight", "gate_proj.bias", 1) + upBias := strings.Replace(outName, "gate_up_proj.weight", "up_proj.bias", 1) + tensors = append(tensors, + TensorSpec{Name: gateBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSGateUpBias}, + TensorSpec{Name: upBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSUpBias}, + ) + sources = append(sources, biasName) + } + return group, tensors, metadata, sources, true, nil + } + + tensors := []TensorSpec{ + {Name: outName, Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSPackedExpertWeight}, + {Name: outName + ".scale", Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSPackedExpertScale}, + } + if biasName := strings.TrimSuffix(blocksName, "_blocks") + "_bias"; inv.Has(biasName) { + tensors = append(tensors, TensorSpec{Name: strings.Replace(outName, ".weight", ".bias", 1), Sources: []SourceTensor{inv.Tensors[biasName]}}) + sources = append(sources, biasName) + } + return group, tensors, metadata, sources, true, nil +} + +func (t *gptossImportTransform) defaultMetadata() map[string]string { + qt := sourceQuantType(t.defaultQuant.mode, t.defaultQuant.bits) + if qt == "" { + qt = "mxfp4" + } + metadata := map[string]string{"quant_type": qt} + if t.defaultQuant.groupSize > 0 { + metadata["group_size"] = strconv.Itoa(t.defaultQuant.groupSize) + } else if qt == "mxfp4" { + metadata["group_size"] = "32" + } + return metadata +} + +func validateGPTOSSPackedMXFP4Inputs(name string, blocks, scales *safetensors.TensorData) error { + if blocks == nil || scales == nil { + return fmt.Errorf("gpt-oss expert tensor %q requires blocks and scales", name) + } + if blocks.Dtype != "U8" { + return fmt.Errorf("gpt-oss expert blocks %q dtype = %q, want U8", blocks.Name, blocks.Dtype) + } + if scales.Dtype != "U8" { + return fmt.Errorf("gpt-oss expert scales %q dtype = %q, want U8", scales.Name, scales.Dtype) + } + if len(blocks.Shape) != 4 { + return fmt.Errorf("gpt-oss expert blocks %q shape = %v, want [experts out groups 16]", blocks.Name, blocks.Shape) + } + if len(scales.Shape) != 3 { + return fmt.Errorf("gpt-oss expert scales %q shape = %v, want [experts out groups]", scales.Name, scales.Shape) + } + if blocks.Shape[0] != scales.Shape[0] || blocks.Shape[1] != scales.Shape[1] || blocks.Shape[2] != scales.Shape[2] { + return fmt.Errorf("gpt-oss expert tensor %q shape mismatch: blocks=%v scales=%v", name, blocks.Shape, scales.Shape) + } + if blocks.Shape[3] != 16 { + return fmt.Errorf("gpt-oss expert blocks %q trailing shape = %v, want [... 16]", blocks.Name, blocks.Shape) + } + return nil +} + +func preservePackedExpertProjection(name string, blocks, scales *safetensors.TensorData) ([]*safetensors.TensorData, error) { + if err := validateGPTOSSPackedMXFP4Inputs(name, blocks, scales); err != nil { + return nil, err + } + + sourceBlockBytes, err := io.ReadAll(blocks.Reader()) + if err != nil { + return nil, fmt.Errorf("read gpt-oss expert blocks %q: %w", blocks.Name, err) + } + sourceScaleBytes, err := io.ReadAll(scales.Reader()) + if err != nil { + return nil, fmt.Errorf("read gpt-oss expert scales %q: %w", scales.Name, err) + } + blockBytes, err := preserveGPTOSSMXFP4Blocks(blocks.Name, sourceBlockBytes, int(blocks.Shape[2])) + if err != nil { + return nil, err + } + scaleBytes := convertGPTOSSMXFP4Scales(sourceScaleBytes) + + blockShape := []int32{blocks.Shape[0], blocks.Shape[1], blocks.Shape[2] * 4} + scaleShape := []int32{scales.Shape[0], scales.Shape[1], scales.Shape[2]} + return []*safetensors.TensorData{ + safetensors.NewTensorDataFromBytes(name, "U32", blockShape, blockBytes), + safetensors.NewTensorDataFromBytes(name+".scale", "U8", scaleShape, scaleBytes), + }, nil +} + +func preserveAndSplitGateUpTensor(name string, blocks, scales *safetensors.TensorData) ([]*safetensors.TensorData, error) { + if err := validateGPTOSSPackedMXFP4Inputs(name, blocks, scales); err != nil { + return nil, err + } + + experts, outDim, groups := int(blocks.Shape[0]), int(blocks.Shape[1]), int(blocks.Shape[2]) + if outDim%2 != 0 { + return nil, fmt.Errorf("gpt-oss expert tensor %q output dim = %d, want even gate/up rows", name, outDim) + } + mid := outDim / 2 + + sourceBlockBytes, err := io.ReadAll(blocks.Reader()) + if err != nil { + return nil, fmt.Errorf("read gpt-oss expert blocks %q: %w", blocks.Name, err) + } + sourceScaleBytes, err := io.ReadAll(scales.Reader()) + if err != nil { + return nil, fmt.Errorf("read gpt-oss expert scales %q: %w", scales.Name, err) + } + blockBytes, err := preserveGPTOSSMXFP4Blocks(blocks.Name, sourceBlockBytes, groups) + if err != nil { + return nil, err + } + scaleBytes := convertGPTOSSMXFP4Scales(sourceScaleBytes) + + rowBlockBytes := groups * 16 + rowScaleBytes := groups + wantBlockBytes := experts * outDim * rowBlockBytes + wantScaleBytes := experts * outDim * rowScaleBytes + if len(blockBytes) != wantBlockBytes { + return nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want %d", blocks.Name, len(blockBytes), wantBlockBytes) + } + if len(scaleBytes) != wantScaleBytes { + return nil, fmt.Errorf("gpt-oss expert scales %q byte length = %d, want %d", scales.Name, len(scaleBytes), wantScaleBytes) + } + + gateBlocks := make([]byte, experts*mid*rowBlockBytes) + upBlocks := make([]byte, experts*mid*rowBlockBytes) + gateScales := make([]byte, experts*mid*rowScaleBytes) + upScales := make([]byte, experts*mid*rowScaleBytes) + for e := range experts { + for row := range outDim { + dstRow := row / 2 + srcBlock := (e*outDim + row) * rowBlockBytes + dstBlock := (e*mid + dstRow) * rowBlockBytes + srcScale := (e*outDim + row) * rowScaleBytes + dstScale := (e*mid + dstRow) * rowScaleBytes + if row%2 == 0 { + copy(gateBlocks[dstBlock:dstBlock+rowBlockBytes], blockBytes[srcBlock:srcBlock+rowBlockBytes]) + copy(gateScales[dstScale:dstScale+rowScaleBytes], scaleBytes[srcScale:srcScale+rowScaleBytes]) + } else { + copy(upBlocks[dstBlock:dstBlock+rowBlockBytes], blockBytes[srcBlock:srcBlock+rowBlockBytes]) + copy(upScales[dstScale:dstScale+rowScaleBytes], scaleBytes[srcScale:srcScale+rowScaleBytes]) + } + } + } + + gateName := strings.Replace(name, "gate_up_proj", "gate_proj", 1) + upName := strings.Replace(name, "gate_up_proj", "up_proj", 1) + blockShape := []int32{int32(experts), int32(mid), int32(groups * 4)} + scaleShape := []int32{int32(experts), int32(mid), int32(groups)} + return []*safetensors.TensorData{ + safetensors.NewTensorDataFromBytes(gateName, "U32", blockShape, gateBlocks), + safetensors.NewTensorDataFromBytes(gateName+".scale", "U8", scaleShape, gateScales), + safetensors.NewTensorDataFromBytes(upName, "U32", blockShape, upBlocks), + safetensors.NewTensorDataFromBytes(upName+".scale", "U8", scaleShape, upScales), + }, nil +} + +func preserveGPTOSSMXFP4Blocks(name string, source []byte, groups int) ([]byte, error) { + if groups <= 0 { + return nil, fmt.Errorf("gpt-oss expert blocks %q group count must be positive, got %d", name, groups) + } + if len(source)%(groups*16) != 0 { + return nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want multiple of row bytes %d", name, len(source), groups*16) + } + + out := make([]byte, len(source)) + copy(out, source) + return out, nil +} + +func convertGPTOSSMXFP4Scales(source []byte) []byte { + out := make([]byte, len(source)) + copy(out, source) + return out +} + +var gptossMXFP4Values = [16]float32{0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6} + +func decodeGPTOSSMXFP4Scale(scale byte) float32 { + return math.Float32frombits(uint32(scale) << 23) +} + +func decodeGPTOSSMXFP4TensorValues(name string, blocks, scales *safetensors.TensorData) ([]float32, []int32, error) { + if blocks == nil || scales == nil { + return nil, nil, fmt.Errorf("gpt-oss expert tensor %q requires blocks and scales", name) + } + if blocks.Dtype != "U8" { + return nil, nil, fmt.Errorf("gpt-oss expert blocks %q dtype = %q, want U8", blocks.Name, blocks.Dtype) + } + if scales.Dtype != "U8" { + return nil, nil, fmt.Errorf("gpt-oss expert scales %q dtype = %q, want U8", scales.Name, scales.Dtype) + } + if len(blocks.Shape) != 4 { + return nil, nil, fmt.Errorf("gpt-oss expert blocks %q shape = %v, want [experts out groups 16]", blocks.Name, blocks.Shape) + } + if len(scales.Shape) != 3 { + return nil, nil, fmt.Errorf("gpt-oss expert scales %q shape = %v, want [experts out groups]", scales.Name, scales.Shape) + } + if blocks.Shape[0] != scales.Shape[0] || blocks.Shape[1] != scales.Shape[1] || blocks.Shape[2] != scales.Shape[2] { + return nil, nil, fmt.Errorf("gpt-oss expert tensor %q shape mismatch: blocks=%v scales=%v", name, blocks.Shape, scales.Shape) + } + if blocks.Shape[3] != 16 { + return nil, nil, fmt.Errorf("gpt-oss expert blocks %q trailing shape = %v, want [... 16]", blocks.Name, blocks.Shape) + } + + blockBytes, err := io.ReadAll(blocks.Reader()) + if err != nil { + return nil, nil, fmt.Errorf("read gpt-oss expert blocks %q: %w", blocks.Name, err) + } + scaleBytes, err := io.ReadAll(scales.Reader()) + if err != nil { + return nil, nil, fmt.Errorf("read gpt-oss expert scales %q: %w", scales.Name, err) + } + + groupCount := int(blocks.Shape[0] * blocks.Shape[1] * blocks.Shape[2]) + if len(blockBytes) != groupCount*16 { + return nil, nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want %d", blocks.Name, len(blockBytes), groupCount*16) + } + if len(scaleBytes) != groupCount { + return nil, nil, fmt.Errorf("gpt-oss expert scales %q byte length = %d, want %d", scales.Name, len(scaleBytes), groupCount) + } + + values := make([]float32, groupCount*32) + for i := range groupCount { + src := blockBytes[i*16 : (i+1)*16] + + scale := decodeGPTOSSMXFP4Scale(scaleBytes[i]) + base := i * 32 + for j, packed := range src { + values[base+2*j] = gptossMXFP4Values[packed&0x0F] * scale + values[base+2*j+1] = gptossMXFP4Values[packed>>4] * scale + } + } + + if validateGPTOSSDequant() { + for i, v := range values { + if math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) { + return nil, nil, fmt.Errorf("gpt-oss expert tensor %q dequantized invalid value at %d", name, i) + } + } + } + + shape := []int32{blocks.Shape[0], blocks.Shape[1], blocks.Shape[2] * 32} + return values, shape, nil +} + +func dequantizeGPTOSSMXFP4Tensor(name string, blocks, scales *safetensors.TensorData) (*safetensors.TensorData, error) { + values, shape, err := decodeGPTOSSMXFP4TensorValues(name, blocks, scales) + if err != nil { + return nil, err + } + + raw, err := EncodeFloatTensor("BF16", values) + if err != nil { + return nil, fmt.Errorf("encode gpt-oss expert tensor %q as BF16: %w", name, err) + } + + return safetensors.NewTensorDataFromBytes(name, "BF16", shape, raw), nil +} + +func splitGateUpBiasTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { + if td == nil { + return nil, nil + } + if td.Dtype != "BF16" { + return nil, fmt.Errorf("gpt-oss expert tensor %q dtype = %q, want BF16", td.Name, td.Dtype) + } + if len(td.Shape) != 2 { + return nil, fmt.Errorf("gpt-oss expert tensor %q shape = %v, want [experts out]", td.Name, td.Shape) + } + experts, outDim := int(td.Shape[0]), int(td.Shape[1]) + if outDim%2 != 0 { + return nil, fmt.Errorf("gpt-oss expert tensor %q output dim = %d, want even gate/up rows", td.Name, outDim) + } + mid := outDim / 2 + + raw, err := io.ReadAll(td.Reader()) + if err != nil { + return nil, fmt.Errorf("read gpt-oss expert tensor %q: %w", td.Name, err) + } + values, err := DecodeFloatTensor(td.Dtype, raw) + if err != nil { + return nil, fmt.Errorf("decode gpt-oss expert tensor %q: %w", td.Name, err) + } + + gateVals := make([]float32, experts*mid) + upVals := make([]float32, experts*mid) + for e := range experts { + for row := range outDim { + src := e*outDim + row + dst := e*mid + row/2 + if row%2 == 0 { + gateVals[dst] = values[src] + } else { + upVals[dst] = values[src] + } + } + } + + gateRaw, err := EncodeFloatTensor("BF16", gateVals) + if err != nil { + return nil, fmt.Errorf("encode gate expert bias %q: %w", td.Name, err) + } + upRaw, err := EncodeFloatTensor("BF16", upVals) + if err != nil { + return nil, fmt.Errorf("encode up expert bias %q: %w", td.Name, err) + } + + gateName := strings.Replace(td.Name, "gate_up_proj", "gate_proj", 1) + upName := strings.Replace(td.Name, "gate_up_proj", "up_proj", 1) + shape := []int32{int32(experts), int32(mid)} + return []*safetensors.TensorData{ + safetensors.NewTensorDataFromBytes(gateName, "BF16", shape, gateRaw), + safetensors.NewTensorDataFromBytes(upName, "BF16", shape, upRaw), + }, nil +} + +func (t *gptossImportTransform) canonicalTensorName(name string) string { + switch name { + case "model.embed_tokens.weight": + return "embedding.weight" + case "model.embed_tokens.scales": + return "embedding.weight.scale" + case "model.embed_tokens.biases": + return "embedding.weight.bias" + case "model.norm.weight": + return "output_norm.weight" + case "lm_head.weight": + return "output.weight" + case "lm_head.scales": + return "output.weight.scale" + case "lm_head.biases": + return "output.weight.bias" + } + + const layerPrefix = "model.layers." + if !strings.HasPrefix(name, layerPrefix) { + return name + } + + remainder := strings.TrimPrefix(name, layerPrefix) + layer, suffix, ok := strings.Cut(remainder, ".") + if !ok || layer == "" { + return name + } + + prefix := "blocks." + layer + "." + switch suffix { + case "input_layernorm.weight": + return prefix + "attn_norm.weight" + case "self_attn.q_proj.weight": + return prefix + "q_proj.weight" + case "self_attn.q_proj.bias": + return prefix + "q_proj.bias" + case "self_attn.q_proj.scales": + return prefix + "q_proj.weight.scale" + case "self_attn.q_proj.biases": + return prefix + "q_proj.weight.bias" + case "self_attn.k_proj.weight": + return prefix + "k_proj.weight" + case "self_attn.k_proj.bias": + return prefix + "k_proj.bias" + case "self_attn.k_proj.scales": + return prefix + "k_proj.weight.scale" + case "self_attn.k_proj.biases": + return prefix + "k_proj.weight.bias" + case "self_attn.v_proj.weight": + return prefix + "v_proj.weight" + case "self_attn.v_proj.bias": + return prefix + "v_proj.bias" + case "self_attn.v_proj.scales": + return prefix + "v_proj.weight.scale" + case "self_attn.v_proj.biases": + return prefix + "v_proj.weight.bias" + case "self_attn.o_proj.weight": + return prefix + "attn_out.weight" + case "self_attn.o_proj.bias": + return prefix + "attn_out.bias" + case "self_attn.o_proj.scales": + return prefix + "attn_out.weight.scale" + case "self_attn.o_proj.biases": + return prefix + "attn_out.weight.bias" + case "self_attn.sinks": + return prefix + "attn_sinks" + case "post_attention_layernorm.weight": + return prefix + "ffn_norm.weight" + case "mlp.router.weight": + return prefix + "router.weight" + case "mlp.router.bias": + return prefix + "router.bias" + case "mlp.router.scales": + return prefix + "router.weight.scale" + case "mlp.router.biases": + return prefix + "router.weight.bias" + case "mlp.experts.gate_up_proj_blocks": + return prefix + "experts.gate_up_proj.weight" + case "mlp.experts.gate_up_proj_scales": + return prefix + "experts.gate_up_proj.weight" + case "mlp.experts.gate_up_proj_bias": + return prefix + "experts.gate_up_proj.bias" + case "mlp.experts.down_proj_blocks": + return prefix + "experts.down_proj.weight" + case "mlp.experts.down_proj_scales": + return prefix + "experts.down_proj.weight" + case "mlp.experts.down_proj_bias": + return prefix + "experts.down_proj.bias" + case "mlp.experts.gate_proj.weight", + "mlp.experts.gate_proj.scales", + "mlp.experts.gate_proj.biases": + return prefix + "experts.gate_proj.weight" + case "mlp.experts.gate_proj.bias": + return prefix + "experts.gate_proj.bias" + case "mlp.experts.up_proj.weight", + "mlp.experts.up_proj.scales", + "mlp.experts.up_proj.biases": + return prefix + "experts.up_proj.weight" + case "mlp.experts.up_proj.bias": + return prefix + "experts.up_proj.bias" + case "mlp.experts.down_proj.weight", + "mlp.experts.down_proj.scales", + "mlp.experts.down_proj.biases": + return prefix + "experts.down_proj.weight" + case "mlp.experts.down_proj.bias": + return prefix + "experts.down_proj.bias" + default: + return name + } +} diff --git a/x/create/plan.go b/x/create/plan.go index d754859aa69..35379bbd4bf 100644 --- a/x/create/plan.go +++ b/x/create/plan.go @@ -46,6 +46,23 @@ const ( // sources are the N weights followed by the N scales, in expert-index order; // the result is a BF16 tensor, which Quantize (if set) then re-quantizes. TransformDecodeStackFP8 Transform = "decode_stack_fp8" + + // TransformGPTOSSPackedExpertWeight converts one native GPT-OSS MXFP4 + // block/scale pair into the runtime-ready packed weight tensor. + TransformGPTOSSPackedExpertWeight Transform = "gptoss_packed_expert_weight" + + // TransformGPTOSSPackedExpertScale converts one native GPT-OSS MXFP4 + // block/scale pair into the runtime-ready scale companion. + TransformGPTOSSPackedExpertScale Transform = "gptoss_packed_expert_scale" + + // TransformGPTOSSGateUpWeight and TransformGPTOSSUpWeight split the + // interleaved native gate/up MXFP4 tensor into runtime-ready projections. + TransformGPTOSSGateUpWeight Transform = "gptoss_gate_up_gate_weight" + TransformGPTOSSUpWeight Transform = "gptoss_gate_up_up_weight" + TransformGPTOSSGateUpScale Transform = "gptoss_gate_up_gate_scale" + TransformGPTOSSUpScale Transform = "gptoss_gate_up_up_scale" + TransformGPTOSSGateUpBias Transform = "gptoss_gate_up_gate_bias" + TransformGPTOSSUpBias Transform = "gptoss_gate_up_up_bias" ) // TensorSpec describes one output tensor within a blob: the source tensor(s) From 355028110999bc5455d07fb379045d8741558397 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:22:18 +0000 Subject: [PATCH 13/58] gptoss: add the MLX model configuration contract Co-authored-by: Codex --- x/models/gptoss/gptoss.go | 283 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 x/models/gptoss/gptoss.go diff --git a/x/models/gptoss/gptoss.go b/x/models/gptoss/gptoss.go new file mode 100644 index 00000000000..644d69f9495 --- /dev/null +++ b/x/models/gptoss/gptoss.go @@ -0,0 +1,283 @@ +// Package gptoss provides the gpt-oss text model implementation for MLX. +package gptoss + +import ( + "encoding/json" + "fmt" + "math" + "strings" + + "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/tokenizer" +) + +func init() { + base.Register("GptOssForCausalLM", NewModel) +} + +// RopeScaling carries the GPT-OSS rope scaling block. +type RopeScaling struct { + Factor float32 `json:"factor"` + OriginalMaxPositionEmbeddings int32 `json:"original_max_position_embeddings"` + RopeType string `json:"rope_type,omitempty"` + BetaFast float32 `json:"beta_fast,omitempty"` + BetaSlow float32 `json:"beta_slow,omitempty"` + Truncate bool `json:"truncate,omitempty"` +} + +// Quantization carries optional quantization metadata from config.json. +type Quantization struct { + Bits int `json:"bits"` + GroupSize int `json:"group_size"` + Mode string `json:"mode"` + QuantMethod string `json:"quant_method"` +} + +// Config holds the gpt-oss model configuration. +type Config struct { + Architecture string `json:"-"` + ModelType string `json:"model_type"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + NumLocalExperts int32 `json:"num_local_experts"` + NumExpertsPerTok int32 `json:"num_experts_per_tok"` + SlidingWindow int32 `json:"sliding_window"` + RopeTheta float32 `json:"rope_theta"` + RopeScaling RopeScaling `json:"rope_scaling"` + RMSNormEps float32 `json:"rms_norm_eps"` + VocabSize int32 `json:"vocab_size"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + + Quantization Quantization `json:"quantization"` + QuantizationConfig Quantization `json:"quantization_config"` + + QuantGroupSize int `json:"-"` + QuantBits int `json:"-"` + QuantMode string `json:"-"` + TensorQuant map[string]*model.TensorQuantInfo `json:"-"` + QuantMethod string `json:"-"` +} + +// Model is the gpt-oss text-only model. +type Model struct { + tok *tokenizer.Tokenizer + *Config +} + +// NewModel creates a gpt-oss 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 + } + if cfg.QuantMethod == "" { + cfg.QuantMethod = strings.ToLower(qt) + } + } else { + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams("") + } + cfg.TensorQuant = root.AllTensorQuant() + if cfg.QuantMethod == "" { + if cfg.QuantizationConfig.QuantMethod != "" { + cfg.QuantMethod = strings.ToLower(cfg.QuantizationConfig.QuantMethod) + } else if cfg.Quantization.QuantMethod != "" { + cfg.QuantMethod = strings.ToLower(cfg.Quantization.QuantMethod) + } + } + + 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 + } + if specialTokensMapData, err := root.Manifest.ReadConfig("special_tokens_map.json"); err == nil { + tokConfig.SpecialTokensMapJSON = specialTokensMapData + } + + tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) + if err != nil { + return nil, fmt.Errorf("parse tokenizer: %w", err) + } + + return &Model{ + Config: &cfg, + tok: tok, + }, nil +} + +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 + active := configData + if textRaw, ok := raw["text_config"]; ok { + active = textRaw + } + if err := json.Unmarshal(active, &cfg); err != nil { + return Config{}, fmt.Errorf("parse config: %w", err) + } + + var archConfig struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + } + if err := json.Unmarshal(configData, &archConfig); err != nil { + return Config{}, fmt.Errorf("parse architecture: %w", err) + } + if len(archConfig.Architectures) > 0 && archConfig.Architectures[0] != "" { + cfg.Architecture = archConfig.Architectures[0] + } else { + cfg.Architecture = archConfig.ModelType + } + if cfg.Architecture == "" { + return Config{}, fmt.Errorf("missing architecture in config.json") + } + + if cfg.HiddenSize <= 0 { + return Config{}, fmt.Errorf("invalid hidden_size: %d", cfg.HiddenSize) + } + if cfg.IntermediateSize <= 0 { + return Config{}, fmt.Errorf("invalid intermediate_size: %d", cfg.IntermediateSize) + } + 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 { + return Config{}, fmt.Errorf("invalid num_key_value_heads: %d", cfg.NumKeyValueHeads) + } + if cfg.HeadDim <= 0 { + return Config{}, fmt.Errorf("invalid head_dim: %d", cfg.HeadDim) + } + if cfg.NumLocalExperts <= 0 { + return Config{}, fmt.Errorf("invalid num_local_experts: %d", cfg.NumLocalExperts) + } + if cfg.NumExpertsPerTok <= 0 { + return Config{}, fmt.Errorf("invalid num_experts_per_tok: %d", cfg.NumExpertsPerTok) + } + if cfg.SlidingWindow <= 0 { + return Config{}, fmt.Errorf("invalid sliding_window: %d", cfg.SlidingWindow) + } + if cfg.RopeTheta <= 0 { + return Config{}, fmt.Errorf("invalid rope_theta: %f", cfg.RopeTheta) + } + if cfg.RopeScaling.Factor <= 0 { + return Config{}, fmt.Errorf("invalid rope_scaling.factor: %f", cfg.RopeScaling.Factor) + } + if cfg.RopeScaling.OriginalMaxPositionEmbeddings <= 0 { + return Config{}, fmt.Errorf("invalid rope_scaling.original_max_position_embeddings: %d", cfg.RopeScaling.OriginalMaxPositionEmbeddings) + } + if cfg.RMSNormEps <= 0 { + return Config{}, fmt.Errorf("invalid rms_norm_eps: %f", cfg.RMSNormEps) + } + if cfg.VocabSize <= 0 { + return Config{}, fmt.Errorf("invalid vocab_size: %d", cfg.VocabSize) + } + if cfg.MaxPositionEmbeddings <= 0 { + cfg.MaxPositionEmbeddings = int32(math.Round(float64(cfg.RopeScaling.Factor) * float64(cfg.RopeScaling.OriginalMaxPositionEmbeddings))) + } + if cfg.MaxPositionEmbeddings <= 0 { + cfg.MaxPositionEmbeddings = cfg.SlidingWindow + } + if cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return Config{}, fmt.Errorf("num_attention_heads (%d) must be divisible by num_key_value_heads (%d)", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) + } + + if cfg.QuantizationConfig.QuantMethod != "" { + cfg.QuantMethod = strings.ToLower(cfg.QuantizationConfig.QuantMethod) + } else if cfg.Quantization.QuantMethod != "" { + cfg.QuantMethod = strings.ToLower(cfg.Quantization.QuantMethod) + } + + return cfg, nil +} + +// Forward is intentionally skeletal for Phase 2. +func (m *Model) Forward(_ *batch.Batch, _ []cache.Cache) (*mlx.Array, *mlx.Array) { + return nil, nil +} + +// Unembed is intentionally skeletal for Phase 2. +func (m *Model) Unembed(_ *mlx.Array) *mlx.Array { + return nil +} + +// NumLayers returns the configured layer count. +func (m *Model) NumLayers() int { + if m == nil || m.Config == nil { + return 0 + } + return int(m.NumHiddenLayers) +} + +// NewCaches declares the alternating sliding-window and full-attention cache +// slots owned by GPT-OSS. +func (m *Model) NewCaches() []cache.Cache { + caches := make([]cache.Cache, m.NumLayers()) + for i := range caches { + if i%2 == 0 { + caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow)) + } else { + caches[i] = cache.NewKVCache() + } + } + return caches +} + +// Tokenizer returns the loaded tokenizer. +func (m *Model) Tokenizer() *tokenizer.Tokenizer { + if m == nil { + return nil + } + return m.tok +} + +// MaxContextLength returns the derived context length. +func (m *Model) MaxContextLength() int { + if m == nil || m.Config == nil { + return 0 + } + if m.MaxPositionEmbeddings > 0 { + return int(m.MaxPositionEmbeddings) + } + return 0 +} + +// LoadWeights is intentionally skeletal for Phase 2. +func (m *Model) LoadWeights(map[string]*mlx.Array) error { + return fmt.Errorf("gpt-oss weight loading is not implemented yet") +} From f00bf638f79ffacd3e3f3a65cea43666473e6587 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:24:30 +0000 Subject: [PATCH 14/58] mlxrunner: register the GPT-OSS model runtime Co-authored-by: Codex --- x/mlxrunner/imports.go | 1 + 1 file changed, 1 insertion(+) diff --git a/x/mlxrunner/imports.go b/x/mlxrunner/imports.go index 540d8352d74..5084ee3b198 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -6,6 +6,7 @@ import ( _ "github.com/ollama/ollama/x/models/gemma4" _ "github.com/ollama/ollama/x/models/glimmer" _ "github.com/ollama/ollama/x/models/glm4_moe_lite" + _ "github.com/ollama/ollama/x/models/gptoss" _ "github.com/ollama/ollama/x/models/laguna" _ "github.com/ollama/ollama/x/models/llama" _ "github.com/ollama/ollama/x/models/nemotron_h" From 428cd93b3f6963bba6f55a7ea645e3c2544e710b Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:26:50 +0000 Subject: [PATCH 15/58] create: recognize GPT-OSS safetensors checkpoints Co-authored-by: Codex --- x/create/client/create.go | 7 +++++++ x/create/client/create_test.go | 15 +++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/x/create/client/create.go b/x/create/client/create.go index ae35caa2018..e3b87db74d9 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -590,6 +590,11 @@ func isQwen4Family(s string) bool { strings.Contains(s, "qwen4_exp") } +func isGPTOSSFamily(s string) bool { + s = strings.ToLower(s) + return strings.Contains(s, "gptoss") || strings.Contains(s, "gpt_oss") || strings.Contains(s, "gpt-oss") +} + func qwen35RendererName(modelDir string) string { template := readChatTemplate(modelDir) if strings.Contains(template, "resolved_reasoning_effort") && @@ -676,6 +681,8 @@ func parserNameForIdentifier(modelDir, s string) string { return lagunaRendererParserName(modelDir) case strings.Contains(s, "cohere2moe") || strings.Contains(s, "cohere2_moe"): return "cohere" + case isGPTOSSFamily(s): + return "harmony" case strings.Contains(s, "glm4") || strings.Contains(s, "glm-4"): return "glm-4.7" case strings.Contains(s, "deepseek"): diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 607f6b7d40f..e48950f63e5 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -816,6 +816,21 @@ func TestGetParserName(t *testing.T) { configJSON: `{"model_type": "qwen3"}`, want: "qwen3", }, + { + name: "gpt-oss architecture", + configJSON: `{"architectures":["GptOssForCausalLM"],"model_type":"gpt_oss"}`, + want: "harmony", + }, + { + name: "gpt-oss model type", + configJSON: `{"model_type":"gpt-oss"}`, + want: "harmony", + }, + { + name: "gpt-oss nested llm model type", + configJSON: `{"model_type":"wrapper","llm_config":{"model_type":"gpt_oss"}}`, + want: "harmony", + }, { name: "laguna model", configJSON: `{"architectures": ["LagunaForCausalLM"], "model_type": "laguna"}`, From 2239da46621377c53533f544c70ec1793ff57985 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:29:10 +0000 Subject: [PATCH 16/58] create: preserve native GPT-OSS MXFP4 tensors Co-authored-by: Codex --- x/create/transform.go | 45 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/x/create/transform.go b/x/create/transform.go index ceb761cc898..329e441456b 100644 --- a/x/create/transform.go +++ b/x/create/transform.go @@ -49,6 +49,51 @@ func applyByteTransform(ts TensorSpec, sources []*safetensors.TensorData) (*safe case TransformStackExperts: return stackExpertTensors(ts.Name, ts.OutDtype, ts.OutShape, sources) + case TransformGPTOSSPackedExpertWeight, TransformGPTOSSPackedExpertScale: + if len(sources) != 2 { + return nil, fmt.Errorf("transform %s expects block+scale sources, got %d", ts.Transform, len(sources)) + } + out, err := preservePackedExpertProjection(ts.Name, sources[0], sources[1]) + if err != nil { + return nil, err + } + if ts.Transform == TransformGPTOSSPackedExpertWeight { + return out[0].WithName(ts.Name), nil + } + return out[1].WithName(ts.Name), nil + + case TransformGPTOSSGateUpWeight, TransformGPTOSSUpWeight, TransformGPTOSSGateUpScale, TransformGPTOSSUpScale: + if len(sources) != 2 { + return nil, fmt.Errorf("transform %s expects block+scale sources, got %d", ts.Transform, len(sources)) + } + out, err := preserveAndSplitGateUpTensor(ts.Name, sources[0], sources[1]) + if err != nil { + return nil, err + } + switch ts.Transform { + case TransformGPTOSSGateUpWeight: + return out[0].WithName(ts.Name), nil + case TransformGPTOSSGateUpScale: + return out[1].WithName(ts.Name), nil + case TransformGPTOSSUpWeight: + return out[2].WithName(ts.Name), nil + default: + return out[3].WithName(ts.Name), nil + } + + case TransformGPTOSSGateUpBias, TransformGPTOSSUpBias: + if len(sources) != 1 { + return nil, fmt.Errorf("transform %s expects 1 source, got %d", ts.Transform, len(sources)) + } + out, err := splitGateUpBiasTensor(sources[0]) + if err != nil { + return nil, err + } + if ts.Transform == TransformGPTOSSGateUpBias { + return out[0].WithName(ts.Name), nil + } + return out[1].WithName(ts.Name), nil + default: return nil, fmt.Errorf("transform %q requires the MLX writer path", ts.Transform) } From c74120e28f4976294f6d4988223d9ac7b31bb4b6 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:41:38 +0000 Subject: [PATCH 17/58] mlx: fuse GPT-OSS MXFP4 MoE decode Co-authored-by: Codex --- x/mlxrunner/mlx/gptoss_moe.go | 730 +++++++++++++++++++++++++++++ x/mlxrunner/mlx/gptoss_moe_test.go | 224 +++++++++ 2 files changed, 954 insertions(+) create mode 100644 x/mlxrunner/mlx/gptoss_moe.go create mode 100644 x/mlxrunner/mlx/gptoss_moe_test.go diff --git a/x/mlxrunner/mlx/gptoss_moe.go b/x/mlxrunner/mlx/gptoss_moe.go new file mode 100644 index 00000000000..46aee9142aa --- /dev/null +++ b/x/mlxrunner/mlx/gptoss_moe.go @@ -0,0 +1,730 @@ +package mlx + +// #include +// #include "generated.h" +import "C" + +import ( + "math" + "sync" + "unsafe" +) + +func validMoEPositiveInt(v int) bool { + return v > 0 && v <= math.MaxInt32 +} + +// checkedMoESpan returns an element span whose maximum zero-based index fits +// the uint32 index type used by the Metal kernels. +func checkedMoESpan(factors ...uint64) (uint64, bool) { + span := uint64(1) + for _, factor := range factors { + if factor == 0 || span > math.MaxUint64/factor { + return 0, false + } + span *= factor + } + if span > uint64(math.MaxUint32)+1 { + return 0, false + } + return span, true +} + +func validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK int, down bool) bool { + values := []int{experts, batch, numRows, numColVecs, topK} + factors := make([]uint64, len(values)) + for i, value := range values { + if !validMoEPositiveInt(value) { + return false + } + factors[i] = uint64(value) + } + e, b, r, c, k := factors[0], factors[1], factors[2], factors[3], factors[4] + + // Complete spans prove expertStride, expert/row/channel terms, companion + // offsets, selector offsets, and output offsets without uint32 wrapping. + for _, product := range [][]uint64{ + {r, c}, + {e, r, c}, + {e, r}, + {b, k}, + {b, k, r}, + } { + if _, ok := checkedMoESpan(product...); !ok { + return false + } + } + + inputSpan := []uint64{b, c, 8} + if down { + inputSpan = []uint64{b, k, c, 8} + } + _, ok := checkedMoESpan(inputSpan...) + return ok +} + +func validateMoEExpertIDs(expertIds *Array, experts int) bool { + if expertIds == nil || expertIds.DType() != DTypeUint32 || expertIds.Size() == 0 || !validMoEPositiveInt(experts) { + return false + } + Eval(expertIds) + data := C.mlx_array_data_uint32(expertIds.ctx) + if data == nil { + return false + } + for _, expertID := range unsafe.Slice(data, expertIds.Size()) { + if uint64(expertID) >= uint64(experts) { + return false + } + } + return true +} + +func validMoEArray(a *Array, dtype DType, dims ...int) bool { + if a == nil || a.DType() != dtype || len(dims) != a.NumDims() { + return false + } + size := 1 + for i, dim := range dims { + if !validMoEPositiveInt(dim) || a.Dim(i) != dim || size > math.MaxInt/dim { + return false + } + size *= dim + } + return a.Size() == size +} + +func validateMoEProjection(weight, scales, bias *Array, numRows, numColVecs int) (int, bool) { + if weight == nil || scales == nil || bias == nil || !validMoEPositiveInt(numRows) || !validMoEPositiveInt(numColVecs) || numColVecs < 32 || numColVecs > math.MaxInt/32 || numColVecs > math.MaxInt/4 { + return 0, false + } + dims := weight.Dims() + if len(dims) != 3 || !validMoEPositiveInt(dims[0]) { + return 0, false + } + experts := dims[0] + if !validMoEArray(weight, DTypeUint32, experts, numRows, numColVecs*4) || + !validMoEArray(scales, DTypeUint8, experts, numRows, numColVecs) || + !validMoEArray(bias, DTypeBFloat16, experts, numRows) { + return 0, false + } + // The Metal kernel is created with ensure_row_contiguous=true. Together + // with these exact shapes and sizes, MLX supplies the linear row strides + // used by the kernel instead of exposing unchecked source strides here. + return experts, true +} + +func validateMoEGateUpInputs( + input, gateWeight, gateScales, gateBias, + upWeight, upScales, upBias, expertIds *Array, + numRows, numColVecs, topK int, +) (int, bool) { + if !validMoEPositiveInt(topK) { + return 0, false + } + experts, ok := validateMoEProjection(gateWeight, gateScales, gateBias, numRows, numColVecs) + if !ok || topK > experts { + return 0, false + } + upExperts, ok := validateMoEProjection(upWeight, upScales, upBias, numRows, numColVecs) + if !ok || upExperts != experts { + return 0, false + } + inDims := input.Dims() + if len(inDims) != 2 || !validMoEPositiveInt(inDims[0]) { + return 0, false + } + batch := inDims[0] + if !validMoEArray(input, DTypeFloat32, batch, numColVecs*32) || + !validMoEArray(expertIds, DTypeUint32, batch, topK) { + return 0, false + } + if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, false) || + !validateMoEExpertIDs(expertIds, experts) { + return 0, false + } + return batch, true +} + +func validateMoEDownInputs(input, weight, scales, bias, expertIds *Array, numRows, numColVecs, topK int) (int, bool) { + if !validMoEPositiveInt(topK) { + return 0, false + } + experts, ok := validateMoEProjection(weight, scales, bias, numRows, numColVecs) + if !ok || topK > experts { + return 0, false + } + inDims := input.Dims() + if len(inDims) != 3 || !validMoEPositiveInt(inDims[0]) { + return 0, false + } + batch := inDims[0] + if !validMoEArray(input, DTypeFloat32, batch, topK, numColVecs*32) || + !validMoEArray(expertIds, DTypeUint32, batch, topK) { + return 0, false + } + if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, true) || + !validateMoEExpertIDs(expertIds, experts) { + return 0, false + } + return batch, true +} + +var ( + moeSwiGLUOnce sync.Once + moeSwiGLUKernel C.mlx_fast_metal_kernel + moeSwiGLUDisabled bool + + moeDownOnce sync.Once + moeDownKernel C.mlx_fast_metal_kernel + moeDownDisabled bool +) + +// moeSwiGLUSource is a fused gate+up+SwiGLU Metal kernel for MXFP4 MoE experts. +// Adapted from the gpt-oss reference kernel gptoss_f32_mf4w_moe_matmul_swiglu. +// +// Each threadgroup processes nSg rows total (even simdgroups compute gate, +// odd compute up). After a barrier, simdgroup pairs apply SwiGLU to produce +// nSg/2 output channels per threadgroup. +// +// Template args: NumColVecs (int), NumRows (int), NumTopK (int) +// Inputs (named): input(float), gate_w(uint), gate_s(bfloat), gate_b(bfloat), +// +// up_w(uint), up_s(bfloat), up_b(bfloat), +// expert_ids(uint), swiglu_params(float) +// +// Output (named): output(float) +const moeSwiGLUSource = ` +uint3 gid = threadgroup_position_in_grid; +uint sgTid = thread_index_in_simdgroup; +uint sgIdx = simdgroup_index_in_threadgroup; +uint nSg = simdgroups_per_threadgroup; +constexpr uint sgSize = 32; + +threadgroup float tg_buf[32]; + +uint outCh = gid.x * (nSg / 2) + (sgIdx / 2); +uint expertId = expert_ids[gid.y * NumTopK + gid.z]; + +uint rowStride = (uint)NumColVecs; +uint expertStride = (uint)NumRows * rowStride; + +const device uint4* w; +const device uint8_t* ws; +const device bfloat* wb; + +if (sgIdx % 2 == 0) { + w = reinterpret_cast(gate_w) + expertId * expertStride + outCh * rowStride + sgTid; + ws = gate_s + expertId * (uint)NumRows * (uint)NumColVecs + outCh * (uint)NumColVecs + sgTid; + wb = gate_b + expertId * (uint)NumRows + outCh; +} else { + w = reinterpret_cast(up_w) + expertId * expertStride + outCh * rowStride + sgTid; + ws = up_s + expertId * (uint)NumRows * (uint)NumColVecs + outCh * (uint)NumColVecs + sgTid; + wb = up_b + expertId * (uint)NumRows + outCh; +} + +auto inp = reinterpret_cast(input) + 8 * (gid.y * (uint)NumColVecs + sgTid); + +uint numIter = ((uint)NumColVecs - sgTid + (sgSize - 1)) / sgSize; + +float4 sum4 = 0.0f; +do { + // Read 32 packed FP4 E2M1 values (uint4 = 4 × uint32) and E8M0 scale. + // Each uint32 holds 8 nibbles: bits [3:0],[7:4],...,[31:28]. + // Dequant: (nibble & 7) << 9 → half, * 16384, sign from bit 3. + const uint4 pk = *w; + const float wscale = as_type(uint(*ws) << 23); + + // Helper macro: extract 8 FP4 values from a uint32 into two float4 vectors. + // lo = values from nibbles 0-3 (bits 0-15), hi = values from nibbles 4-7 (bits 16-31). +#define DEQUANT_U32(u, lo, hi) { \ + half _h0 = as_type(ushort((u & 0x7u) << 9)) * 16384.0h; \ + half _h1 = as_type(ushort(((u >> 4) & 0x7u) << 9)) * 16384.0h; \ + half _h2 = as_type(ushort(((u >> 8) & 0x7u) << 9)) * 16384.0h; \ + half _h3 = as_type(ushort(((u >> 12) & 0x7u) << 9)) * 16384.0h; \ + lo = float4((u & 0x8u) ? -float(_h0) : float(_h0), \ + ((u >> 4) & 0x8u) ? -float(_h1) : float(_h1), \ + ((u >> 8) & 0x8u) ? -float(_h2) : float(_h2), \ + ((u >> 12) & 0x8u) ? -float(_h3) : float(_h3)); \ + half _h4 = as_type(ushort(((u >> 16) & 0x7u) << 9)) * 16384.0h; \ + half _h5 = as_type(ushort(((u >> 20) & 0x7u) << 9)) * 16384.0h; \ + half _h6 = as_type(ushort(((u >> 24) & 0x7u) << 9)) * 16384.0h; \ + half _h7 = as_type(ushort(((u >> 28) & 0x7u) << 9)) * 16384.0h; \ + hi = float4(((u >> 16) & 0x8u) ? -float(_h4) : float(_h4), \ + ((u >> 20) & 0x8u) ? -float(_h5) : float(_h5), \ + ((u >> 24) & 0x8u) ? -float(_h6) : float(_h6), \ + ((u >> 28) & 0x8u) ? -float(_h7) : float(_h7)); \ +} + + float4 wv0123, wv4567, wv89AB, wvCDEF, wvGHIJ, wvKLMN, wvOPQR, wvSTUV; + DEQUANT_U32(pk.x, wv0123, wv4567) // values 0-7 + DEQUANT_U32(pk.y, wv89AB, wvCDEF) // values 8-15 + DEQUANT_U32(pk.z, wvGHIJ, wvKLMN) // values 16-23 + DEQUANT_U32(pk.w, wvOPQR, wvSTUV) // values 24-31 +#undef DEQUANT_U32 + + const float4 i0123 = inp[0]; + const float4 i4567 = inp[1]; + const float4 i89AB = inp[2]; + const float4 iCDEF = inp[3]; + const float4 iGHIJ = inp[4]; + const float4 iKLMN = inp[5]; + const float4 iOPQR = inp[6]; + const float4 iSTUV = inp[7]; + + float4 psum0 = i0123 * wv0123; + float4 psum1 = i4567 * wv4567; + psum0 = fma(i89AB, wv89AB, psum0); + psum1 = fma(iCDEF, wvCDEF, psum1); + psum0 = fma(iGHIJ, wvGHIJ, psum0); + psum1 = fma(iKLMN, wvKLMN, psum1); + psum0 = fma(iOPQR, wvOPQR, psum0); + psum1 = fma(iSTUV, wvSTUV, psum1); + sum4 = fma(psum0, wscale, sum4); + sum4 = fma(psum1, wscale, sum4); + + w += sgSize; + ws += sgSize; + inp += 8 * sgSize; +} while (--numIter != 0); + +const float2 sum2 = sum4.xy + sum4.zw; +float sum = sum2.x + sum2.y; +sum = simd_sum(sum); + +if (simd_is_first()) { + sum += static_cast(*wb); + tg_buf[sgIdx] = sum; +} +threadgroup_barrier(mem_flags::mem_threadgroup); + +uint tid = sgIdx * sgSize + sgTid; +if (tid * 2 < nSg) { + const float2 gu = reinterpret_cast(tg_buf)[tid]; + const float smin = swiglu_params[0]; + const float smax = swiglu_params[1]; + const float gate_val = min(gu.x, smax); + const float up_val = clamp(gu.y, smin, smax); + const float alpha = 1.702f; + const float swish = gate_val / (1.0f + precise::exp(-alpha * gate_val)); + const float result = fma(swish, up_val, swish); + + uint ch = gid.x * (nSg / 2) + tid; + output[gid.y * (uint)NumTopK * (uint)NumRows + gid.z * (uint)NumRows + ch] = result; +} +` + +func initMoESwiGLUKernel() { + inputs, freeInputs, ok := cStringVector([]string{ + "input", "gate_w", "gate_s", "gate_b", + "up_w", "up_s", "up_b", + "expert_ids", "swiglu_params", + }) + if !ok { + moeSwiGLUDisabled = true + freeInputs() + return + } + defer freeInputs() + + outputs, freeOutputs, ok := cStringVector([]string{"output"}) + if !ok { + moeSwiGLUDisabled = true + freeOutputs() + return + } + defer freeOutputs() + + cName := C.CString("gptoss_moe_swiglu") + defer C.free(unsafe.Pointer(cName)) + cSource := C.CString(moeSwiGLUSource) + defer C.free(unsafe.Pointer(cSource)) + cHeader := C.CString("") + defer C.free(unsafe.Pointer(cHeader)) + + moeSwiGLUKernel = C.mlx_fast_metal_kernel_new( + cName, + inputs, + outputs, + cSource, + cHeader, + C.bool(true), // ensure_row_contiguous + C.bool(false), // atomic_outputs + ) +} + +// MoEFusedGateUpSwiGLU runs a fused gate+up+SwiGLU kernel for MXFP4 MoE experts. +// It returns (result, true) on success, or (nil, false) if the kernel is unavailable +// or the inputs are incompatible. +// +// Parameters: +// - input: float32 [batch, hiddenSize] — flattened input hidden states +// - gateWeight, gateScales, gateBias: MXFP4 gate projection weights +// - upWeight, upScales, upBias: MXFP4 up projection weights +// - expertIds: uint32 [batch, topK] — expert indices per token +// - numRows: output dimension of gate/up projections +// - numColVecs: hiddenSize / 32 (MXFP4 groups per row) +// - topK: number of active experts per token +// - swigluMin, swigluMax: SwiGLU clamp parameters +func MoEFusedGateUpSwiGLU( + input, gateWeight, gateScales, gateBias, + upWeight, upScales, upBias, expertIds *Array, + numRows, numColVecs, topK int, + swigluMin, swigluMax float32, +) (*Array, bool) { + if moeSwiGLUDisabled { + return nil, false + } + if input == nil || gateWeight == nil || gateScales == nil || gateBias == nil || + upWeight == nil || upScales == nil || upBias == nil || expertIds == nil { + return nil, false + } + + batch, ok := validateMoEGateUpInputs(input, gateWeight, gateScales, gateBias, upWeight, upScales, upBias, expertIds, numRows, numColVecs, topK) + if !ok { + return nil, false + } + + moeSwiGLUOnce.Do(initMoESwiGLUKernel) + if moeSwiGLUDisabled { + return nil, false + } + + // Choose simdgroup count (must be even, numRows divisible by nSg/2) + nSg := 8 + for nSg > 2 && numRows%(nSg/2) != 0 { + nSg /= 2 + } + if nSg < 2 || numRows%(nSg/2) != 0 { + return nil, false + } + if numRows > math.MaxInt32/(32*nSg) { + return nil, false + } + + // Configure template args + cfg := C.mlx_fast_metal_kernel_config_new() + defer C.mlx_fast_metal_kernel_config_free(cfg) + + for _, tpl := range []struct { + name string + value int + }{ + {"NumColVecs", numColVecs}, + {"NumRows", numRows}, + {"NumTopK", topK}, + } { + cn := C.CString(tpl.name) + rc := C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, cn, C.int(tpl.value)) + C.free(unsafe.Pointer(cn)) + if rc != 0 { + moeSwiGLUDisabled = true + return nil, false + } + } + + // Output shape: [batch, topK, numRows] + outShape := []C.int{C.int(batch), C.int(topK), C.int(numRows)} + if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(outShape), C.size_t(len(outShape)), C.mlx_dtype(DTypeFloat32)) != 0 { + moeSwiGLUDisabled = true + return nil, false + } + + // MLX dispatch_threads takes total thread count, not threadgroup count. + // We need numRows/(nSg/2) threadgroups, each with 32*nSg threads. + tgSize := 32 * nSg + numTGx := numRows / (nSg / 2) + gridX := numTGx * tgSize + if C.mlx_fast_metal_kernel_config_set_grid(cfg, C.int(gridX), C.int(batch), C.int(topK)) != 0 { + moeSwiGLUDisabled = true + return nil, false + } + + // Threadgroup: (32 * nSg, 1, 1) + if C.mlx_fast_metal_kernel_config_set_thread_group(cfg, C.int(tgSize), 1, 1) != 0 { + moeSwiGLUDisabled = true + return nil, false + } + + // Build swiglu params array + swigluParams := FromValues([]float32{swigluMin, swigluMax}, 2) + + // Assemble inputs + inputs := []C.mlx_array{ + input.ctx, + gateWeight.ctx, + gateScales.ctx, + gateBias.ctx, + upWeight.ctx, + upScales.ctx, + upBias.ctx, + expertIds.ctx, + swigluParams.ctx, + } + inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) + defer C.mlx_vector_array_free(inVec) + + outVec := C.mlx_vector_array_new() + defer C.mlx_vector_array_free(outVec) + if C.mlx_fast_metal_kernel_apply(&outVec, moeSwiGLUKernel, inVec, cfg, DefaultStream().ctx) != 0 { + moeSwiGLUDisabled = true + return nil, false + } + if int(C.mlx_vector_array_size(outVec)) < 1 { + return nil, false + } + + out := New("MOE_FUSED_SWIGLU") + C.mlx_vector_array_get(&out.ctx, outVec, 0) + return out, true +} + +// moeDownSource is a fused down-projection Metal kernel for MXFP4 MoE experts. +// Adapted from the gpt-oss reference kernel gptoss_f32_mf4w_moe_matmul. +// +// Each simdgroup computes one output channel (dot-product of input row against +// one MXFP4-packed weight row, plus bias). One threadgroup produces nSg output +// channels. No SwiGLU — this is a plain matmul+bias. +// +// The input is per-expert: input[gid.y * NumTopK * NumInCols + gid.z * NumInCols + ...] +// where NumInCols is the intermediate_size (in float32 elements). +// +// Template args: NumColVecs (int), NumRows (int), NumTopK (int) +// Inputs (named): input(float), down_w(uint), down_s(bfloat), down_b(bfloat), +// +// expert_ids(uint) +// +// Output (named): output(float) +const moeDownSource = ` +uint3 gid = threadgroup_position_in_grid; +uint sgTid = thread_index_in_simdgroup; +uint sgIdx = simdgroup_index_in_threadgroup; +uint nSg = simdgroups_per_threadgroup; +constexpr uint sgSize = 32; + +uint outCh = gid.x * nSg + sgIdx; +uint expertId = expert_ids[gid.y * NumTopK + gid.z]; + +uint rowStride = (uint)NumColVecs; +uint expertStride = (uint)NumRows * rowStride; + +const device uint4* w = reinterpret_cast(down_w) + expertId * expertStride + outCh * rowStride + sgTid; +const device uint8_t* ws = down_s + expertId * (uint)NumRows * (uint)NumColVecs + outCh * (uint)NumColVecs + sgTid; +const device bfloat* wb = down_b + expertId * (uint)NumRows + outCh; + +auto inp = reinterpret_cast(input) + 8 * (gid.y * NumTopK * (uint)NumColVecs + gid.z * (uint)NumColVecs + sgTid); + +uint numIter = ((uint)NumColVecs - sgTid + (sgSize - 1)) / sgSize; + +float4 sum4 = 0.0f; +do { + const uint4 pk = *w; + const float wscale = as_type(uint(*ws) << 23); + +#define DEQUANT_U32(u, lo, hi) { \ + half _h0 = as_type(ushort((u & 0x7u) << 9)) * 16384.0h; \ + half _h1 = as_type(ushort(((u >> 4) & 0x7u) << 9)) * 16384.0h; \ + half _h2 = as_type(ushort(((u >> 8) & 0x7u) << 9)) * 16384.0h; \ + half _h3 = as_type(ushort(((u >> 12) & 0x7u) << 9)) * 16384.0h; \ + lo = float4((u & 0x8u) ? -float(_h0) : float(_h0), \ + ((u >> 4) & 0x8u) ? -float(_h1) : float(_h1), \ + ((u >> 8) & 0x8u) ? -float(_h2) : float(_h2), \ + ((u >> 12) & 0x8u) ? -float(_h3) : float(_h3)); \ + half _h4 = as_type(ushort(((u >> 16) & 0x7u) << 9)) * 16384.0h; \ + half _h5 = as_type(ushort(((u >> 20) & 0x7u) << 9)) * 16384.0h; \ + half _h6 = as_type(ushort(((u >> 24) & 0x7u) << 9)) * 16384.0h; \ + half _h7 = as_type(ushort(((u >> 28) & 0x7u) << 9)) * 16384.0h; \ + hi = float4(((u >> 16) & 0x8u) ? -float(_h4) : float(_h4), \ + ((u >> 20) & 0x8u) ? -float(_h5) : float(_h5), \ + ((u >> 24) & 0x8u) ? -float(_h6) : float(_h6), \ + ((u >> 28) & 0x8u) ? -float(_h7) : float(_h7)); \ +} + + float4 wv0123, wv4567, wv89AB, wvCDEF, wvGHIJ, wvKLMN, wvOPQR, wvSTUV; + DEQUANT_U32(pk.x, wv0123, wv4567) + DEQUANT_U32(pk.y, wv89AB, wvCDEF) + DEQUANT_U32(pk.z, wvGHIJ, wvKLMN) + DEQUANT_U32(pk.w, wvOPQR, wvSTUV) +#undef DEQUANT_U32 + + const float4 i0123 = inp[0]; + const float4 i4567 = inp[1]; + const float4 i89AB = inp[2]; + const float4 iCDEF = inp[3]; + const float4 iGHIJ = inp[4]; + const float4 iKLMN = inp[5]; + const float4 iOPQR = inp[6]; + const float4 iSTUV = inp[7]; + + float4 psum0 = i0123 * wv0123; + float4 psum1 = i4567 * wv4567; + psum0 = fma(i89AB, wv89AB, psum0); + psum1 = fma(iCDEF, wvCDEF, psum1); + psum0 = fma(iGHIJ, wvGHIJ, psum0); + psum1 = fma(iKLMN, wvKLMN, psum1); + psum0 = fma(iOPQR, wvOPQR, psum0); + psum1 = fma(iSTUV, wvSTUV, psum1); + sum4 = fma(psum0, wscale, sum4); + sum4 = fma(psum1, wscale, sum4); + + w += sgSize; + ws += sgSize; + inp += 8 * sgSize; +} while (--numIter != 0); + +const float2 sum2 = sum4.xy + sum4.zw; +float sum = sum2.x + sum2.y; +sum = simd_sum(sum); + +if (simd_is_first()) { + sum += static_cast(*wb); + output[gid.y * (uint)NumTopK * (uint)NumRows + gid.z * (uint)NumRows + outCh] = sum; +} +` + +func initMoEDownKernel() { + inputs, freeInputs, ok := cStringVector([]string{ + "input", "down_w", "down_s", "down_b", "expert_ids", + }) + if !ok { + moeDownDisabled = true + freeInputs() + return + } + defer freeInputs() + + outputs, freeOutputs, ok := cStringVector([]string{"output"}) + if !ok { + moeDownDisabled = true + freeOutputs() + return + } + defer freeOutputs() + + cName := C.CString("gptoss_moe_down") + defer C.free(unsafe.Pointer(cName)) + cSource := C.CString(moeDownSource) + defer C.free(unsafe.Pointer(cSource)) + cHeader := C.CString("") + defer C.free(unsafe.Pointer(cHeader)) + + moeDownKernel = C.mlx_fast_metal_kernel_new( + cName, + inputs, + outputs, + cSource, + cHeader, + C.bool(true), // ensure_row_contiguous + C.bool(false), // atomic_outputs + ) +} + +// MoEFusedDown runs a fused down-projection kernel for MXFP4 MoE experts. +// It returns (result, true) on success, or (nil, false) if the kernel is unavailable +// or the inputs are incompatible. +// +// Parameters: +// - input: float32 [batch, topK, intermediateSize] — per-expert SwiGLU output +// - downWeight, downScales, downBias: MXFP4 down projection weights +// - expertIds: uint32 [batch, topK] — expert indices per token +// - numRows: output dimension (hiddenSize) +// - numColVecs: intermediateSize / 32 (MXFP4 groups per row) +// - topK: number of active experts per token +func MoEFusedDown( + input, downWeight, downScales, downBias, expertIds *Array, + numRows, numColVecs, topK int, +) (*Array, bool) { + if moeDownDisabled { + return nil, false + } + if input == nil || downWeight == nil || downScales == nil || downBias == nil || expertIds == nil { + return nil, false + } + + batch, ok := validateMoEDownInputs(input, downWeight, downScales, downBias, expertIds, numRows, numColVecs, topK) + if !ok { + return nil, false + } + + moeDownOnce.Do(initMoEDownKernel) + if moeDownDisabled { + return nil, false + } + + // Each simdgroup computes one output channel. Choose nSg so numRows is divisible. + nSg := 8 + for nSg > 1 && numRows%nSg != 0 { + nSg /= 2 + } + if nSg < 1 || numRows%nSg != 0 { + return nil, false + } + if numRows > math.MaxInt32/(32*nSg) { + return nil, false + } + + cfg := C.mlx_fast_metal_kernel_config_new() + defer C.mlx_fast_metal_kernel_config_free(cfg) + + for _, tpl := range []struct { + name string + value int + }{ + {"NumColVecs", numColVecs}, + {"NumRows", numRows}, + {"NumTopK", topK}, + } { + cn := C.CString(tpl.name) + rc := C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, cn, C.int(tpl.value)) + C.free(unsafe.Pointer(cn)) + if rc != 0 { + moeDownDisabled = true + return nil, false + } + } + + // Output shape: [batch, topK, numRows] + outShape := []C.int{C.int(batch), C.int(topK), C.int(numRows)} + if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(outShape), C.size_t(len(outShape)), C.mlx_dtype(DTypeFloat32)) != 0 { + moeDownDisabled = true + return nil, false + } + + // Grid: numRows/nSg threadgroups in X, batch in Y, topK in Z + // Each threadgroup has 32*nSg threads + tgSize := 32 * nSg + numTGx := numRows / nSg + gridX := numTGx * tgSize + if C.mlx_fast_metal_kernel_config_set_grid(cfg, C.int(gridX), C.int(batch), C.int(topK)) != 0 { + moeDownDisabled = true + return nil, false + } + if C.mlx_fast_metal_kernel_config_set_thread_group(cfg, C.int(tgSize), 1, 1) != 0 { + moeDownDisabled = true + return nil, false + } + + inputs := []C.mlx_array{ + input.ctx, + downWeight.ctx, + downScales.ctx, + downBias.ctx, + expertIds.ctx, + } + inVec := C.mlx_vector_array_new_data(unsafe.SliceData(inputs), C.size_t(len(inputs))) + defer C.mlx_vector_array_free(inVec) + + outVec := C.mlx_vector_array_new() + defer C.mlx_vector_array_free(outVec) + if C.mlx_fast_metal_kernel_apply(&outVec, moeDownKernel, inVec, cfg, DefaultStream().ctx) != 0 { + moeDownDisabled = true + return nil, false + } + if int(C.mlx_vector_array_size(outVec)) < 1 { + return nil, false + } + + out := New("MOE_FUSED_DOWN") + C.mlx_vector_array_get(&out.ctx, outVec, 0) + return out, true +} diff --git a/x/mlxrunner/mlx/gptoss_moe_test.go b/x/mlxrunner/mlx/gptoss_moe_test.go new file mode 100644 index 00000000000..75c2cc59f57 --- /dev/null +++ b/x/mlxrunner/mlx/gptoss_moe_test.go @@ -0,0 +1,224 @@ +package mlx + +import ( + "math" + "testing" +) + +func TestCheckedMoESpan(t *testing.T) { + tests := []struct { + name string + factors []uint64 + want uint64 + ok bool + }{ + {name: "single", factors: []uint64{1}, want: 1, ok: true}, + {name: "maximum index", factors: []uint64{65536, 65536}, want: uint64(math.MaxUint32) + 1, ok: true}, + {name: "first past maximum index", factors: []uint64{65536, 65537}, ok: false}, + {name: "zero", factors: []uint64{1, 0}, ok: false}, + {name: "uint64 overflow", factors: []uint64{math.MaxUint64, 2}, ok: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := checkedMoESpan(tt.factors...) + if ok != tt.ok || got != tt.want { + t.Fatalf("checkedMoESpan(%v) = (%d, %v), want (%d, %v)", tt.factors, got, ok, tt.want, tt.ok) + } + }) + } +} + +func TestValidateMoEIndexSpans(t *testing.T) { + tests := []struct { + name string + experts, batch, rows, colVecs, topK int + down, want bool + }{ + {name: "valid ordinary gate", experts: 8, batch: 4, rows: 256, colVecs: 64, topK: 4, want: true}, + {name: "valid exact uint32 weight span", experts: 1, batch: 1, rows: 65536, colVecs: 65536, topK: 1, want: true}, + {name: "expert stride overflow", experts: 1, batch: 1, rows: 65536, colVecs: 65537, topK: 1}, + {name: "full expert weight scale overflow", experts: 65537, batch: 1, rows: 65536, colVecs: 1, topK: 1}, + {name: "full bias overflow", experts: 65536, batch: 1, rows: 65537, colVecs: 1, topK: 1}, + {name: "selector overflow", experts: 1, batch: 65536, rows: 1, colVecs: 1, topK: 65537}, + {name: "output overflow", experts: 1, batch: 65536, rows: 2, colVecs: 1, topK: 65536}, + {name: "gate input factor overflow", experts: 1, batch: 32768, rows: 1, colVecs: 16385, topK: 1}, + {name: "valid ordinary down", experts: 8, batch: 4, rows: 256, colVecs: 64, topK: 4, down: true, want: true}, + {name: "down input factor overflow", experts: 1, batch: 32768, rows: 1, colVecs: 16385, topK: 1, down: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := validateMoEIndexSpans(tt.experts, tt.batch, tt.rows, tt.colVecs, tt.topK, tt.down); got != tt.want { + t.Fatalf("validateMoEIndexSpans() = %v, want %v", got, tt.want) + } + }) + } +} + +type moeTestInputs struct { + input, weight, scales, bias, expertIDs *Array + upWeight, upScales, upBias *Array +} + +func newMoETestInputs(expertIDs []uint32) moeTestInputs { + const experts, rows, colVecs, topK = 2, 4, 32, 1 + batch := len(expertIDs) / topK + return moeTestInputs{ + input: Zeros(DTypeFloat32, batch, colVecs*32), + weight: Zeros(DTypeUint32, experts, rows, colVecs*4), + scales: Zeros(DTypeUint8, experts, rows, colVecs), + bias: Zeros(DTypeBFloat16, experts, rows), + expertIDs: FromValues(expertIDs, batch, topK), + upWeight: Zeros(DTypeUint32, experts, rows, colVecs*4), + upScales: Zeros(DTypeUint8, experts, rows, colVecs), + upBias: Zeros(DTypeBFloat16, experts, rows), + } +} + +func TestValidateMoEInputs(t *testing.T) { + skipIfNoMLX(t) + withMLXThread(t, func() { + valid := newMoETestInputs([]uint32{1}) + if _, ok := validateMoEGateUpInputs(valid.input, valid.weight, valid.scales, valid.bias, valid.upWeight, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1); !ok { + t.Fatal("valid gate/up inputs rejected") + } + + downInput := Zeros(DTypeFloat32, 1, 1, 32*32) + if _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1); !ok { + t.Fatal("valid down inputs rejected") + } + + for _, tt := range []struct { + name string + fn func() bool + }{ + {name: "expert equal to count", fn: func() bool { + bad := newMoETestInputs([]uint32{2}) + _, ok := validateMoEGateUpInputs(bad.input, bad.weight, bad.scales, bad.bias, bad.upWeight, bad.upScales, bad.upBias, bad.expertIDs, 4, 32, 1) + return ok + }}, + {name: "maximum expert", fn: func() bool { + bad := newMoETestInputs([]uint32{math.MaxUint32}) + _, ok := validateMoEDownInputs(downInput, bad.weight, bad.scales, bad.bias, bad.expertIDs, 4, 32, 1) + return ok + }}, + {name: "wrong selector dtype", fn: func() bool { + ids := FromValues([]int32{0}, 1, 1) + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, ids, 4, 32, 1) + return ok + }}, + {name: "wrong input rank", fn: func() bool { + input := Zeros(DTypeFloat32, 1, 1, 32*32) + _, ok := validateMoEGateUpInputs(input, valid.weight, valid.scales, valid.bias, valid.upWeight, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "wrong input dtype", fn: func() bool { + input := Zeros(DTypeBFloat16, 1, 32*32) + _, ok := validateMoEGateUpInputs(input, valid.weight, valid.scales, valid.bias, valid.upWeight, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "wrong weight dtype", fn: func() bool { + weight := Zeros(DTypeUint8, 2, 4, 32*4) + _, ok := validateMoEDownInputs(downInput, weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "wrong weight rank", fn: func() bool { + weight := Zeros(DTypeUint32, 2, 4*32*4) + _, ok := validateMoEDownInputs(downInput, weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "gate up expert mismatch", fn: func() bool { + up := Zeros(DTypeUint32, 3, 4, 32*4) + _, ok := validateMoEGateUpInputs(valid.input, valid.weight, valid.scales, valid.bias, up, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "malformed scale", fn: func() bool { + scale := Zeros(DTypeUint8, 2, 4, 31) + _, ok := validateMoEDownInputs(downInput, valid.weight, scale, valid.bias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "malformed bias", fn: func() bool { + bias := Zeros(DTypeBFloat16, 2, 5) + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, bias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "missing scale", fn: func() bool { + _, ok := validateMoEDownInputs(downInput, valid.weight, nil, valid.bias, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "missing bias", fn: func() bool { + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, nil, valid.expertIDs, 4, 32, 1) + return ok + }}, + {name: "wrong selector shape", fn: func() bool { + ids := FromValues([]uint32{0}, 1) + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, ids, 4, 32, 1) + return ok + }}, + {name: "zero topk", fn: func() bool { + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 0) + return ok + }}, + {name: "oversized topk", fn: func() bool { + _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 3) + return ok + }}, + } { + t.Run(tt.name, func(t *testing.T) { + if tt.fn() { + t.Fatal("malformed inputs accepted") + } + }) + } + }) +} + +func TestMoEFusedZeroReferenceAndValidationIsolation(t *testing.T) { + skipIfNoMLX(t) + withMLXThread(t, func() { + if !MetalIsAvailable() { + t.Skip("Metal is not available") + } + + bad := newMoETestInputs([]uint32{2}) + if out, ok := MoEFusedGateUpSwiGLU(bad.input, bad.weight, bad.scales, bad.bias, bad.upWeight, bad.upScales, bad.upBias, bad.expertIDs, 4, 32, 1, -7, 7); ok || out != nil { + t.Fatal("out-of-range gate/up selector accepted") + } + + valid := newMoETestInputs([]uint32{1}) + gateOut, ok := MoEFusedGateUpSwiGLU(valid.input, valid.weight, valid.scales, valid.bias, valid.upWeight, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1, -7, 7) + if !ok { + t.Fatal("valid gate/up call failed after malformed call") + } + Eval(gateOut) + assertMoEZeroReference(t, gateOut, []int{1, 1, 4}) + + downInput := Zeros(DTypeFloat32, 1, 1, 32*32) + if out, ok := MoEFusedDown(downInput, valid.weight, valid.scales, valid.bias, bad.expertIDs, 4, 32, 1); ok || out != nil { + t.Fatal("out-of-range down selector accepted") + } + downOut, ok := MoEFusedDown(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1) + if !ok { + t.Fatal("valid down call failed after malformed call") + } + Eval(downOut) + assertMoEZeroReference(t, downOut, []int{1, 1, 4}) + }) +} + +func assertMoEZeroReference(t *testing.T, got *Array, wantShape []int) { + t.Helper() + if dims := got.Dims(); len(dims) != len(wantShape) { + t.Fatalf("shape = %v, want %v", dims, wantShape) + } else { + for i := range dims { + if dims[i] != wantShape[i] { + t.Fatalf("shape = %v, want %v", dims, wantShape) + } + } + } + for i, value := range got.Floats() { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) || math.Abs(float64(value)) > 1e-6 { + t.Fatalf("output[%d] = %v, want zero reference within 1e-6", i, value) + } + } +} From 135d09f8e7c6dc2beabd8983175f6d6f0626a0b0 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:45:10 +0000 Subject: [PATCH 18/58] tokenizer: preserve ordered GPT-OSS EOS tokens Co-authored-by: Codex --- x/tokenizer/tokenizer.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/x/tokenizer/tokenizer.go b/x/tokenizer/tokenizer.go index b4c8a1c0eb5..d798b700747 100644 --- a/x/tokenizer/tokenizer.go +++ b/x/tokenizer/tokenizer.go @@ -90,6 +90,23 @@ func (t *Tokenizer) EOSTokens() []int32 { return t.vocab.EOS } +// SetEOSTokens replaces the tokenizer's EOS token sequence, preserving the +// declaration order and duplicate non-negative IDs. +func (t *Tokenizer) SetEOSTokens(ids ...int32) { + if t == nil || t.vocab == nil { + return + } + + eos := make([]int32, 0, len(ids)) + for _, id := range ids { + if id >= 0 { + eos = append(eos, id) + } + } + + t.vocab.EOS = eos +} + // PAD returns the padding token ID, or -1 if not set func (t *Tokenizer) PAD() int32 { return t.vocab.PAD From 47f285e39da64bebbeec327a11fab28ccc0e4fb4 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 15:56:20 +0000 Subject: [PATCH 19/58] test: add GPT-OSS forward-reference parity Co-authored-by: Codex --- x/models/gptoss/forward_reference_test.go | 439 ++++++++++++++++++++++ 1 file changed, 439 insertions(+) create mode 100644 x/models/gptoss/forward_reference_test.go diff --git a/x/models/gptoss/forward_reference_test.go b/x/models/gptoss/forward_reference_test.go new file mode 100644 index 00000000000..0b328417c17 --- /dev/null +++ b/x/models/gptoss/forward_reference_test.go @@ -0,0 +1,439 @@ +//go:build gptoss_forward_reference + +package gptoss + +import ( + "encoding/json" + "fmt" + "math" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +func TestForwardReferenceShortEmbeddings(t *testing.T) { + testForwardReferenceEmbeddings(t, "short") +} + +func TestForwardReferenceCacheDecodeEmbeddings(t *testing.T) { + testForwardReferenceEmbeddings(t, "cache") +} + +func testForwardReferenceEmbeddings(t *testing.T, caseName string) { + t.Helper() + modelDir := envDirOrSkip(t, "GPTOSS_MODEL_DIR") + refDir := envDirOrSkip(t, "GPTOSS_REF_DIR") + gptossForwardReferenceSkipIfNoMLX(t) + + ref := loadGPTOSSForwardReference(t, refDir, caseName) + defer ref.Close() + + model := loadGPTOSSEmbeddingReferenceModel(t, modelDir) + defer model.Close() + + inputIDs := flattenReferenceInputIDs(t, ref.Manifest.InputIDs) + tokens := mlx.FromValues(inputIDs, len(ref.Manifest.InputIDs), len(ref.Manifest.InputIDs[0])) + assertReferenceInputIDs(t, ref, inputIDs) + + got := model.Embedding.Forward(tokens) + want := ref.Tensor(t, "model.embed_tokens") + compareForwardReferenceArrays(t, "model.embed_tokens", got, want, 0) +} + +type gptossEmbeddingReferenceModel struct { + Embedding *nn.Embedding + file *mlx.SafetensorsFile +} + +func (m *gptossEmbeddingReferenceModel) Close() { + if m != nil && m.file != nil { + m.file.Free() + } +} + +func loadGPTOSSEmbeddingReferenceModel(t *testing.T, modelDir string) *gptossEmbeddingReferenceModel { + t.Helper() + + configData, err := os.ReadFile(filepath.Join(modelDir, "config.json")) + if err != nil { + t.Skipf("GPTOSS_MODEL_DIR is missing config.json: %v", err) + } + cfg, err := parseConfig(configData) + if err != nil { + t.Fatalf("parse GPTOSS_MODEL_DIR config.json: %v", err) + } + + file, weight := loadHFSafetensor(t, modelDir, "model.embed_tokens.weight") + if err := gptossForwardReferenceValidateTensorShape("model.embed_tokens.weight", weight, []int{int(cfg.VocabSize), int(cfg.HiddenSize)}, "vocab_size x hidden_size"); err != nil { + file.Free() + t.Fatalf("validate embedding tensor: %v", err) + } + + return &gptossEmbeddingReferenceModel{ + Embedding: nn.NewEmbedding(weight), + file: file, + } +} + +func loadHFSafetensor(t *testing.T, modelDir, tensorName string) (*mlx.SafetensorsFile, *mlx.Array) { + t.Helper() + + indexPath := filepath.Join(modelDir, "model.safetensors.index.json") + data, err := os.ReadFile(indexPath) + if err != nil { + t.Skipf("GPTOSS_MODEL_DIR is missing model.safetensors.index.json: %v", err) + } + + var index struct { + WeightMap map[string]string `json:"weight_map"` + } + if err := json.Unmarshal(data, &index); err != nil { + t.Fatalf("parse %s: %v", indexPath, err) + } + + fileName, ok := index.WeightMap[tensorName] + if !ok { + t.Skipf("GPTOSS_MODEL_DIR index does not contain %q", tensorName) + } + + path := filepath.Join(modelDir, fileName) + if _, err := os.Stat(path); err != nil { + t.Skipf("GPTOSS_MODEL_DIR is missing %s for %q: %v", fileName, tensorName, err) + } + + file, err := mlx.LoadSafetensorsNative(path) + if err != nil { + t.Fatalf("load %s: %v", path, err) + } + tensor := file.Get(tensorName) + if tensor == nil { + file.Free() + t.Fatalf("safetensors file %s did not contain %q", path, tensorName) + } + + return file, tensor +} + +type gptossForwardReference struct { + Manifest gptossReferenceManifest + file *mlx.SafetensorsFile +} + +type gptossReferenceManifest struct { + InputIDs []gptossReferenceTokenRow `json:"input_ids"` + Tensors map[string]gptossReferenceTensorSpec `json:"tensors"` +} + +type gptossReferenceTensorSpec struct { + DType string `json:"dtype"` + Shape []int `json:"shape"` +} + +type gptossReferenceTokenRow []int32 + +func (r *gptossForwardReference) Close() { + if r != nil && r.file != nil { + r.file.Free() + } +} + +func (r *gptossForwardReference) Tensor(t *testing.T, name string) *mlx.Array { + t.Helper() + spec, ok := r.Manifest.Tensors[name] + if !ok { + t.Fatalf("reference manifest does not declare tensor %q", name) + } + tensor := r.file.Get(name) + if tensor == nil { + t.Fatalf("reference safetensors does not contain %q", name) + } + if err := gptossForwardReferenceValidateMetadata(name, spec, tensor.DType().String(), tensor.Dims()); err != nil { + t.Fatal(err) + } + return tensor +} + +func gptossForwardReferenceValidateMetadata(name string, spec gptossReferenceTensorSpec, dtype string, shape []int) error { + if spec.DType == "" || !strings.EqualFold(spec.DType, dtype) { + return fmt.Errorf("reference tensor %q dtype = %q, manifest declares %q", name, dtype, spec.DType) + } + if !slices.Equal(spec.Shape, shape) { + return fmt.Errorf("reference tensor %q shape = %v, manifest declares %v", name, shape, spec.Shape) + } + return nil +} + +func loadGPTOSSForwardReference(t *testing.T, refDir, caseName string) *gptossForwardReference { + t.Helper() + + caseDir := filepath.Join(refDir, caseName) + manifestPath := filepath.Join(caseDir, "activations.safetensors.manifest.json") + data, err := os.ReadFile(manifestPath) + if err != nil { + t.Skipf("GPTOSS_REF_DIR is missing %s reference manifest: %v", caseName, err) + } + + var manifest gptossReferenceManifest + if err := json.Unmarshal(data, &manifest); err != nil { + t.Fatalf("parse %s: %v", manifestPath, err) + } + if len(manifest.InputIDs) == 0 || len(manifest.InputIDs[0]) == 0 { + t.Fatalf("%s reference manifest has empty input_ids", caseName) + } + + refPath := filepath.Join(caseDir, "activations.safetensors") + if _, err := os.Stat(refPath); err != nil { + t.Skipf("GPTOSS_REF_DIR is missing %s reference safetensors: %v", caseName, err) + } + file, err := mlx.LoadSafetensorsNative(refPath) + if err != nil { + t.Fatalf("load %s: %v", refPath, err) + } + + return &gptossForwardReference{ + Manifest: manifest, + file: file, + } +} + +func flattenReferenceInputIDs(t *testing.T, rows []gptossReferenceTokenRow) []int32 { + t.Helper() + if len(rows) == 0 { + t.Fatal("reference input_ids has no rows") + } + width := len(rows[0]) + if width == 0 { + t.Fatal("reference input_ids has no columns") + } + + out := make([]int32, 0, len(rows)*width) + for rowIndex, row := range rows { + if len(row) != width { + t.Fatalf("reference input_ids row %d width = %d, want %d", rowIndex, len(row), width) + } + out = append(out, row...) + } + return out +} + +func assertReferenceInputIDs(t *testing.T, ref *gptossForwardReference, want []int32) { + t.Helper() + + gotTensor := ref.Tensor(t, "input_ids") + got := materializedInts(gotTensor) + wantInts := make([]int, len(want)) + for i, token := range want { + wantInts[i] = int(token) + } + if !slices.Equal(got, wantInts) { + t.Fatalf("reference input_ids tensor = %v, want manifest ids %v", got, wantInts) + } +} + +func materializedInts(a *mlx.Array) []int { + if a == nil { + return nil + } + cloned := a.Clone() + mlx.Eval(cloned) + return cloned.Ints() +} + +func compareForwardReferenceArrays(t *testing.T, name string, got, want *mlx.Array, absTol float64) { + t.Helper() + + if got == nil || !got.Valid() { + t.Fatalf("%s got tensor is invalid", name) + } + if want == nil || !want.Valid() { + t.Fatalf("%s reference tensor is invalid", name) + } + if !slices.Equal(got.Dims(), want.Dims()) { + t.Fatalf("%s dims = %v, want %v", name, got.Dims(), want.Dims()) + } + + gotVals := gptossForwardReferenceMaterializedFloats(got.AsType(mlx.DTypeFloat32)) + wantVals := gptossForwardReferenceMaterializedFloats(want.AsType(mlx.DTypeFloat32)) + if len(gotVals) != len(wantVals) { + t.Fatalf("%s length = %d, want %d", name, len(gotVals), len(wantVals)) + } + + maxDiff, maxIndex, err := gptossForwardReferenceCompareValues(gotVals, wantVals, absTol) + if err != nil { + t.Fatalf("%s: %v", name, err) + } + t.Logf("%s matched %d values exactly; max diff %v at flat index %d", name, len(wantVals), maxDiff, maxIndex) +} + +func gptossForwardReferenceCompareValues(got, want []float32, absTol float64) (float64, int, error) { + if len(got) != len(want) { + return 0, -1, fmt.Errorf("length = %d, want %d", len(got), len(want)) + } + var maxDiff float64 + maxIndex := -1 + for i := range want { + if math.IsNaN(float64(got[i])) || math.IsInf(float64(got[i]), 0) || + math.IsNaN(float64(want[i])) || math.IsInf(float64(want[i]), 0) { + return 0, -1, fmt.Errorf("non-finite value at %d: got %v, want %v", i, got[i], want[i]) + } + diff := math.Abs(float64(got[i] - want[i])) + if diff > maxDiff { + maxDiff = diff + maxIndex = i + } + if diff > absTol { + return maxDiff, maxIndex, fmt.Errorf("value[%d] = %v, want %v (diff %v, max tol %v)", i, got[i], want[i], diff, absTol) + } + } + return maxDiff, maxIndex, nil +} + +func TestForwardReferenceMetadataValidation(t *testing.T) { + spec := gptossReferenceTensorSpec{DType: "F32", Shape: []int{2, 3}} + for _, tt := range []struct { + name string + dtype string + shape []int + wantErr bool + }{ + {name: "exact", dtype: "F32", shape: []int{2, 3}}, + {name: "dtype mismatch", dtype: "BF16", shape: []int{2, 3}, wantErr: true}, + {name: "shape mismatch", dtype: "F32", shape: []int{3, 2}, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + err := gptossForwardReferenceValidateMetadata("tensor", spec, tt.dtype, tt.shape) + if (err != nil) != tt.wantErr { + t.Fatalf("error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestForwardReferenceValueComparison(t *testing.T) { + for _, tt := range []struct { + name string + got []float32 + want []float32 + wantErr bool + }{ + {name: "exact", got: []float32{-1, 0, 1}, want: []float32{-1, 0, 1}}, + {name: "finite zero tolerance mismatch", got: []float32{1}, want: []float32{1.0001}, wantErr: true}, + {name: "nan", got: []float32{float32(math.NaN())}, want: []float32{0}, wantErr: true}, + {name: "reference nan", got: []float32{0}, want: []float32{float32(math.NaN())}, wantErr: true}, + {name: "positive infinity", got: []float32{float32(math.Inf(1))}, want: []float32{0}, wantErr: true}, + {name: "negative infinity", got: []float32{float32(math.Inf(-1))}, want: []float32{0}, wantErr: true}, + {name: "matching positive infinities", got: []float32{float32(math.Inf(1))}, want: []float32{float32(math.Inf(1))}, wantErr: true}, + {name: "matching negative infinities", got: []float32{float32(math.Inf(-1))}, want: []float32{float32(math.Inf(-1))}, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + _, _, err := gptossForwardReferenceCompareValues(tt.got, tt.want, 0) + if (err != nil) != tt.wantErr { + t.Fatalf("error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func gptossForwardReferenceSkipIfNoMLX(t *testing.T) { + t.Helper() + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } +} + +func gptossForwardReferenceValidateTensorShape(name string, tensor *mlx.Array, want []int, description string) error { + if tensor == nil || !tensor.Valid() { + return fmt.Errorf("tensor %q is invalid", name) + } + if got := tensor.Dims(); !slices.Equal(got, want) { + return fmt.Errorf("tensor %q shape = %v, want %v (%s)", name, got, want, description) + } + return nil +} + +func gptossForwardReferenceMaterializedFloats(a *mlx.Array) []float32 { + if a == nil { + return nil + } + cloned := a.Clone() + mlx.Eval(cloned) + return cloned.Floats() +} + +func envDirOrSkip(t *testing.T, name string) string { + t.Helper() + value := os.Getenv(name) + if value == "" { + t.Skipf("%s not set; set %s to a local GPT-OSS model/reference directory", name, name) + } + candidates, err := envDirCandidates(value) + if err != nil { + t.Fatalf("resolve %s=%q: %v", name, value, err) + } + var statErr error + for _, candidate := range candidates { + info, err := os.Stat(candidate) + if err != nil { + statErr = err + continue + } + if !info.IsDir() { + t.Skipf("%s=%s resolved to %s, which is not a directory", name, value, candidate) + } + return candidate + } + t.Skipf("%s=%s does not exist; checked %v: %v", name, value, candidates, statErr) + return "" +} + +func envDirCandidates(value string) ([]string, error) { + if filepath.IsAbs(value) { + return []string{value}, nil + } + + cwdAbs, absErr := filepath.Abs(value) + if absErr != nil { + return nil, absErr + } + + root := findGoModRoot() + if root == "" { + return []string{cwdAbs}, nil + } + rootAbs := filepath.Clean(filepath.Join(root, value)) + candidates := []string{cwdAbs} + if rootAbs != cwdAbs { + candidates = append(candidates, rootAbs) + } + + if workdir := os.Getenv("WORKDIR"); workdir != "" { + workdirAbs := filepath.Clean(filepath.Join(workdir, value)) + if !slices.Contains(candidates, workdirAbs) { + candidates = append(candidates, workdirAbs) + } + } + + return candidates, nil +} + +func findGoModRoot() string { + dir, err := os.Getwd() + if err != nil { + return "" + } + for { + if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil { + return dir + } + parent := filepath.Dir(dir) + if parent == dir { + return "" + } + dir = parent + } +} From fc8dd078070f07ac56b72f87c80da01a11adc9cd Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 16:04:33 +0000 Subject: [PATCH 20/58] mlx: carry attention sinks through shared SDPA Co-authored-by: Codex --- x/mlxrunner/mlx/fast.go | 18 +++++++++++++++++- x/models/nn/sdpa.go | 13 +++++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/x/mlxrunner/mlx/fast.go b/x/mlxrunner/mlx/fast.go index 27d5724ede0..af501e4b8d2 100644 --- a/x/mlxrunner/mlx/fast.go +++ b/x/mlxrunner/mlx/fast.go @@ -4,11 +4,27 @@ package mlx import "C" import ( + "fmt" "unsafe" ) -func FastScaledDotProductAttention(q, k, v *Array, scale float32, mode string, mask *Array) *Array { +func FastScaledDotProductAttention(q, k, v *Array, scale float32, mode string, mask *Array, sinkArr ...*Array) *Array { sinks := New("") + if len(sinkArr) > 1 { + panic("mlx.FastScaledDotProductAttention: at most one sinks array is allowed") + } + if len(sinkArr) == 1 && sinkArr[0] != nil { + sinks = sinkArr[0] + if q == nil || !q.Valid() || sinks == nil || !sinks.Valid() { + panic("mlx.FastScaledDotProductAttention: query and sinks must be valid arrays") + } + if q.NumDims() != 4 { + panic(fmt.Sprintf("mlx.FastScaledDotProductAttention: query with sinks must have rank 4, got shape %v", q.Dims())) + } + if sinks.NumDims() != 1 || sinks.Dim(0) != q.Dim(1) || sinks.DType() != q.DType() { + panic(fmt.Sprintf("mlx.FastScaledDotProductAttention: sinks must have shape [heads]=[%d] and dtype %s for query %v, got shape %v dtype %s", q.Dim(1), q.DType(), q.Dims(), sinks.Dims(), sinks.DType())) + } + } cMode := C.CString(mode) defer C.free(unsafe.Pointer(cMode)) diff --git a/x/models/nn/sdpa.go b/x/models/nn/sdpa.go index 3aaa2a908d8..fdce97df104 100644 --- a/x/models/nn/sdpa.go +++ b/x/models/nn/sdpa.go @@ -19,6 +19,10 @@ type sdpaConfig struct { // Optional model-supplied logical mask. mask AttentionMask + + // Optional per-head attention sinks. This is deliberately call-scoped and + // excluded from dispatchInputs because mask resolution does not depend on it. + sinks *mlx.Array } // WithKVHistory supplies a cache's per-layer view of K and V. The @@ -47,6 +51,11 @@ func WithMask(m AttentionMask) SDPAOption { return func(c *sdpaConfig) { c.mask = m } } +// WithSinks supplies per-head attention sinks for the current SDPA call. +func WithSinks(sinks *mlx.Array) SDPAOption { + return func(c *sdpaConfig) { c.sinks = sinks } +} + // ScaledDotProductAttention runs the fast SDPA kernel against q and // the keys/values supplied via exactly one of WithKV or // WithKVHistory. Automatically applies any Q/K padding masking required @@ -85,12 +94,12 @@ func ScaledDotProductAttention(b *batch.Batch, q *mlx.Array, scale float32, opts if cached, ok := b.Memo.Get(inputs); ok { d := cached.(sdpaDispatch) - return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr) + return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr, cfg.sinks) } d := inputs.resolve() b.Memo.Put(inputs, d) - return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr) + return mlx.FastScaledDotProductAttention(q, k, v, scale, d.mode, d.arr, cfg.sinks) } // sdpaDispatch is the resolved kernel call for a given SDPA key — From 6baa1f25fd450691ce603cc3d22c619916f0a4a3 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 17:55:37 +0000 Subject: [PATCH 21/58] gptoss: execute dense, expert, and cached MLX paths Co-authored-by: Codex --- x/models/gptoss/gptoss.go | 1266 +++++++++++++++++++++++++- x/models/gptoss/gptoss_row10_test.go | 368 ++++++++ 2 files changed, 1607 insertions(+), 27 deletions(-) create mode 100644 x/models/gptoss/gptoss_row10_test.go diff --git a/x/models/gptoss/gptoss.go b/x/models/gptoss/gptoss.go index 644d69f9495..ad5774151ab 100644 --- a/x/models/gptoss/gptoss.go +++ b/x/models/gptoss/gptoss.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "math" + "slices" "strings" "github.com/ollama/ollama/x/mlxrunner/batch" @@ -12,6 +13,7 @@ import ( "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" ) @@ -19,7 +21,9 @@ func init() { base.Register("GptOssForCausalLM", NewModel) } -// RopeScaling carries the GPT-OSS rope scaling block. +var harmonyEOSTokens = []string{"<|endoftext|>", "<|return|>", "<|call|>"} + +// RopeScaling carries the gpt-oss rope scaling block. type RopeScaling struct { Factor float32 `json:"factor"` OriginalMaxPositionEmbeddings int32 `json:"original_max_position_embeddings"` @@ -65,14 +69,215 @@ type Config struct { QuantMode string `json:"-"` TensorQuant map[string]*model.TensorQuantInfo `json:"-"` QuantMethod string `json:"-"` + checked gptossCheckedDimensions `json:"-"` +} + +const ( + gptossMaxLoopEntries = 1 << 20 + gptossMaxDimension = 1 << 24 +) + +type gptossCheckedDimensions struct { + layers, experts, expertsPerToken int + hidden, intermediate, headDim int + attentionHeads, kvHeads int + vocab, slidingWindow, context int + q, kv, gateUp int +} + +func checkedGPTOSSProduct(limit uint64, values ...int32) (int, error) { + product := uint64(1) + for _, value := range values { + if value <= 0 || product > ^uint64(0)/uint64(value) { + return 0, fmt.Errorf("invalid dimension product %v", values) + } + product *= uint64(value) + } + if product > limit || product > uint64(^uint(0)>>1) { + return 0, fmt.Errorf("dimension product %v = %d exceeds limit %d", values, product, limit) + } + return int(product), nil +} + +func validateGPTOSSConfigDimensions(cfg *Config) (gptossCheckedDimensions, error) { + if cfg == nil { + return gptossCheckedDimensions{}, fmt.Errorf("missing config") + } + loopScalars := []struct { + name string + value int32 + }{ + {"num_hidden_layers", cfg.NumHiddenLayers}, + {"num_local_experts", cfg.NumLocalExperts}, + {"num_experts_per_tok", cfg.NumExpertsPerTok}, + {"num_attention_heads", cfg.NumAttentionHeads}, + {"num_key_value_heads", cfg.NumKeyValueHeads}, + } + for _, scalar := range loopScalars { + if scalar.value <= 0 || scalar.value > gptossMaxLoopEntries { + return gptossCheckedDimensions{}, fmt.Errorf("%s = %d exceeds supported range 1..%d", scalar.name, scalar.value, gptossMaxLoopEntries) + } + } + dimensionScalars := []struct { + name string + value int32 + }{ + {"hidden_size", cfg.HiddenSize}, + {"intermediate_size", cfg.IntermediateSize}, + {"head_dim", cfg.HeadDim}, + {"vocab_size", cfg.VocabSize}, + {"sliding_window", cfg.SlidingWindow}, + {"max_position_embeddings", cfg.MaxPositionEmbeddings}, + {"rope_scaling.original_max_position_embeddings", cfg.RopeScaling.OriginalMaxPositionEmbeddings}, + } + for _, scalar := range dimensionScalars { + if scalar.value <= 0 || scalar.value > gptossMaxDimension { + return gptossCheckedDimensions{}, fmt.Errorf("%s = %d exceeds supported range 1..%d", scalar.name, scalar.value, gptossMaxDimension) + } + } + if cfg.RopeTheta <= 0 || cfg.RMSNormEps <= 0 || cfg.RopeScaling.Factor <= 0 { + return gptossCheckedDimensions{}, fmt.Errorf("rope_theta, rms_norm_eps, and rope_scaling.factor must be positive") + } + if cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return gptossCheckedDimensions{}, fmt.Errorf("num_attention_heads must be divisible by num_key_value_heads") + } + if cfg.NumExpertsPerTok > cfg.NumLocalExperts { + return gptossCheckedDimensions{}, fmt.Errorf("num_experts_per_tok must not exceed num_local_experts") + } + q, err := checkedGPTOSSProduct(gptossMaxDimension, cfg.NumAttentionHeads, cfg.HeadDim) + if err != nil { + return gptossCheckedDimensions{}, fmt.Errorf("query dimension: %w", err) + } + kv, err := checkedGPTOSSProduct(gptossMaxDimension, cfg.NumKeyValueHeads, cfg.HeadDim) + if err != nil { + return gptossCheckedDimensions{}, fmt.Errorf("key/value dimension: %w", err) + } + gateUp, err := checkedGPTOSSProduct(gptossMaxDimension, 2, cfg.IntermediateSize) + if err != nil { + return gptossCheckedDimensions{}, fmt.Errorf("gate/up dimension: %w", err) + } + if _, err := checkedGPTOSSProduct(gptossMaxLoopEntries, cfg.NumHiddenLayers, cfg.NumLocalExperts); err != nil { + return gptossCheckedDimensions{}, fmt.Errorf("layer/expert work: %w", err) + } + return gptossCheckedDimensions{ + layers: int(cfg.NumHiddenLayers), experts: int(cfg.NumLocalExperts), expertsPerToken: int(cfg.NumExpertsPerTok), + hidden: int(cfg.HiddenSize), intermediate: int(cfg.IntermediateSize), headDim: int(cfg.HeadDim), + attentionHeads: int(cfg.NumAttentionHeads), kvHeads: int(cfg.NumKeyValueHeads), vocab: int(cfg.VocabSize), + slidingWindow: int(cfg.SlidingWindow), context: int(cfg.MaxPositionEmbeddings), q: q, kv: kv, gateUp: gateUp, + }, nil +} + +// RopeParameters returns the runtime rope settings derived from config. +func (c *Config) RopeParameters() (base, scale float32, originalContext int) { + if c == nil { + return 0, 1, 0 + } + base = c.RopeTheta + scale = 1 + if c.RopeScaling.Factor > 0 { + scale = 1 / c.RopeScaling.Factor + } + if c.RopeScaling.OriginalMaxPositionEmbeddings > 0 { + originalContext = int(c.RopeScaling.OriginalMaxPositionEmbeddings) + } + return base, scale, originalContext } // Model is the gpt-oss text-only model. type Model struct { + EmbedTokens nn.EmbeddingLayer + Layers []*Layer + Norm *nn.RMSNorm + LMHead nn.LinearLayer + tok *tokenizer.Tokenizer *Config } +// Layer is a single gpt-oss decoder block. +type Layer struct { + AttentionNorm *nn.RMSNorm + Attention *Attention + FFNNorm *nn.RMSNorm + Router nn.LinearLayer + Experts *Experts +} + +// Attention implements the split gpt-oss attention path. +type Attention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + Sinks *mlx.Array + RoPEFreqs *mlx.Array + RoPEScale float32 +} + +// Experts holds the loaded gpt-oss MoE expert projections. +type Experts struct { + GateUp *ExpertPair + Down *ExpertProjection +} + +// ExpertPair stores the split gate and up projections from the packed expert tensor. +type ExpertPair struct { + Gate *ExpertProjection + Up *ExpertProjection +} + +// ExpertProjection wraps a per-expert weight matrix. +// When Scales is non-nil, the weight is quantized and Forward uses GatherQMM. +type ExpertProjection struct { + Weight *mlx.Array + Bias *mlx.Array + Scales *mlx.Array + QBiases *mlx.Array + GroupSize int + Bits int + Mode string + Transpose bool // true for MLX-native [experts, out, packed_in] layout +} + +type fusedExpertDescriptor struct { + present, hasQBias, transpose bool + mode string + bits, groupSize int + weightType, scaleType, biasType mlx.DType + weightDims, scaleDims, biasDims []int +} + +func fusedExpertDescriptorFor(p *ExpertProjection) fusedExpertDescriptor { + if p == nil || p.Weight == nil || !p.Weight.Valid() || p.Scales == nil || !p.Scales.Valid() || p.Bias == nil || !p.Bias.Valid() { + return fusedExpertDescriptor{} + } + return fusedExpertDescriptor{ + present: true, hasQBias: p.QBiases != nil, transpose: p.Transpose, mode: p.Mode, bits: p.Bits, groupSize: p.GroupSize, + weightType: p.Weight.DType(), scaleType: p.Scales.DType(), biasType: p.Bias.DType(), + weightDims: p.Weight.Dims(), scaleDims: p.Scales.Dims(), biasDims: p.Bias.Dims(), + } +} + +func fusedExpertDescriptorEligible(d fusedExpertDescriptor) bool { + return d.present && !d.hasQBias && d.mode == "mxfp4" && d.bits == 4 && d.groupSize == 32 && d.transpose && + d.weightType == mlx.DTypeUint32 && d.scaleType == mlx.DTypeUint8 && d.biasType == mlx.DTypeBFloat16 && + len(d.weightDims) == 3 && len(d.scaleDims) == 3 && len(d.biasDims) == 2 +} + +func fusedExpertProjectionEligible(p *ExpertProjection) bool { + return fusedExpertDescriptorEligible(fusedExpertDescriptorFor(p)) +} + +func fusedExpertPairEligible(gate, up *ExpertProjection) bool { + g, u := fusedExpertDescriptorFor(gate), fusedExpertDescriptorFor(up) + return fusedExpertDescriptorsPairEligible(g, u) +} + +func fusedExpertDescriptorsPairEligible(g, u fusedExpertDescriptor) bool { + return fusedExpertDescriptorEligible(g) && fusedExpertDescriptorEligible(u) && + slices.Equal(g.weightDims, u.weightDims) && slices.Equal(g.scaleDims, u.scaleDims) && slices.Equal(g.biasDims, u.biasDims) +} + // NewModel creates a gpt-oss model from a manifest root. func NewModel(root *model.Root) (base.Model, error) { configData, err := root.Manifest.ReadConfig("config.json") @@ -127,13 +332,31 @@ func NewModel(root *model.Root) (base.Model, error) { if err != nil { return nil, fmt.Errorf("parse tokenizer: %w", err) } + ensureHarmonyEOSTokens(tok) return &Model{ Config: &cfg, + Layers: make([]*Layer, cfg.checked.layers), tok: tok, }, nil } +func ensureHarmonyEOSTokens(tok *tokenizer.Tokenizer) { + if tok == nil { + return + } + + ids := append([]int32(nil), tok.EOSTokens()...) + for _, name := range harmonyEOSTokens { + if id, ok := tok.GetSpecialToken(name); ok { + ids = append(ids, id) + } + } + if len(ids) > 0 { + tok.SetEOSTokens(ids...) + } +} + func parseConfig(configData []byte) (Config, error) { var raw map[string]json.RawMessage if err := json.Unmarshal(configData, &raw); err != nil { @@ -189,6 +412,9 @@ func parseConfig(configData []byte) (Config, error) { if cfg.NumExpertsPerTok <= 0 { return Config{}, fmt.Errorf("invalid num_experts_per_tok: %d", cfg.NumExpertsPerTok) } + if cfg.NumExpertsPerTok > cfg.NumLocalExperts { + return Config{}, fmt.Errorf("num_experts_per_tok (%d) must be <= num_local_experts (%d)", cfg.NumExpertsPerTok, cfg.NumLocalExperts) + } if cfg.SlidingWindow <= 0 { return Config{}, fmt.Errorf("invalid sliding_window: %d", cfg.SlidingWindow) } @@ -222,62 +448,1048 @@ func parseConfig(configData []byte) (Config, error) { } else if cfg.Quantization.QuantMethod != "" { cfg.QuantMethod = strings.ToLower(cfg.Quantization.QuantMethod) } + checked, err := validateGPTOSSConfigDimensions(&cfg) + if err != nil { + return Config{}, err + } + cfg.checked = checked return cfg, nil } -// Forward is intentionally skeletal for Phase 2. -func (m *Model) Forward(_ *batch.Batch, _ []cache.Cache) (*mlx.Array, *mlx.Array) { - return nil, nil +// NumLayers returns the configured layer count. +func (m *Model) NumLayers() int { + if m == nil || m.Config == nil { + return 0 + } + if m.checked.layers > 0 { + return m.checked.layers + } + return int(m.NumHiddenLayers) } -// Unembed is intentionally skeletal for Phase 2. -func (m *Model) Unembed(_ *mlx.Array) *mlx.Array { - return nil +// Tokenizer returns the loaded tokenizer. +func (m *Model) Tokenizer() *tokenizer.Tokenizer { + if m == nil { + return nil + } + return m.tok } -// NumLayers returns the configured layer count. -func (m *Model) NumLayers() int { +// MaxContextLength returns the derived context length. +func (m *Model) MaxContextLength() int { if m == nil || m.Config == nil { return 0 } - return int(m.NumHiddenLayers) + if m.MaxPositionEmbeddings > 0 { + return int(m.MaxPositionEmbeddings) + } + return 0 } -// NewCaches declares the alternating sliding-window and full-attention cache -// slots owned by GPT-OSS. +// NewCaches returns one cache per layer, matching the classic gpt-oss +// alternating sliding-window / causal parity. func (m *Model) NewCaches() []cache.Cache { caches := make([]cache.Cache, m.NumLayers()) for i := range caches { if i%2 == 0 { caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow)) - } else { - caches[i] = cache.NewKVCache() + continue } + caches[i] = cache.NewKVCache() } return caches } -// Tokenizer returns the loaded tokenizer. -func (m *Model) Tokenizer() *tokenizer.Tokenizer { - if m == nil { +// swiGLUAlphaLimit is a compiled kernel implementing the gpt-oss MoE activation: +// +// swish = min(gate, 7) * sigmoid(1.702 * min(gate, 7)) +// result = swish * (clamp(up, -7, 7) + 1) +// +// The Compile2 wrapper traces the elementwise chain once and fuses it into a +// single Metal kernel, eliminating ~10 per-op dispatches per MoE block. +var swiGLUAlphaLimit = mlx.Compile2( + "gptoss_swiglu_alpha_limit", + func(gate, up *mlx.Array) *mlx.Array { + dt := gate.DType() + alpha := mlx.FromValue[float32](1.702).AsType(dt) + limit := mlx.FromValue[float32](7).AsType(dt) + negLimit := mlx.Neg(limit) + one := mlx.FromValue[float32](1).AsType(dt) + + clippedGate := mlx.Minimum(gate, limit) + clippedUp := mlx.Clip(up, negLimit, limit) + + swish := clippedGate.Multiply(mlx.Mul(clippedGate, alpha).Sigmoid()) + return swish.Multiply(clippedUp.Add(one)) + }, + mlx.Shapeless(), +) + +func expertSlice(t *mlx.Array, expert int32) *mlx.Array { + if t == nil || !t.Valid() { return nil } - return m.tok + + dims := t.Dims() + if len(dims) == 0 { + return nil + } + + start := make([]int32, len(dims)) + stop := make([]int32, len(dims)) + start[0] = expert + stop[0] = expert + 1 + for i := 1; i < len(dims); i++ { + stop[i] = int32(dims[i]) + } + + return mlx.Squeeze(mlx.SliceStartStop(t, start, stop), 0) } -// MaxContextLength returns the derived context length. -func (m *Model) MaxContextLength() int { +func interleavedIndices(count int, offset int32) *mlx.Array { + indices := make([]int32, count) + for i := range count { + indices[i] = int32(i*2) + offset + } + return mlx.FromValues(indices, count) +} + +func splitGateUpInterleaved(dense, bias *mlx.Array, mid int) (gateWeight, upWeight, gateBias, upBias *mlx.Array) { + if dense == nil || !dense.Valid() || bias == nil || !bias.Valid() || mid <= 0 { + return nil, nil, nil, nil + } + + even := interleavedIndices(mid, 0) + odd := interleavedIndices(mid, 1) + gateWeight = mlx.Take(dense, even, 0) + upWeight = mlx.Take(dense, odd, 0) + gateBias = mlx.Take(bias, even, 0) + upBias = mlx.Take(bias, odd, 0) + return gateWeight, upWeight, gateBias, upBias +} + +func buildGPTOSSRoPEFreqs(cfg *Config) (*mlx.Array, float32) { + if cfg == nil || cfg.HeadDim <= 0 || cfg.RopeTheta <= 0 { + return nil, 1 + } + params := &nn.RopeParameters{ + RopeTheta: cfg.RopeTheta, + RopeType: cfg.RopeScaling.RopeType, + Factor: cfg.RopeScaling.Factor, + OriginalMaxPositionEmbeddings: cfg.RopeScaling.OriginalMaxPositionEmbeddings, + BetaFast: cfg.RopeScaling.BetaFast, + BetaSlow: cfg.RopeScaling.BetaSlow, + } + return nn.BuildYarnRopeFreqs(int(cfg.HeadDim), cfg.RopeTheta, params) +} + +func requireTensor(tensors map[string]*mlx.Array, name string) (*mlx.Array, error) { + t := tensors[name] + if t == nil || !t.Valid() { + return nil, fmt.Errorf("missing tensor %q", name) + } + return t, nil +} + +func validateTensorShape(name string, t *mlx.Array, want []int, wantExpr string) error { + if t == nil || !t.Valid() { + return fmt.Errorf("missing tensor %q", name) + } + + got := t.Dims() + if len(got) != len(want) { + return fmt.Errorf("tensor %q shape %v, want %v (%s)", name, got, want, wantExpr) + } + for i := range want { + if got[i] != want[i] { + return fmt.Errorf("tensor %q shape %v, want %v (%s)", name, got, want, wantExpr) + } + } + return nil +} + +func validateLayerTensorShape(layer int, name string, t *mlx.Array, want []int, wantExpr string) error { + if err := validateTensorShape(name, t, want, wantExpr); err != nil { + return fmt.Errorf("layer %d: %w", layer, err) + } + return nil +} + +func validateLayerTensorDType(layer int, name string, t *mlx.Array, want mlx.DType, wantExpr string) error { + if t == nil || !t.Valid() { + return fmt.Errorf("layer %d: missing tensor %q", layer, name) + } + if got := t.DType(); got != want { + return fmt.Errorf("layer %d: tensor %q dtype %s, want %s (%s)", layer, name, got, want, wantExpr) + } + return nil +} + +func prepareGPTOSSAttentionSinks(layer int, name string, source *mlx.Array, heads int) (*mlx.Array, error) { + if err := validateLayerTensorShape(layer, name, source, []int{heads}, "num_attention_heads"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, name, source, mlx.DTypeBFloat16, "GPT-OSS checkpoint attention sinks"); err != nil { + return nil, err + } + prepared := mlx.Contiguous(source, false) + if prepared == nil || !prepared.Valid() { + return nil, fmt.Errorf("layer %d: failed to prepare tensor %q as contiguous BF16", layer, name) + } + mlx.Eval(prepared) + if err := validateLayerTensorShape(layer, name+" (prepared)", prepared, []int{heads}, "prepared num_attention_heads"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, name+" (prepared)", prepared, mlx.DTypeBFloat16, "prepared GPT-OSS attention sinks"); err != nil { + return nil, err + } + finiteProbe := prepared.AsType(mlx.DTypeFloat32) + mlx.Eval(finiteProbe) + for i, value := range finiteProbe.Floats() { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return nil, fmt.Errorf("layer %d: tensor %q contains non-finite value at %d", layer, name, i) + } + } + return prepared, nil +} + +func validateGPTOSSAttentionSinkDispatch(query, sinks *mlx.Array, heads int) error { + if query == nil || !query.Valid() { + return fmt.Errorf("query must be a valid rank-4 array") + } + if query.NumDims() != 4 { + return fmt.Errorf("query must be rank-4, got shape %v dtype %s", query.Dims(), query.DType()) + } + if query.Dim(1) != heads { + return fmt.Errorf("query head axis = %d, want configured heads = %d (shape %v)", query.Dim(1), heads, query.Dims()) + } + if sinks == nil || !sinks.Valid() { + return fmt.Errorf("prepared sinks must be a valid rank-1 array with %d heads", heads) + } + if sinks.NumDims() != 1 || sinks.Dim(0) != heads { + return fmt.Errorf("prepared sinks must be rank-1 with %d heads, got shape %v dtype %s", heads, sinks.Dims(), sinks.DType()) + } + if sinks.DType() != query.DType() { + return fmt.Errorf("prepared sink dtype %s must equal actual query dtype %s", sinks.DType(), query.DType()) + } + return nil +} + +func validateLinearLayerShape( + layer int, + tensors map[string]*mlx.Array, + path string, + wantOut, wantIn int, + wantExpr string, + cfg *Config, +) error { + weightName := path + ".weight" + weight, err := requireTensor(tensors, weightName) + if err != nil { + return fmt.Errorf("layer %d: %w", layer, err) + } + + scales := tensors[weightName+"_scale"] + if scales == nil { + return validateLayerTensorShape(layer, weightName, weight, []int{wantOut, wantIn}, wantExpr) + } + + if len(weight.Dims()) != 2 { + return fmt.Errorf("layer %d: tensor %q dims %v, want quantized matrix for %s", layer, weightName, weight.Dims(), wantExpr) + } + if weight.Dim(0) != wantOut { + return fmt.Errorf("layer %d: tensor %q output dim %d, want %d (%s)", layer, weightName, weight.Dim(0), wantOut, wantExpr) + } + + _, bits, mode := model.ResolveLinearQuantParams( + cfg.QuantGroupSize, + cfg.QuantBits, + cfg.QuantMode, + cfg.TensorQuant, + weightName, + weight, + scales, + ) + if mode == "affine" { + if _, inferredBits, ok := model.InferAffineQuantParamsFromShapes(weight, scales, bits); !ok || inferredBits != bits { + return fmt.Errorf("layer %d: tensor %q has unsupported affine quantized shapes %v / %v for %s", layer, weightName, weight.Dims(), scales.Dims(), wantExpr) + } + } + + return nil +} + +func loadLinearLayer(tensors map[string]*mlx.Array, linears model.LinearFactory, cfg *Config, layer int, path string, wantOut, wantIn int, wantExpr string) (nn.LinearLayer, error) { + if err := validateLinearLayerShape( + layer, + tensors, + path, + wantOut, + wantIn, + wantExpr, + cfg, + ); err != nil { + return nil, err + } + + biasName := path + ".bias" + bias, err := requireTensor(tensors, biasName) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + if err := validateLayerTensorShape(layer, biasName, bias, []int{wantOut}, fmt.Sprintf("%s bias length", path)); err != nil { + return nil, err + } + + layerLinear := linears.Make(path) + if layerLinear == nil { + return nil, fmt.Errorf("layer %d: failed to construct linear layer from %q", layer, path+".weight") + } + if got := layerLinear.OutputDim(); int(got) != wantOut { + return nil, fmt.Errorf("layer %d: linear %q output dim = %d, want %d (%s)", layer, path, got, wantOut, wantExpr) + } + return layerLinear, nil +} + +func loadLayerNormTensor(tensors map[string]*mlx.Array, layer int, name string, want int, wantExpr string) (*nn.RMSNorm, error) { + weight, err := requireTensor(tensors, name) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + if err := validateLayerTensorShape(layer, name, weight, []int{want}, wantExpr); err != nil { + return nil, err + } + return nn.NewRMSNorm(weight, 0), nil +} + +func loadExpertPair(tensors map[string]*mlx.Array, layer int, prefix string, wantOut, wantIn int, cfg *Config) (*ExpertPair, error) { + pair, err := loadDirectExpertPair(tensors, layer, prefix, wantOut/2, wantIn, cfg) + if err != nil { + return nil, err + } + if pair != nil { + return pair, nil + } + + weightName := prefix + ".weight" + biasName := prefix + ".bias" + + weight, err := requireTensor(tensors, weightName) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + bias, err := requireTensor(tensors, biasName) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + + if err := validateLayerTensorShape(layer, weightName, weight, []int{int(cfg.NumLocalExperts), wantOut, wantIn}, fmt.Sprintf("num_local_experts x %s x hidden", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorShape(layer, biasName, bias, []int{int(cfg.NumLocalExperts), wantOut}, fmt.Sprintf("num_local_experts x %s bias", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, weightName, weight, mlx.DTypeBFloat16, "offline-dequantized expert weights"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, biasName, bias, mlx.DTypeBFloat16, "offline-dequantized expert bias"); err != nil { + return nil, err + } + + if wantOut%2 != 0 { + return nil, fmt.Errorf("layer %d: %s output dim must be even, got %d", layer, prefix, wantOut) + } + + mid := wantOut / 2 + gateWeights := make([]*mlx.Array, 0, cfg.NumLocalExperts) + upWeights := make([]*mlx.Array, 0, cfg.NumLocalExperts) + gateBiases := make([]*mlx.Array, 0, cfg.NumLocalExperts) + upBiases := make([]*mlx.Array, 0, cfg.NumLocalExperts) + for e := range cfg.NumLocalExperts { + expertWeight := expertSlice(weight, e) + expertBias := expertSlice(bias, e) + + expertWeightName := fmt.Sprintf("%s.expert[%d]", weightName, e) + if err := validateLayerTensorShape(layer, expertWeightName, expertWeight, []int{wantOut, wantIn}, fmt.Sprintf("%s expert slice", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorShape(layer, expertWeightName+".bias", expertBias, []int{wantOut}, fmt.Sprintf("%s expert bias", prefix)); err != nil { + return nil, err + } + + gateWeight, upWeight, gateBias, upBias := splitGateUpInterleaved(expertWeight, expertBias, mid) + if gateWeight == nil || upWeight == nil || gateBias == nil || upBias == nil { + return nil, fmt.Errorf("layer %d: failed to split interleaved gate/up expert tensor %q", layer, expertWeightName) + } + gateWeight = mlx.Transpose(gateWeight, 1, 0) + upWeight = mlx.Transpose(upWeight, 1, 0) + gateWeight = mlx.Contiguous(gateWeight, false) + upWeight = mlx.Contiguous(upWeight, false) + gateBias = mlx.Contiguous(gateBias, false) + upBias = mlx.Contiguous(upBias, false) + + gateWeights = append(gateWeights, gateWeight) + upWeights = append(upWeights, upWeight) + gateBiases = append(gateBiases, gateBias) + upBiases = append(upBiases, upBias) + } + + gateWeight := mlx.Stack(gateWeights, 0) + upWeight := mlx.Stack(upWeights, 0) + gateBias := mlx.Stack(gateBiases, 0) + upBias := mlx.Stack(upBiases, 0) + mlx.Eval(gateWeight, upWeight, gateBias, upBias) + + return &ExpertPair{ + Gate: &ExpertProjection{ + Weight: gateWeight, + Bias: gateBias, + }, + Up: &ExpertProjection{ + Weight: upWeight, + Bias: upBias, + }, + }, nil +} + +func loadExpertProjection(tensors map[string]*mlx.Array, layer int, prefix string, wantOut, wantIn int, cfg *Config) (*ExpertProjection, error) { + proj, err := loadDirectExpertProjection(tensors, layer, prefix, wantOut, wantIn, cfg) + if err != nil { + return nil, err + } + if proj != nil { + return proj, nil + } + + weightName := prefix + ".weight" + biasName := prefix + ".bias" + + weight, err := requireTensor(tensors, weightName) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + bias, err := requireTensor(tensors, biasName) + if err != nil { + return nil, fmt.Errorf("layer %d: %w", layer, err) + } + + if err := validateLayerTensorShape(layer, weightName, weight, []int{int(cfg.NumLocalExperts), wantOut, wantIn}, fmt.Sprintf("num_local_experts x %s x hidden", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorShape(layer, biasName, bias, []int{int(cfg.NumLocalExperts), wantOut}, fmt.Sprintf("num_local_experts x %s bias", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, weightName, weight, mlx.DTypeBFloat16, "offline-dequantized expert weights"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, biasName, bias, mlx.DTypeBFloat16, "offline-dequantized expert bias"); err != nil { + return nil, err + } + + weights := make([]*mlx.Array, 0, cfg.NumLocalExperts) + biases := make([]*mlx.Array, 0, cfg.NumLocalExperts) + for e := range cfg.NumLocalExperts { + expertWeight := expertSlice(weight, e) + expertBias := expertSlice(bias, e) + + expertWeightName := fmt.Sprintf("%s.expert[%d]", weightName, e) + if err := validateLayerTensorShape(layer, expertWeightName, expertWeight, []int{wantOut, wantIn}, fmt.Sprintf("%s expert slice", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorShape(layer, expertWeightName+".bias", expertBias, []int{wantOut}, fmt.Sprintf("%s expert bias", prefix)); err != nil { + return nil, err + } + + expertWeight = mlx.Transpose(expertWeight, 1, 0) + expertWeight = mlx.Contiguous(expertWeight, false) + weights = append(weights, expertWeight) + biases = append(biases, expertBias) + } + + weightStack := mlx.Stack(weights, 0) + biasStack := mlx.Stack(biases, 0) + mlx.Eval(weightStack, biasStack) + + return &ExpertProjection{ + Weight: weightStack, + Bias: biasStack, + }, nil +} + +func loadDirectExpertPair(tensors map[string]*mlx.Array, layer int, legacyPrefix string, wantOut, wantIn int, cfg *Config) (*ExpertPair, error) { + gatePrefix := strings.Replace(legacyPrefix, "gate_up_proj", "gate_proj", 1) + upPrefix := strings.Replace(legacyPrefix, "gate_up_proj", "up_proj", 1) + gate, err := loadDirectExpertProjection(tensors, layer, gatePrefix, wantOut, wantIn, cfg) + if err != nil { + return nil, err + } + up, err := loadDirectExpertProjection(tensors, layer, upPrefix, wantOut, wantIn, cfg) + if err != nil { + return nil, err + } + switch { + case gate == nil && up == nil: + return nil, nil + case gate == nil: + return nil, fmt.Errorf("layer %d: missing direct gate expert tensors for %q", layer, legacyPrefix) + case up == nil: + return nil, fmt.Errorf("layer %d: missing direct up expert tensors for %q", layer, legacyPrefix) + } + return &ExpertPair{Gate: gate, Up: up}, nil +} + +func loadDirectExpertProjection(tensors map[string]*mlx.Array, layer int, prefix string, wantOut, wantIn int, cfg *Config) (*ExpertProjection, error) { + weightName := prefix + ".weight" + biasName := prefix + ".bias" + weight := tensors[weightName] + bias := tensors[biasName] + scales := tensors[weightName+"_scale"] + qbiases := tensors[weightName+"_qbias"] + if weight == nil && bias == nil && scales == nil && qbiases == nil { + return nil, nil + } + if weight == nil || bias == nil { + return nil, fmt.Errorf("layer %d: incomplete direct expert tensors for %q", layer, prefix) + } + + if err := validateLayerTensorShape(layer, biasName, bias, []int{int(cfg.NumLocalExperts), wantOut}, fmt.Sprintf("num_local_experts x out bias for %s", prefix)); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, biasName, bias, mlx.DTypeBFloat16, "runtime-ready offline expert bias"); err != nil { + return nil, err + } + + if scales != nil { + groupSize, bits, mode := model.ResolveLinearQuantParams( + cfg.QuantGroupSize, + cfg.QuantBits, + cfg.QuantMode, + cfg.TensorQuant, + weightName, + weight, + scales, + ) + if mode != "mxfp4" || bits != 4 || groupSize != 32 { + return nil, fmt.Errorf("layer %d: tensor %q quantization = %s/%d/group%d, want mxfp4/4/group32", layer, weightName, mode, bits, groupSize) + } + if wantIn <= 0 || wantIn%32 != 0 { + return nil, fmt.Errorf("layer %d: tensor %q input dim %d must be positive and divisible by 32", layer, weightName, wantIn) + } + if qbiases != nil { + return nil, fmt.Errorf("layer %d: native MXFP4 tensor %q must not have quantization bias", layer, weightName) + } + if err := validateLayerTensorShape(layer, weightName, weight, []int{int(cfg.NumLocalExperts), wantOut, wantIn / 8}, "experts x out x packed-input words"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, weightName, weight, mlx.DTypeUint32, "native MXFP4 packed weights"); err != nil { + return nil, err + } + if err := validateLayerTensorShape(layer, weightName+"_scale", scales, []int{int(cfg.NumLocalExperts), wantOut, wantIn / 32}, "experts x out x MXFP4 groups"); err != nil { + return nil, err + } + if err := validateLayerTensorDType(layer, weightName+"_scale", scales, mlx.DTypeUint8, "native MXFP4 scales"); err != nil { + return nil, err + } + + return &ExpertProjection{ + Weight: weight, + Bias: bias, + Scales: scales, + QBiases: qbiases, + GroupSize: groupSize, + Bits: bits, + Mode: mode, + Transpose: true, // MLX-native quantized layout: [experts, out, packed_in] + }, nil + } + if qbiases != nil { + return nil, fmt.Errorf("layer %d: orphan quantization bias for %q", layer, weightName) + } + + if err := validateLayerTensorDType(layer, weightName, weight, mlx.DTypeBFloat16, "runtime-ready offline expert weights"); err != nil { + return nil, err + } + if len(weight.Dims()) != 3 || weight.Dim(0) != int(cfg.NumLocalExperts) { + return nil, fmt.Errorf("layer %d: tensor %q dims %v, want direct expert layout [num_local_experts, out, in]", layer, weightName, weight.Dims()) + } + + switch { + case weight.Dim(1) == wantOut && weight.Dim(2) == wantIn: + weight = mlx.Transpose(weight, 0, 2, 1) + weight = mlx.Contiguous(weight, false) + mlx.Eval(weight) + return &ExpertProjection{Weight: weight, Bias: bias}, nil + case weight.Dim(1) == wantIn && weight.Dim(2) == wantOut: + // Legacy GPT-OSS direct BF16 layout kept weights pre-transposed for GatherMM. + return &ExpertProjection{Weight: weight, Bias: bias}, nil + default: + return nil, fmt.Errorf("layer %d: tensor %q dims %v, want direct expert layout [%d %d %d] or legacy [%d %d %d]", + layer, weightName, weight.Dims(), + int(cfg.NumLocalExperts), wantOut, wantIn, + int(cfg.NumLocalExperts), wantIn, wantOut, + ) + } +} + +// LoadWeights assigns dense tensors and structural placeholders to the model. +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { if m == nil || m.Config == nil { - return 0 + return fmt.Errorf("missing gpt-oss config") } - if m.MaxPositionEmbeddings > 0 { - return int(m.MaxPositionEmbeddings) + checked, err := validateGPTOSSConfigDimensions(m.Config) + if err != nil { + return err } - return 0 + m.checked = checked + if len(m.Layers) == 0 { + m.Layers = make([]*Layer, m.NumLayers()) + } + + linears := model.NewLinearFactory(tensors, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + + embeddingWeight, err := requireTensor(tensors, "embedding.weight") + if err != nil { + return err + } + if tensors["embedding.weight_scale"] == nil { + if err := validateTensorShape("embedding.weight", embeddingWeight, []int{int(m.VocabSize), int(m.HiddenSize)}, "vocab_size x hidden_size"); err != nil { + return err + } + } + embedTokens := model.MakeEmbeddingLayer(tensors, "embedding", m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if embedTokens == nil { + return fmt.Errorf("failed to construct embedding layer from %q", "embedding.weight") + } + m.EmbedTokens = embedTokens + + outputNormWeight, err := requireTensor(tensors, "output_norm.weight") + if err != nil { + return err + } + if err := validateTensorShape("output_norm.weight", outputNormWeight, []int{int(m.HiddenSize)}, "hidden_size"); err != nil { + return err + } + m.Norm = nn.NewRMSNorm(outputNormWeight, m.RMSNormEps) + + if _, err := requireTensor(tensors, "output.weight"); err != nil { + return err + } + if err := validateLinearLayerShape(-1, tensors, "output", int(m.VocabSize), int(m.HiddenSize), "vocab_size x hidden_size", m.Config); err != nil { + return err + } + m.LMHead = linears.Make("output") + if m.LMHead == nil { + return fmt.Errorf("failed to construct linear layer from %q", "output.weight") + } + + expectedQ := checked.q + expectedKV := checked.kv + ropeFreqs, ropeScale := buildGPTOSSRoPEFreqs(m.Config) + for i := range m.Layers { + prefix := fmt.Sprintf("blocks.%d", i) + + attnNorm, err := loadLayerNormTensor(tensors, i, prefix+".attn_norm.weight", int(m.HiddenSize), "hidden_size") + if err != nil { + return err + } + + qProj, err := loadLinearLayer(tensors, linears, m.Config, i, prefix+".q_proj", expectedQ, int(m.HiddenSize), "num_attention_heads * head_dim x hidden_size") + if err != nil { + return err + } + kProj, err := loadLinearLayer(tensors, linears, m.Config, i, prefix+".k_proj", expectedKV, int(m.HiddenSize), "num_key_value_heads * head_dim x hidden_size") + if err != nil { + return err + } + vProj, err := loadLinearLayer(tensors, linears, m.Config, i, prefix+".v_proj", expectedKV, int(m.HiddenSize), "num_key_value_heads * head_dim x hidden_size") + if err != nil { + return err + } + oProj, err := loadLinearLayer(tensors, linears, m.Config, i, prefix+".attn_out", int(m.HiddenSize), expectedQ, "hidden_size x num_attention_heads * head_dim") + if err != nil { + return err + } + + sinkSource, err := requireTensor(tensors, prefix+".attn_sinks") + if err != nil { + return fmt.Errorf("layer %d: %w", i, err) + } + sinks, err := prepareGPTOSSAttentionSinks(i, prefix+".attn_sinks", sinkSource, checked.attentionHeads) + if err != nil { + return err + } + + ffnNorm, err := loadLayerNormTensor(tensors, i, prefix+".ffn_norm.weight", int(m.HiddenSize), "hidden_size") + if err != nil { + return err + } + + router, err := loadLinearLayer(tensors, linears, m.Config, i, prefix+".router", int(m.NumLocalExperts), int(m.HiddenSize), "num_local_experts x hidden_size") + if err != nil { + return err + } + + gateUp, err := loadExpertPair(tensors, i, prefix+".experts.gate_up_proj", checked.gateUp, checked.hidden, m.Config) + if err != nil { + return err + } + down, err := loadExpertProjection(tensors, i, prefix+".experts.down_proj", checked.hidden, checked.intermediate, m.Config) + if err != nil { + return err + } + m.Layers[i] = &Layer{ + AttentionNorm: attnNorm, + Attention: &Attention{ + QProj: qProj, + KProj: kProj, + VProj: vProj, + OProj: oProj, + Sinks: sinks, + RoPEFreqs: ropeFreqs, + RoPEScale: ropeScale, + }, + FFNNorm: ffnNorm, + Router: router, + Experts: &Experts{ + GateUp: gateUp, + Down: down, + }, + } + } + + return nil +} + +func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) { + if m == nil || m.Config == nil || m.EmbedTokens == nil || m.Norm == nil || b == nil || b.InputIDs == nil { + return nil, nil + } + + dims := b.InputIDs.Dims() + if len(dims) != 2 { + panic(fmt.Sprintf("gpt-oss forward requires 2D token input, got %v", dims)) + } + + batchSize, seqLen := dims[0], dims[1] + out := m.forwardDense(b, caches, batchSize, seqLen) + return out, out +} + +func (m *Model) forwardDense(b *batch.Batch, caches []cache.Cache, batchSize, seqLen int) *mlx.Array { + 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.ForwardBatch(h, b, c, batchSize, seqLen, m.Config, i) + } + + return m.Norm.Forward(h, m.RMSNormEps) +} + +// Unembed projects hidden states back into vocabulary space. +func (m *Model) Unembed(x *mlx.Array) *mlx.Array { + if m == nil || m.LMHead == nil || x == nil { + return nil + } + return m.LMHead.Forward(x) +} + +func (l *Layer) Forward(x *mlx.Array, c cache.Cache, batchSize, seqLen int, cfg *Config, layerIndex int) *mlx.Array { + return l.ForwardBatch(x, batchForForward(c, seqLen), c, batchSize, seqLen, cfg, layerIndex) +} + +func (l *Layer) ForwardBatch(x *mlx.Array, b *batch.Batch, c cache.Cache, batchSize, seqLen int, cfg *Config, layerIndex int) *mlx.Array { + if l == nil || l.Attention == nil || l.AttentionNorm == nil || l.FFNNorm == nil || l.Router == nil || l.Experts == nil || x == nil || cfg == nil { + panic("gpt-oss layer is not fully loaded") + } + residual := x + x = l.AttentionNorm.Forward(x, cfg.RMSNormEps) + x = l.Attention.ForwardBatch(x, b, c, batchSize, seqLen, cfg, layerIndex) + if x == nil || !x.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention output is invalid", layerIndex)) + } + + h := residual.Add(x) + if h == nil || !h.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d residual add output is invalid", layerIndex)) + } + + x = l.FFNNorm.Forward(h, cfg.RMSNormEps) + router := l.Router.Forward(x) + if router == nil || !router.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d router output is invalid", layerIndex)) + } + + x = l.Experts.Forward(x, router, cfg, layerIndex) + if x == nil || !x.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d expert output is invalid", layerIndex)) + } + + return h.Add(x) +} + +func (a *Attention) Forward(x *mlx.Array, c cache.Cache, batchSize, seqLen int, cfg *Config, layerIndex int) *mlx.Array { + return a.ForwardBatch(x, batchForForward(c, seqLen), c, batchSize, seqLen, cfg, layerIndex) +} + +func batchForForward(c cache.Cache, seqLen int) *batch.Batch { + offset := 0 + if c != nil { + offset = c.Offset() + } + return &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, seqLen), + SeqOffsets: []int32{int32(offset)}, + SeqQueryLens: []int32{int32(seqLen)}, + } +} + +func (a *Attention) ForwardBatch(x *mlx.Array, b *batch.Batch, c cache.Cache, batchSize, seqLen int, cfg *Config, layerIndex int) *mlx.Array { + if a == nil || a.QProj == nil || a.KProj == nil || a.VProj == nil || a.OProj == nil || x == nil || cfg == nil { + return x + } + query := a.QProj.Forward(x) + key := a.KProj.Forward(x) + value := a.VProj.Forward(x) + if query == nil || key == nil || value == nil || !query.Valid() || !key.Valid() || !value.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention projections are invalid", layerIndex)) + } + + batchDim := int32(batchSize) + seq := int32(seqLen) + numHeads := cfg.NumAttentionHeads + numKVHeads := cfg.NumKeyValueHeads + headDim := cfg.HeadDim + + query = mlx.Reshape(query, batchDim, seq, numHeads, headDim) + key = mlx.Reshape(key, batchDim, seq, numKVHeads, headDim) + value = mlx.Reshape(value, batchDim, seq, numKVHeads, headDim) + if query == nil || key == nil || value == nil || !query.Valid() || !key.Valid() || !value.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention reshape is invalid", layerIndex)) + } + + query = mlx.Transpose(query, 0, 2, 1, 3) + key = mlx.Transpose(key, 0, 2, 1, 3) + value = mlx.Transpose(value, 0, 2, 1, 3) + if query == nil || key == nil || value == nil || !query.Valid() || !key.Valid() || !value.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention transpose is invalid", layerIndex)) + } + + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + attentionScale := float32(1.0 / math.Sqrt(float64(cfg.HeadDim))) + if a.RoPEFreqs != nil && a.RoPEFreqs.Valid() { + query = mlx.RoPEWithFreqs(query, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, a.RoPEFreqs) + key = mlx.RoPEWithFreqs(key, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, a.RoPEFreqs) + ropeScale := a.RoPEScale + if ropeScale == 0 { + ropeScale = 1 + } + attentionScale *= ropeScale * ropeScale + } else { + ropeBase, ropeScale, _ := cfg.RopeParameters() + query = mlx.RoPEWithBase(query, int(cfg.HeadDim), false, ropeBase, ropeScale, positions) + key = mlx.RoPEWithBase(key, int(cfg.HeadDim), false, ropeBase, ropeScale, positions) + } + if query == nil || key == nil || !query.Valid() || !key.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention RoPE is invalid", layerIndex)) + } + if err := validateGPTOSSAttentionSinkDispatch(query, a.Sinks, int(cfg.NumAttentionHeads)); err != nil { + panic(fmt.Sprintf("gpt-oss layer %d attention sink dispatch: %v", layerIndex, err)) + } + + var kv nn.SDPAOption + if c != nil { + attnCache, ok := c.(cache.Attention) + if !ok { + panic(fmt.Sprintf("gpt-oss layer %d cache does not support attention", layerIndex)) + } + history := attnCache.Update(b, key, value) + if history == nil || history.K() == nil || history.V() == nil || !history.K().Valid() || !history.V().Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention cache update is invalid", layerIndex)) + } + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(key, value, b.SeqQueryLens) + } + + attention := nn.ScaledDotProductAttention( + b, + query, + attentionScale, + kv, + nn.WithMask(nn.CausalMask()), + nn.WithSinks(a.Sinks), + ) + if attention == nil || !attention.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention sdpa is invalid", layerIndex)) + } + attention = mlx.Transpose(attention, 0, 2, 1, 3) + attention = mlx.Reshape(attention, batchDim, seq, int32(cfg.checked.q)) + if attention == nil || !attention.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention output reshape is invalid", layerIndex)) + } + attention = a.OProj.Forward(attention) + if attention == nil || !attention.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d attention output projection is invalid", layerIndex)) + } + return attention } -// LoadWeights is intentionally skeletal for Phase 2. -func (m *Model) LoadWeights(map[string]*mlx.Array) error { - return fmt.Errorf("gpt-oss weight loading is not implemented yet") +func (p *ExpertProjection) Forward(x, indices *mlx.Array, sorted bool) *mlx.Array { + if p == nil || p.Weight == nil || x == nil || indices == nil { + return nil + } + + var out *mlx.Array + if p.Scales != nil { + out = mlx.GatherQMM(x, p.Weight, p.Scales, p.QBiases, nil, indices, p.Transpose, p.GroupSize, p.Bits, p.Mode, sorted) + } else { + if x.DType() != p.Weight.DType() { + x = x.AsType(p.Weight.DType()) + } + // Keep dense GatherMM on the generic path. The sorted dense fast path + // fails GPT-OSS expert parity with Xcode 26.5 / MetalToolchain 17.6. + out = mlx.GatherMM(x, p.Weight, nil, indices, false) + } + + if p.Bias == nil || !p.Bias.Valid() { + return out + } + + bias := p.Bias.TakeAxis(indices, 0) + bias = mlx.ExpandDims(bias, 2) + return mlx.Add(out, bias) +} + +func (e *Experts) Forward(x, router *mlx.Array, cfg *Config, layerIndex int) *mlx.Array { + if e == nil || e.GateUp == nil || e.GateUp.Gate == nil || e.GateUp.Up == nil || e.Down == nil || x == nil || router == nil || cfg == nil { + panic("gpt-oss expert path is not fully loaded") + } + if !x.Valid() || !router.Valid() { + panic("gpt-oss expert path received invalid tensors") + } + + dims := x.Dims() + if len(dims) != 3 { + panic(fmt.Sprintf("gpt-oss expert path expects 3D hidden states, got %v", dims)) + } + + B, L := int32(dims[0]), int32(dims[1]) + topK := cfg.NumExpertsPerTok + + neg := mlx.Neg(router) + inds := mlx.Argpartition(neg, int(topK)-1, -1) + shape := inds.Dims() + inds = mlx.SliceStartStop(inds, []int32{0, 0, 0}, []int32{int32(shape[0]), int32(shape[1]), topK}) + + scores := mlx.TakeAlongAxis(router, inds, -1) + scores = mlx.SoftmaxAxis(scores, -1, true) + + var xFlat *mlx.Array + if B == 1 && L == 1 { + xFlat = mlx.Reshape(x, 1, 1, 1, cfg.HiddenSize) + } else { + xExpanded := mlx.ExpandDims(mlx.ExpandDims(x, -2), -2) + xFlat = mlx.Reshape(xExpanded, B*L, 1, 1, cfg.HiddenSize) + } + idxFlat := mlx.Reshape(inds, B*L, topK) + + doSort := B*L >= 24 + 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) + } + + // Try fully fused gate+up+SwiGLU+down kernel path for single-token MXFP4 MoE (decode). + canFuse := !doSort && B*L == 1 && + fusedExpertPairEligible(e.GateUp.Gate, e.GateUp.Up) + + var down *mlx.Array + if canFuse { + gateW := e.GateUp.Gate + upW := e.GateUp.Up + wDims := gateW.Weight.Dims() + if len(wDims) == 3 { + numRows := wDims[1] + numColVecs := int(cfg.HiddenSize) / 32 + inputFlat := mlx.Reshape(x, B*L, cfg.HiddenSize).AsType(mlx.DTypeFloat32) + + swiGLUOut, ok := mlx.MoEFusedGateUpSwiGLU( + inputFlat, + gateW.Weight, gateW.Scales, gateW.Bias, + upW.Weight, upW.Scales, upW.Bias, + idxFlat, + numRows, numColVecs, int(topK), + -7.0, 7.0, // swiglu clamp limits + ) + if ok && swiGLUOut != nil { + // Try fused down projection on the float32 SwiGLU output. + if fusedExpertProjectionEligible(e.Down) { + downDims := e.Down.Weight.Dims() + if len(downDims) == 3 { + downNumRows := int(cfg.HiddenSize) + downNumColVecs := numRows / 32 // intermediateSize / 32 + + downOut, dok := mlx.MoEFusedDown( + swiGLUOut, + e.Down.Weight, e.Down.Scales, e.Down.Bias, + idxFlat, + downNumRows, downNumColVecs, int(topK), + ) + if dok && downOut != nil { + // downOut is [batch, topK, hiddenSize] in float32 + down = downOut.AsType(mlx.DTypeBFloat16) + } + } + } + // If fused down failed, fall back to GatherQMM for just the down proj + if down == nil { + hidden := mlx.Reshape(swiGLUOut, B*L, topK, 1, int32(numRows)).AsType(mlx.DTypeBFloat16) + down = e.Down.Forward(hidden, idxFlat, doSort) + if down == nil || !down.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d expert down projection is invalid", layerIndex)) + } + down = mlx.Squeeze(down, 2) + } + } + } + } + + // Fully unfused fallback path. + if down == nil { + gate := e.GateUp.Gate.Forward(xFlat, idxFlat, doSort) + up := e.GateUp.Up.Forward(xFlat, idxFlat, doSort) + if gate == nil || !gate.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d expert gate projection is invalid", layerIndex)) + } + if up == nil || !up.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d expert up projection is invalid", layerIndex)) + } + hidden := swiGLUAlphaLimit(gate, up) + + downResult := e.Down.Forward(hidden, idxFlat, doSort) + if downResult == nil || !downResult.Valid() { + panic(fmt.Sprintf("gpt-oss layer %d expert down projection is invalid", layerIndex)) + } + + if doSort { + down = mlx.Reshape( + mlx.Take(mlx.Squeeze(mlx.Squeeze(downResult, 2), 1), invOrder, 0), + B*L, topK, cfg.HiddenSize, + ) + } else { + down = mlx.Squeeze(downResult, 2) + } + } + + down = mlx.Reshape(down, B, L, topK, cfg.HiddenSize) + return mlx.Sum(mlx.Mul(down, mlx.ExpandDims(scores, -1)), 2, false) } diff --git a/x/models/gptoss/gptoss_row10_test.go b/x/models/gptoss/gptoss_row10_test.go new file mode 100644 index 00000000000..614a1223581 --- /dev/null +++ b/x/models/gptoss/gptoss_row10_test.go @@ -0,0 +1,368 @@ +package gptoss + +import ( + "context" + "fmt" + "math" + "runtime" + "strings" + "sync" + "testing" + + "github.com/ollama/ollama/x/internal/mlxthread" + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" +) + +func row10ValidConfig() Config { + return Config{ + NumHiddenLayers: 2, HiddenSize: 64, IntermediateSize: 64, NumAttentionHeads: 2, + NumKeyValueHeads: 1, HeadDim: 32, NumLocalExperts: 2, NumExpertsPerTok: 1, + SlidingWindow: 128, VocabSize: 256, MaxPositionEmbeddings: 1024, + RopeTheta: 10000, RMSNormEps: 1e-5, RopeScaling: RopeScaling{Factor: 1, OriginalMaxPositionEmbeddings: 1024}, + } +} + +func TestConfigDimensions(t *testing.T) { + base := row10ValidConfig() + tests := []struct { + name string + mutate func(*Config) + wantErr string + }{ + {name: "valid", mutate: func(*Config) {}}, + {name: "zero scalar", mutate: func(c *Config) { c.HiddenSize = 0 }, wantErr: "hidden_size"}, + {name: "negative scalar", mutate: func(c *Config) { c.NumLocalExperts = -1 }, wantErr: "num_local_experts"}, + {name: "loop ceiling accepted", mutate: func(c *Config) { c.NumHiddenLayers = gptossMaxLoopEntries; c.NumLocalExperts = 1 }}, + {name: "loop ceiling rejected", mutate: func(c *Config) { c.NumHiddenLayers = gptossMaxLoopEntries + 1 }, wantErr: "num_hidden_layers"}, + {name: "dimension ceiling accepted", mutate: func(c *Config) { c.HiddenSize = gptossMaxDimension }}, + {name: "dimension ceiling rejected", mutate: func(c *Config) { c.HiddenSize = gptossMaxDimension + 1 }, wantErr: "hidden_size"}, + {name: "query product accepted", mutate: func(c *Config) { c.NumAttentionHeads = 1 << 12; c.HeadDim = 1 << 12 }, wantErr: ""}, + {name: "query product rejected", mutate: func(c *Config) { c.NumAttentionHeads = 1 << 12; c.HeadDim = (1 << 12) + 1 }, wantErr: "query dimension"}, + {name: "int32 wrapping head product", mutate: func(c *Config) { c.NumAttentionHeads = 1 << 16; c.HeadDim = 1 << 16 }, wantErr: "query dimension"}, + {name: "gate up product rejected", mutate: func(c *Config) { c.IntermediateSize = (gptossMaxDimension / 2) + 1 }, wantErr: "gate/up dimension"}, + {name: "layer expert work rejected", mutate: func(c *Config) { c.NumHiddenLayers = 1025; c.NumLocalExperts = 1024 }, wantErr: "layer/expert work"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := base + tt.mutate(&cfg) + _, err := validateGPTOSSConfigDimensions(&cfg) + if tt.wantErr == "" && err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tt.wantErr != "" && (err == nil || !strings.Contains(err.Error(), tt.wantErr)) { + t.Fatalf("error = %v, want containing %q", err, tt.wantErr) + } + }) + } +} + +func TestValidateGPTOSSScalarDimensions(t *testing.T) { + setters := []struct { + name string + set func(*Config, int32) + }{ + {"layers", func(c *Config, v int32) { c.NumHiddenLayers = v }}, + {"hidden", func(c *Config, v int32) { c.HiddenSize = v }}, + {"intermediate", func(c *Config, v int32) { c.IntermediateSize = v }}, + {"attention heads", func(c *Config, v int32) { c.NumAttentionHeads = v }}, + {"kv heads", func(c *Config, v int32) { c.NumKeyValueHeads = v }}, + {"head dim", func(c *Config, v int32) { c.HeadDim = v }}, + {"experts", func(c *Config, v int32) { c.NumLocalExperts = v }}, + {"experts per token", func(c *Config, v int32) { c.NumExpertsPerTok = v }}, + {"sliding window", func(c *Config, v int32) { c.SlidingWindow = v }}, + {"vocab", func(c *Config, v int32) { c.VocabSize = v }}, + {"context", func(c *Config, v int32) { c.MaxPositionEmbeddings = v }}, + {"original context", func(c *Config, v int32) { c.RopeScaling.OriginalMaxPositionEmbeddings = v }}, + } + for _, setter := range setters { + for _, value := range []int32{0, -1} { + t.Run(setter.name, func(t *testing.T) { + cfg := row10ValidConfig() + setter.set(&cfg, value) + if _, err := validateGPTOSSConfigDimensions(&cfg); err == nil { + t.Fatalf("value %d accepted", value) + } + }) + } + } +} + +func row10ValidFusedDescriptor() fusedExpertDescriptor { + return fusedExpertDescriptor{ + present: true, transpose: true, mode: "mxfp4", bits: 4, groupSize: 32, + weightType: mlx.DTypeUint32, scaleType: mlx.DTypeUint8, biasType: mlx.DTypeBFloat16, + weightDims: []int{2, 64, 8}, scaleDims: []int{2, 64, 2}, biasDims: []int{2, 64}, + } +} + +func TestFusedExpertEligibility(t *testing.T) { + base := row10ValidFusedDescriptor() + tests := []struct { + name string + mutate func(*fusedExpertDescriptor) + }{ + {name: "mode", mutate: func(d *fusedExpertDescriptor) { d.mode = "affine" }}, + {name: "bits", mutate: func(d *fusedExpertDescriptor) { d.bits = 8 }}, + {name: "group", mutate: func(d *fusedExpertDescriptor) { d.groupSize = 64 }}, + {name: "qbias", mutate: func(d *fusedExpertDescriptor) { d.hasQBias = true }}, + {name: "transpose", mutate: func(d *fusedExpertDescriptor) { d.transpose = false }}, + {name: "weight dtype", mutate: func(d *fusedExpertDescriptor) { d.weightType = mlx.DTypeUint8 }}, + {name: "scale dtype", mutate: func(d *fusedExpertDescriptor) { d.scaleType = mlx.DTypeBFloat16 }}, + {name: "bias dtype", mutate: func(d *fusedExpertDescriptor) { d.biasType = mlx.DTypeFloat32 }}, + {name: "weight rank", mutate: func(d *fusedExpertDescriptor) { d.weightDims = []int{2, 64} }}, + {name: "scale rank", mutate: func(d *fusedExpertDescriptor) { d.scaleDims = []int{2, 64} }}, + {name: "bias rank", mutate: func(d *fusedExpertDescriptor) { d.biasDims = []int{128} }}, + } + if !fusedExpertDescriptorEligible(base) { + t.Fatal("valid descriptor rejected") + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := base + tt.mutate(&got) + if fusedExpertDescriptorEligible(got) { + t.Fatal("malformed descriptor accepted") + } + }) + } + pairMutations := []struct { + name string + mutate func(*fusedExpertDescriptor) + }{ + {"mode", func(d *fusedExpertDescriptor) { d.mode = "affine" }}, + {"bits", func(d *fusedExpertDescriptor) { d.bits = 8 }}, + {"group", func(d *fusedExpertDescriptor) { d.groupSize = 64 }}, + {"qbias", func(d *fusedExpertDescriptor) { d.hasQBias = true }}, + {"transpose", func(d *fusedExpertDescriptor) { d.transpose = false }}, + {"shape", func(d *fusedExpertDescriptor) { d.weightDims = []int{2, 63, 8} }}, + {"dtype", func(d *fusedExpertDescriptor) { d.scaleType = mlx.DTypeBFloat16 }}, + } + for _, side := range []string{"gate", "up"} { + for _, mutation := range pairMutations { + t.Run(side+" "+mutation.name, func(t *testing.T) { + gate, up := base, base + if side == "gate" { + mutation.mutate(&gate) + } else { + mutation.mutate(&up) + } + if fusedExpertDescriptorsPairEligible(gate, up) { + t.Fatal("asymmetric pair accepted") + } + }) + } + } +} + +var row10MLXMu sync.Mutex + +func row10WithMLX(t *testing.T, fn func()) { + t.Helper() + row10MLXMu.Lock() + defer row10MLXMu.Unlock() + runtime.LockOSThread() + defer runtime.UnlockOSThread() + thread, err := mlxthread.Start("gptoss-row10-test", func() error { + if err := mlx.CheckInit(); err != nil { + return err + } + if mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + return nil + }) + if err != nil { + t.Skipf("MLX not available: %v", err) + } + defer func() { _ = thread.Stop(context.Background(), func() { mlx.Sweep(); mlx.ClearCache() }) }() + if err := thread.Do(context.Background(), func() error { fn(); return nil }); err != nil { + t.Fatal(err) + } +} + +func row10NativeExpertTensors(prefix string, experts, out, in int) map[string]*mlx.Array { + return map[string]*mlx.Array{ + prefix + ".weight": mlx.Zeros(mlx.DTypeUint32, experts, out, in/8), + prefix + ".weight_scale": mlx.Zeros(mlx.DTypeUint8, experts, out, in/32), + prefix + ".bias": mlx.Zeros(mlx.DTypeBFloat16, experts, out), + } +} + +func TestDirectExpertValidation(t *testing.T) { + row10WithMLX(t, func() { + cfg := row10ValidConfig() + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = 32, 4, "mxfp4" + valid := func(ts map[string]*mlx.Array) error { + _, err := loadDirectExpertProjection(ts, 0, "p", 64, 64, &cfg) + return err + } + if err := valid(row10NativeExpertTensors("p", 2, 64, 64)); err != nil { + t.Fatalf("valid native projection rejected: %v", err) + } + for _, tt := range []struct { + name string + mutate func(map[string]*mlx.Array) + }{ + {name: "weight dtype", mutate: func(ts map[string]*mlx.Array) { ts["p.weight"] = mlx.Zeros(mlx.DTypeUint8, 2, 64, 8) }}, + {name: "weight rank", mutate: func(ts map[string]*mlx.Array) { ts["p.weight"] = mlx.Zeros(mlx.DTypeUint32, 2, 512) }}, + {name: "weight experts", mutate: func(ts map[string]*mlx.Array) { ts["p.weight"] = mlx.Zeros(mlx.DTypeUint32, 3, 64, 8) }}, + {name: "weight output", mutate: func(ts map[string]*mlx.Array) { ts["p.weight"] = mlx.Zeros(mlx.DTypeUint32, 2, 63, 8) }}, + {name: "packed axis", mutate: func(ts map[string]*mlx.Array) { ts["p.weight"] = mlx.Zeros(mlx.DTypeUint32, 2, 64, 9) }}, + {name: "scale dtype", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_scale"] = mlx.Zeros(mlx.DTypeBFloat16, 2, 64, 2) }}, + {name: "scale rank", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, 2, 128) }}, + {name: "scale experts", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, 3, 64, 2) }}, + {name: "scale output", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, 2, 63, 2) }}, + {name: "scale groups", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, 2, 64, 3) }}, + {name: "bias dtype", mutate: func(ts map[string]*mlx.Array) { ts["p.bias"] = mlx.Zeros(mlx.DTypeFloat32, 2, 64) }}, + {name: "bias rank", mutate: func(ts map[string]*mlx.Array) { ts["p.bias"] = mlx.Zeros(mlx.DTypeBFloat16, 128) }}, + {name: "bias experts", mutate: func(ts map[string]*mlx.Array) { ts["p.bias"] = mlx.Zeros(mlx.DTypeBFloat16, 3, 64) }}, + {name: "bias axis", mutate: func(ts map[string]*mlx.Array) { ts["p.bias"] = mlx.Zeros(mlx.DTypeBFloat16, 2, 63) }}, + {name: "qbias", mutate: func(ts map[string]*mlx.Array) { ts["p.weight_qbias"] = mlx.Zeros(mlx.DTypeUint8, 2, 64, 2) }}, + {name: "missing weight", mutate: func(ts map[string]*mlx.Array) { delete(ts, "p.weight") }}, + {name: "missing bias", mutate: func(ts map[string]*mlx.Array) { delete(ts, "p.bias") }}, + {name: "missing scale", mutate: func(ts map[string]*mlx.Array) { delete(ts, "p.weight_scale") }}, + {name: "orphan scale", mutate: func(ts map[string]*mlx.Array) { delete(ts, "p.weight"); delete(ts, "p.bias") }}, + {name: "orphan qbias", mutate: func(ts map[string]*mlx.Array) { clear(ts); ts["p.weight_qbias"] = mlx.Zeros(mlx.DTypeUint8, 2, 64, 2) }}, + } { + t.Run(tt.name, func(t *testing.T) { + ts := row10NativeExpertTensors("p", 2, 64, 64) + tt.mutate(ts) + if err := valid(ts); err == nil { + t.Fatal("malformed projection accepted") + } + }) + } + if err := valid(row10NativeExpertTensors("p", 2, 64, 64)); err != nil { + t.Fatalf("valid projection rejected after malformed cases: %v", err) + } + for _, projection := range []struct { + prefix string + wantOut, in int + }{ + {"model.layers.0.mlp.experts.gate_proj", 64, 64}, + {"model.layers.0.mlp.experts.up_proj", 64, 64}, + {"model.layers.0.mlp.experts.down_proj", 64, 64}, + } { + if _, err := loadDirectExpertProjection(row10NativeExpertTensors(projection.prefix, 2, projection.wantOut, projection.in), 0, projection.prefix, projection.wantOut, projection.in, &cfg); err != nil { + t.Fatalf("valid native %s rejected: %v", projection.prefix, err) + } + } + pairPrefix := "model.layers.0.mlp.experts.gate_up_proj" + gatePrefix := strings.Replace(pairPrefix, "gate_up_proj", "gate_proj", 1) + upPrefix := strings.Replace(pairPrefix, "gate_up_proj", "up_proj", 1) + gateOnly := row10NativeExpertTensors(gatePrefix, 2, 64, 64) + if _, err := loadDirectExpertPair(gateOnly, 0, pairPrefix, 64, 64, &cfg); err == nil { + t.Fatal("gate-only native pair accepted") + } + upOnly := row10NativeExpertTensors(upPrefix, 2, 64, 64) + if _, err := loadDirectExpertPair(upOnly, 0, pairPrefix, 64, 64, &cfg); err == nil { + t.Fatal("up-only native pair accepted") + } + }) +} + +func TestAttentionSinkValidation(t *testing.T) { + row10WithMLX(t, func() { + source := mlx.FromValues([]float32{0.25, -0.5}, 2).AsType(mlx.DTypeBFloat16) + prepared, err := prepareGPTOSSAttentionSinks(0, "sink", source, 2) + if err != nil { + t.Fatal(err) + } + if source.DType() != mlx.DTypeBFloat16 { + t.Fatalf("source dtype drifted to %s", source.DType()) + } + if prepared == source || prepared.DType() != mlx.DTypeBFloat16 { + t.Fatalf("prepared sinks = %p dtype %s, source = %p", prepared, prepared.DType(), source) + } + preparedValues := prepared.AsType(mlx.DTypeFloat32) + mlx.Eval(preparedValues) + if got := preparedValues.Floats(); len(got) != 2 || got[0] != 0.25 || got[1] != -0.5 { + t.Fatalf("prepared materialized values = %v", got) + } + preparedIdentity := prepared + for range 3 { + if prepared != preparedIdentity { + t.Fatal("prepared sink identity changed") + } + if err := validateGPTOSSAttentionSinkDispatch(mlx.Zeros(mlx.DTypeBFloat16, 1, 2, 1, 4), prepared, 2); err != nil { + t.Fatalf("prepared dispatch proof: %v", err) + } + } + if _, err := prepareGPTOSSAttentionSinks(0, "sink", mlx.Zeros(mlx.DTypeFloat32, 2), 2); err == nil { + t.Fatal("F32 checkpoint sink accepted") + } + if _, err := prepareGPTOSSAttentionSinks(0, "sink", mlx.Zeros(mlx.DTypeUint8, 2), 2); err == nil { + t.Fatal("U8 checkpoint sink accepted") + } + for _, value := range []float32{float32(math.NaN()), float32(math.Inf(1)), float32(math.Inf(-1))} { + bad := mlx.FromValues([]float32{value, 0}, 2).AsType(mlx.DTypeBFloat16) + if _, err := prepareGPTOSSAttentionSinks(0, "sink", bad, 2); err == nil || !strings.Contains(err.Error(), "non-finite") { + t.Fatalf("non-finite %v error = %v", value, err) + } + } + if err := validateGPTOSSAttentionSinkDispatch(mlx.Zeros(mlx.DTypeFloat32, 1, 2, 1, 4), prepared, 2); err == nil || err.Error() != "prepared sink dtype BF16 must equal actual query dtype F32" { + t.Fatalf("query/sink dtype mismatch error = %v", err) + } + if err := validateGPTOSSAttentionSinkDispatch(mlx.Zeros(mlx.DTypeBFloat16, 1, 2, 4), prepared, 2); err == nil { + t.Fatal("rank-3 query accepted") + } + if err := validateGPTOSSAttentionSinkDispatch(mlx.Zeros(mlx.DTypeBFloat16, 1, 3, 1, 4), prepared, 2); err == nil || err.Error() != "query head axis = 3, want configured heads = 2 (shape [1 3 1 4])" { + t.Fatalf("query head mismatch error = %v", err) + } + if err := validateGPTOSSAttentionSinkDispatch(mlx.Zeros(mlx.DTypeBFloat16, 1, 2, 1, 4), mlx.Zeros(mlx.DTypeBFloat16, 1), 2); err == nil { + t.Fatal("sink head mismatch accepted") + } + }) +} + +func TestSharedSDPASinkDTypeAndMemoIsolation(t *testing.T) { + row10WithMLX(t, func() { + call := func(dtype, sinkType mlx.DType, withSink bool, b *batch.Batch) { + q := mlx.Zeros(dtype, 1, 1, 1, 4) + k := mlx.Zeros(dtype, 1, 1, 1, 4) + v := mlx.Zeros(dtype, 1, 1, 1, 4) + opts := []nn.SDPAOption{nn.WithKV(k, v, []int32{1})} + if withSink { + opts = append(opts, nn.WithSinks(mlx.Zeros(sinkType, 1))) + } + out := nn.ScaledDotProductAttention(b, q, 0.5, opts...) + if out == nil || !out.Valid() { + t.Fatal("SDPA returned invalid output") + } + } + newBatch := func() *batch.Batch { + return &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{1}, + } + } + call(mlx.DTypeFloat32, mlx.DTypeFloat32, true, newBatch()) + call(mlx.DTypeBFloat16, mlx.DTypeBFloat16, true, newBatch()) + for _, mismatch := range []struct{ query, sink mlx.DType }{ + {mlx.DTypeFloat32, mlx.DTypeBFloat16}, + {mlx.DTypeBFloat16, mlx.DTypeFloat32}, + } { + func() { + defer func() { + got := recover() + if got == nil { + t.Fatal("query/sink dtype mismatch accepted") + } + want := fmt.Sprintf("mlx.FastScaledDotProductAttention: sinks must have shape [heads]=[1] and dtype %s for query [1 1 1 4], got shape [1] dtype %s", mismatch.query, mismatch.sink) + if fmt.Sprint(got) != want { + t.Fatalf("panic = %q, want %q", got, want) + } + }() + call(mismatch.query, mismatch.sink, true, newBatch()) + }() + } + shared := newBatch() + call(mlx.DTypeFloat32, mlx.DTypeFloat32, false, shared) + call(mlx.DTypeFloat32, mlx.DTypeFloat32, true, shared) + call(mlx.DTypeFloat32, mlx.DTypeFloat32, false, shared) + }) +} From 0db9c95d0b6dbd45fd8041516c4a213d0ee44076 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 18:12:06 +0000 Subject: [PATCH 22/58] test: cover GPT-OSS runtime and import contracts Co-authored-by: Codex --- x/models/gptoss/forward_reference_test.go | 9 +- x/models/gptoss/gptoss_test.go | 3557 +++++++++++++++++++++ 2 files changed, 3563 insertions(+), 3 deletions(-) create mode 100644 x/models/gptoss/gptoss_test.go diff --git a/x/models/gptoss/forward_reference_test.go b/x/models/gptoss/forward_reference_test.go index 0b328417c17..9858f2789e9 100644 --- a/x/models/gptoss/forward_reference_test.go +++ b/x/models/gptoss/forward_reference_test.go @@ -1,5 +1,3 @@ -//go:build gptoss_forward_reference - package gptoss import ( @@ -241,7 +239,12 @@ func materializedInts(a *mlx.Array) []int { } cloned := a.Clone() mlx.Eval(cloned) - return cloned.Ints() + values := cloned.Ints() + ints := make([]int, len(values)) + for i, value := range values { + ints[i] = int(value) + } + return ints } func compareForwardReferenceArrays(t *testing.T, name string, got, want *mlx.Array, absTol float64) { diff --git a/x/models/gptoss/gptoss_test.go b/x/models/gptoss/gptoss_test.go new file mode 100644 index 00000000000..e3f96579b87 --- /dev/null +++ b/x/models/gptoss/gptoss_test.go @@ -0,0 +1,3557 @@ +package gptoss + +import ( + "fmt" + "math" + "os" + "path/filepath" + "runtime" + "slices" + "strings" + "sync" + "testing" + + "github.com/ollama/ollama/x/imagegen/manifest" + "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" +) + +var mlxTestMu sync.Mutex + +const row11NoSinkSentinel float32 = -1e9 + +func forwardModel(m *Model, tokens *mlx.Array, caches []cache.Cache) *mlx.Array { + hidden, _ := m.Forward(testBatch(tokens, 0), caches) + return hidden +} + +func testBatch(tokens *mlx.Array, offset int) *batch.Batch { + return &batch.Batch{ + InputIDs: tokens, + SeqOffsets: []int32{int32(offset)}, + SeqQueryLens: []int32{int32(tokens.Dim(1))}, + } +} + +func testForwardBatch(tokens *mlx.Array, c cache.Cache) *batch.Batch { + offset := 0 + if c != nil { + offset = c.Offset() + } + return testBatch(tokens, offset) +} + +func row11SliceSequence(x *mlx.Array, pos int) *mlx.Array { + return x.Slice(mlx.Slice(), mlx.Slice(pos, pos+1), mlx.Slice()) +} + +func testGPTOSSRoPEFreqs(cfg *Config) *mlx.Array { + freqs, _ := buildGPTOSSRoPEFreqs(cfg) + return freqs +} + +func testGPTOSSYarnScale(cfg *Config) float32 { + _, scale := buildGPTOSSRoPEFreqs(cfg) + return scale +} + +func refreshRow11CheckedConfig(t *testing.T, cfg *Config) { + t.Helper() + checked, err := validateGPTOSSConfigDimensions(cfg) + if err != nil { + t.Fatalf("validate mutated GPT-OSS config: %v", err) + } + cfg.checked = checked +} + +func TestParseConfig(t *testing.T) { + data := []byte(`{ + "architectures": ["GptOssForCausalLM"], + "model_type": "gpt_oss", + "num_hidden_layers": 24, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 64, + "num_local_experts": 32, + "num_experts_per_tok": 4, + "sliding_window": 128, + "rope_theta": 150000, + "rope_scaling": { + "factor": 32.0, + "original_max_position_embeddings": 4096, + "rope_type": "yarn", + "beta_fast": 32.0, + "beta_slow": 1.0, + "truncate": false + }, + "rms_norm_eps": 0.00001, + "vocab_size": 201088, + "tie_word_embeddings": false, + "quantization_config": { + "quant_method": "mxfp4" + } + }`) + + cfg, err := parseConfig(data) + if err != nil { + t.Fatalf("parseConfig() error = %v", err) + } + + if cfg.Architecture != "GptOssForCausalLM" { + t.Fatalf("Architecture = %q, want %q", cfg.Architecture, "GptOssForCausalLM") + } + if cfg.ModelType != "gpt_oss" { + t.Fatalf("ModelType = %q, want %q", cfg.ModelType, "gpt_oss") + } + if cfg.NumHiddenLayers != 24 || cfg.HiddenSize != 2880 || cfg.IntermediateSize != 2880 { + t.Fatalf("unexpected core dims: %+v", cfg) + } + if cfg.NumAttentionHeads != 64 || cfg.NumKeyValueHeads != 8 || cfg.HeadDim != 64 { + t.Fatalf("unexpected attention dims: %+v", cfg) + } + if cfg.NumLocalExperts != 32 || cfg.NumExpertsPerTok != 4 { + t.Fatalf("unexpected expert dims: %+v", cfg) + } + if cfg.MaxPositionEmbeddings != 131072 { + t.Fatalf("MaxPositionEmbeddings = %d, want 131072", cfg.MaxPositionEmbeddings) + } + if cfg.QuantMethod != "mxfp4" { + t.Fatalf("QuantMethod = %q, want %q", cfg.QuantMethod, "mxfp4") + } +} + +func TestNewModelRegistersGptOss(t *testing.T) { + root := testRoot(t, []byte(`{ + "architectures": ["GptOssForCausalLM"], + "model_type": "gpt_oss", + "num_hidden_layers": 24, + "hidden_size": 2880, + "intermediate_size": 2880, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 64, + "num_local_experts": 32, + "num_experts_per_tok": 4, + "sliding_window": 128, + "rope_theta": 150000, + "rope_scaling": { + "factor": 32.0, + "original_max_position_embeddings": 4096 + }, + "rms_norm_eps": 0.00001, + "vocab_size": 201088, + "tie_word_embeddings": false, + "quantization_config": { + "quant_method": "mxfp4" + } + }`), []byte(`{ + "model": { + "type": "BPE", + "vocab": {"a": 0, "b": 1}, + "merges": [] + }, + "added_tokens": [] + }`)) + + m, err := base.New(root) + if err != nil { + t.Fatalf("base.New() error = %v", err) + } + + got, ok := m.(*Model) + if !ok { + t.Fatalf("base.New() type = %T, want *Model", m) + } + + if got.Tokenizer() == nil { + t.Fatal("Tokenizer() = nil, want loaded tokenizer") + } + if got.NumLayers() != 24 { + t.Fatalf("NumLayers() = %d, want 24", got.NumLayers()) + } + if got.MaxContextLength() != 131072 { + t.Fatalf("MaxContextLength() = %d, want 131072", got.MaxContextLength()) + } + if got.Architecture != "GptOssForCausalLM" { + t.Fatalf("Architecture = %q, want %q", got.Architecture, "GptOssForCausalLM") + } + if got.QuantMethod != "mxfp4" { + t.Fatalf("QuantMethod = %q, want %q", got.QuantMethod, "mxfp4") + } +} + +func TestNewModelRestoresHarmonyStopTokens(t *testing.T) { + root := testRootWithExtraConfigs(t, []byte(`{ + "architectures": ["GptOssForCausalLM"], + "model_type": "gpt_oss", + "num_hidden_layers": 1, + "hidden_size": 64, + "intermediate_size": 128, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 64, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "sliding_window": 128, + "rope_theta": 150000, + "rope_scaling": { + "factor": 32.0, + "original_max_position_embeddings": 4096 + }, + "rms_norm_eps": 0.00001, + "vocab_size": 201088, + "tie_word_embeddings": false + }`), []byte(`{ + "model": { + "type": "BPE", + "vocab": {"a": 0}, + "merges": [] + }, + "added_tokens": [ + {"id": 199998, "content": "<|startoftext|>", "special": true}, + {"id": 199999, "content": "<|endoftext|>", "special": true}, + {"id": 200002, "content": "<|return|>", "special": true}, + {"id": 200012, "content": "<|call|>", "special": true} + ] + }`), map[string][]byte{ + "generation_config.json": []byte(`{ + "eos_token_id": [200002, 199999] + }`), + "tokenizer_config.json": []byte(`{ + "bos_token": "<|startoftext|>", + "eos_token": "<|return|>" + }`), + }) + + m, err := base.New(root) + if err != nil { + t.Fatalf("base.New() error = %v", err) + } + + got := m.(*Model).Tokenizer() + want := []int32{200002, 199999, 199999, 200002, 200012} + if !slices.Equal(got.EOSTokens(), want) { + t.Fatalf("Tokenizer().EOSTokens() = %v, want %v", got.EOSTokens(), want) + } + if got.EOS() != 200002 { + t.Fatalf("Tokenizer().EOS() = %d, want %d", got.EOS(), 200002) + } + + if !got.IsEOS(200012) { + t.Fatal("Tokenizer().IsEOS(<|call|>) = false, want true") + } +} + +func TestLoadWeightsDensePath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + if err := m.LoadWeights(tensors); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + + if m.EmbedTokens == nil { + t.Fatal("EmbedTokens = nil") + } + if m.Norm == nil { + t.Fatal("Norm = nil") + } + if m.LMHead == nil { + t.Fatal("LMHead = nil") + } + + if len(m.Layers) != int(cfg.NumHiddenLayers) { + t.Fatalf("len(Layers) = %d, want %d", len(m.Layers), cfg.NumHiddenLayers) + } + for i, layer := range m.Layers { + if layer == nil { + t.Fatalf("layer %d = nil", i) + } + if layer.AttentionNorm == nil { + t.Fatalf("layer %d AttentionNorm = nil", i) + } + if layer.FFNNorm == nil { + t.Fatalf("layer %d FFNNorm = nil", i) + } + if layer.Router == nil { + t.Fatalf("layer %d Router = nil", i) + } + if layer.Attention == nil { + t.Fatalf("layer %d Attention = nil", i) + } + if layer.Attention.QProj == nil || layer.Attention.KProj == nil || layer.Attention.VProj == nil || layer.Attention.OProj == nil { + t.Fatalf("layer %d attention projections not fully loaded", i) + } + if layer.Attention.Sinks == nil { + t.Fatalf("layer %d Attention.Sinks = nil", i) + } + if layer.Attention.Sinks.DType() != mlx.DTypeBFloat16 { + t.Fatalf("layer %d Attention.Sinks dtype = %s, want BF16 prepared state", i, layer.Attention.Sinks.DType()) + } + if layer.Experts == nil || layer.Experts.GateUp == nil || layer.Experts.GateUp.Gate == nil || layer.Experts.GateUp.Up == nil || layer.Experts.Down == nil { + t.Fatalf("layer %d Experts = %+v, want loaded expert projections", i, layer.Experts) + } + if got := layer.Experts.GateUp.Gate.Weight.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d GateUp.Gate dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if got := layer.Experts.GateUp.Gate.Bias.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d GateUp.Gate bias dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if got := layer.Experts.GateUp.Up.Weight.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d GateUp.Up dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if got := layer.Experts.GateUp.Up.Bias.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d GateUp.Up bias dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if got := layer.Experts.Down.Weight.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d Down dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if got := layer.Experts.Down.Bias.DType(); got != mlx.DTypeBFloat16 { + t.Fatalf("layer %d Down bias dtype = %v, want %v", i, got, mlx.DTypeBFloat16) + } + if dims := layer.Experts.GateUp.Gate.Weight.Dims(); len(dims) != 3 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.HiddenSize) || dims[2] != int(cfg.IntermediateSize) { + t.Fatalf("layer %d GateUp.Gate dims = %v, want [%d %d %d]", i, dims, cfg.NumLocalExperts, cfg.HiddenSize, cfg.IntermediateSize) + } + if dims := layer.Experts.GateUp.Gate.Bias.Dims(); len(dims) != 2 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.IntermediateSize) { + t.Fatalf("layer %d GateUp.Gate bias dims = %v, want [%d %d]", i, dims, cfg.NumLocalExperts, cfg.IntermediateSize) + } + if dims := layer.Experts.GateUp.Up.Weight.Dims(); len(dims) != 3 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.HiddenSize) || dims[2] != int(cfg.IntermediateSize) { + t.Fatalf("layer %d GateUp.Up dims = %v, want [%d %d %d]", i, dims, cfg.NumLocalExperts, cfg.HiddenSize, cfg.IntermediateSize) + } + if dims := layer.Experts.GateUp.Up.Bias.Dims(); len(dims) != 2 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.IntermediateSize) { + t.Fatalf("layer %d GateUp.Up bias dims = %v, want [%d %d]", i, dims, cfg.NumLocalExperts, cfg.IntermediateSize) + } + if dims := layer.Experts.Down.Weight.Dims(); len(dims) != 3 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.IntermediateSize) || dims[2] != int(cfg.HiddenSize) { + t.Fatalf("layer %d Down dims = %v, want [%d %d %d]", i, dims, cfg.NumLocalExperts, cfg.IntermediateSize, cfg.HiddenSize) + } + if dims := layer.Experts.Down.Bias.Dims(); len(dims) != 2 || dims[0] != int(cfg.NumLocalExperts) || dims[1] != int(cfg.HiddenSize) { + t.Fatalf("layer %d Down bias dims = %v, want [%d %d]", i, dims, cfg.NumLocalExperts, cfg.HiddenSize) + } + } + + caches := m.NewCaches() + if len(caches) != int(cfg.NumHiddenLayers) { + t.Fatalf("len(NewCaches()) = %d, want %d", len(caches), cfg.NumHiddenLayers) + } + if _, ok := caches[0].(*cache.RotatingKVCache); !ok { + t.Fatalf("cache[0] = %T, want *cache.RotatingKVCache", caches[0]) + } + if _, ok := caches[1].(*cache.KVCache); !ok { + t.Fatalf("cache[1] = %T, want *cache.KVCache", caches[1]) + } +} + +func TestLoadWeightsAcceptsAffineQuantizedLinears(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.QuantGroupSize = 64 + cfg.QuantBits = 8 + cfg.QuantMode = "affine" + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + quantizeTensorForTest(t, tensors, "output.weight", 64, 8, "affine") + quantizeTensorForTest(t, tensors, "blocks.0.q_proj.weight", 64, 8, "affine") + quantizeTensorForTest(t, tensors, "blocks.0.k_proj.weight", 64, 8, "affine") + quantizeTensorForTest(t, tensors, "blocks.0.v_proj.weight", 64, 8, "affine") + quantizeTensorForTest(t, tensors, "blocks.0.attn_out.weight", 64, 8, "affine") + quantizeTensorForTest(t, tensors, "blocks.0.router.weight", 64, 8, "affine") + + if err := m.LoadWeights(tensors); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + if m.LMHead == nil { + t.Fatal("LMHead = nil") + } +} + +func TestLoadWeightsRejectsUnsupportedQuantizedDirectExperts(t *testing.T) { + t.Run("affine-int8", func(t *testing.T) { + testLoadWeightsRejectsUnsupportedQuantizedDirectExperts(t, 64, 8, "affine") + }) + t.Run("nvfp4", func(t *testing.T) { + testLoadWeightsRejectsUnsupportedQuantizedDirectExperts(t, 16, 4, "nvfp4") + }) +} + +func testLoadWeightsRejectsUnsupportedQuantizedDirectExperts(t *testing.T, groupSize, bits int, mode string) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.QuantGroupSize = groupSize + cfg.QuantBits = bits + cfg.QuantMode = mode + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + tensors := denseTestTensors(t, cfg) + projection := "blocks.0.experts.gate_proj.weight" + quantizeTensorForTest(t, tensors, projection, groupSize, bits, mode) + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want unsupported direct expert quantization rejection") + } + want := fmt.Sprintf("tensor %q quantization = %s/%d/group%d, want mxfp4/4/group32", projection, mode, bits, groupSize) + if !strings.Contains(err.Error(), want) { + t.Fatalf("LoadWeights() error = %q, want containing %q", err, want) + } +} + +func TestLoadWeightsMissingTensorFails(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + delete(tensors, "blocks.0.q_proj.weight") + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want missing tensor failure") + } + if !strings.Contains(err.Error(), "layer 0") || !strings.Contains(err.Error(), "blocks.0.q_proj.weight") { + t.Fatalf("LoadWeights() error = %q, want layer and tensor name", err) + } +} + +func TestLoadWeightsShapeValidationFails(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + tensors["blocks.0.q_proj.weight"] = mlx.FromValues([]float32{ + 1, 0, 0, 0, + 0, 1, 0, 0, + 0, 0, 1, 0, + }, 3, 4) + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want shape validation failure") + } + if !strings.Contains(err.Error(), "blocks.0.q_proj.weight") || !strings.Contains(err.Error(), "shape [3 4]") { + t.Fatalf("LoadWeights() error = %q, want q_proj shape mismatch", err) + } +} + +func TestLoadWeightsExpertDTypeValidationFails(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + tensors["blocks.0.experts.gate_proj.weight"] = tensors["blocks.0.experts.gate_proj.weight"].AsType(mlx.DTypeFloat32) + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want expert dtype validation failure") + } + if !strings.Contains(err.Error(), "blocks.0.experts.gate_proj.weight") || !strings.Contains(err.Error(), "dtype F32, want BF16") { + t.Fatalf("LoadWeights() error = %q, want expert dtype mismatch", err) + } +} + +func TestLoadWeightsExpertMissingTensorFails(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + delete(tensors, "blocks.0.experts.down_proj.bias") + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want expert missing tensor failure") + } + if !strings.Contains(err.Error(), "layer 0") || !strings.Contains(err.Error(), `"blocks.0.experts.down_proj"`) || !strings.Contains(err.Error(), "incomplete direct expert tensors") { + t.Fatalf("LoadWeights() error = %q, want missing expert tensor", err) + } +} + +func TestLoadWeightsMissingDirectUpExpertFailsClearly(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + delete(tensors, "blocks.0.experts.up_proj.weight") + delete(tensors, "blocks.0.experts.up_proj.bias") + delete(tensors, "blocks.0.experts.gate_up_proj.weight") + delete(tensors, "blocks.0.experts.gate_up_proj.bias") + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want direct up expert failure") + } + if !strings.Contains(err.Error(), "missing direct up expert tensors") || !strings.Contains(err.Error(), "blocks.0.experts.gate_up_proj") { + t.Fatalf("LoadWeights() error = %q, want direct up expert tensor failure", err) + } +} + +func TestLoadWeightsExpertShapeValidationFails(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + tensors["blocks.0.experts.down_proj.weight"] = denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)-1, 99).AsType(mlx.DTypeBFloat16) + + err := m.LoadWeights(tensors) + if err == nil { + t.Fatal("LoadWeights() error = nil, want expert shape validation failure") + } + if !strings.Contains(err.Error(), "blocks.0.experts.down_proj.weight") || !strings.Contains(err.Error(), "[2 4 7]") { + t.Fatalf("LoadWeights() error = %q, want expert shape mismatch", err) + } +} + +func TestSplitGateUpInterleavedUsesEvenOddOrdering(t *testing.T) { + skipIfNoMLX(t) + + dense := mlx.FromValues([]float32{ + 0, 1, + 10, 11, + 20, 21, + 30, 31, + 40, 41, + 50, 51, + }, 6, 2).AsType(mlx.DTypeBFloat16) + bias := mlx.FromValues([]float32{0, 10, 20, 30, 40, 50}, 6).AsType(mlx.DTypeBFloat16) + + gateWeight, upWeight, gateBias, upBias := splitGateUpInterleaved(dense, bias, 3) + if gateWeight == nil || upWeight == nil || gateBias == nil || upBias == nil { + t.Fatal("splitGateUpInterleaved() returned nil tensors") + } + + if dims := gateWeight.Dims(); len(dims) != 2 || dims[0] != 3 || dims[1] != 2 { + t.Fatalf("gateWeight dims = %v, want [3 2]", dims) + } + if dims := upWeight.Dims(); len(dims) != 2 || dims[0] != 3 || dims[1] != 2 { + t.Fatalf("upWeight dims = %v, want [3 2]", dims) + } + + gateWeightVals := materializedFloats(gateWeight.AsType(mlx.DTypeFloat32)) + upWeightVals := materializedFloats(upWeight.AsType(mlx.DTypeFloat32)) + gateBiasVals := materializedFloats(gateBias.AsType(mlx.DTypeFloat32)) + upBiasVals := materializedFloats(upBias.AsType(mlx.DTypeFloat32)) + if got := []float32{gateWeightVals[0], gateWeightVals[2], gateWeightVals[4]}; !slices.Equal(got, []float32{0, 20, 40}) { + t.Fatalf("gateWeight first-column values = %v, want [0 20 40]", got) + } + if got := []float32{upWeightVals[0], upWeightVals[2], upWeightVals[4]}; !slices.Equal(got, []float32{10, 30, 50}) { + t.Fatalf("upWeight first-column values = %v, want [10 30 50]", got) + } + if !slices.Equal(gateBiasVals, []float32{0, 20, 40}) { + t.Fatalf("gateBias values = %v, want [0 20 40]", gateBiasVals) + } + if !slices.Equal(upBiasVals, []float32{10, 30, 50}) { + t.Fatalf("upBias values = %v, want [10 30 50]", upBiasVals) + } +} + +func TestNewCachesLayerParity(t *testing.T) { + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + caches := m.NewCaches() + if len(caches) != 2 { + t.Fatalf("len(NewCaches()) = %d, want 2", len(caches)) + } + if _, ok := caches[0].(*cache.RotatingKVCache); !ok { + t.Fatalf("cache[0] = %T, want *cache.RotatingKVCache", caches[0]) + } + if _, ok := caches[1].(*cache.KVCache); !ok { + t.Fatalf("cache[1] = %T, want *cache.KVCache", caches[1]) + } +} + +func TestRopeParametersDerivedFromConfig(t *testing.T) { + cfg := denseTestConfig(t) + + base, scale, originalContext := cfg.RopeParameters() + if base != cfg.RopeTheta { + t.Fatalf("rope base = %v, want %v", base, cfg.RopeTheta) + } + if scale != 0.5 { + t.Fatalf("rope scale = %v, want 0.5", scale) + } + if originalContext != 4 { + t.Fatalf("original context = %d, want 4", originalContext) + } +} + +func TestBuildGPTOSSRoPEFreqsMatchesReference(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HeadDim = 64 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + gotTensor := testGPTOSSRoPEFreqs(&cfg) + if gotTensor == nil || !gotTensor.Valid() { + t.Fatal("testGPTOSSRoPEFreqs() returned invalid tensor") + } + + got := materializedFloats(gotTensor.AsType(mlx.DTypeFloat32)) + want := referenceGPTOSSRoPEDenominators(&cfg) + if len(got) != len(want) { + t.Fatalf("rope frequency length = %d, want %d", len(got), len(want)) + } + + for i := range want { + tol := 1e-5 * math.Max(1, math.Abs(float64(want[i]))) + if diff := math.Abs(float64(got[i] - want[i])); diff > tol { + t.Fatalf("rope frequency[%d] = %v, want %v (diff %v)", i, got[i], want[i], diff) + } + } +} + +func TestForwardRunsCompletePath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseRuntimeTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + caches := m.NewCaches() + tokens := mlx.FromValues([]int32{1, 2}, 1, 2) + out, aux := m.Forward(testBatch(tokens, 0), caches) + + if out == nil || !out.Valid() { + t.Fatal("Forward() returned invalid output") + } + if aux == nil || !aux.Valid() || aux != out { + t.Fatal("Forward() did not return the output as draft-conditioning state") + } + if dims := out.Dims(); len(dims) != 3 || dims[0] != 1 || dims[1] != 2 || dims[2] != int(cfg.HiddenSize) { + t.Fatalf("Forward() dims = %v, want [1 2 %d]", dims, cfg.HiddenSize) + } +} + +func TestForwardLastTokenMatchesPrefillStepPath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + prepareSyntheticF32Sinks(t, m) + + tokens := mlx.FromValues([]int32{1, 2, 3, 4, 5, 6}, 1, 6) + + fullCaches := m.NewCaches() + fullHidden := forwardModel(m, tokens, fullCaches) + fullLogits := m.Unembed(fullHidden) + fullLast := materializedFloats(fullLogits.Slice(mlx.Slice(), mlx.Slice(fullLogits.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepCaches := m.NewCaches() + forwardModel(m, mlx.FromValues([]int32{1, 2, 3, 4, 5}, 1, 5), stepCaches) + stepHidden := forwardModel(m, mlx.FromValues([]int32{6}, 1, 1), stepCaches) + stepLogits := m.Unembed(stepHidden) + stepLast := materializedFloats(stepLogits.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("logit length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + // Batched prefill has a known ~0.04% divergence from stepwise on MLX + // due to different Metal kernel dispatch. This is well within quantization + // tolerance for MXFP4/int8 production models. + for i := range fullLast { + diff := math.Abs(float64(fullLast[i] - stepLast[i])) + mag := math.Max(math.Abs(float64(fullLast[i])), 1e-6) + if diff/mag > 5e-4 { + t.Fatalf("last-token logit[%d] = %v, want %v (diff %v, rel %v)", i, stepLast[i], fullLast[i], diff, diff/mag) + } + } +} + +func TestLayerLastTokenMatchesPrefillStepPath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + prepareSyntheticF32Sinks(t, m) + layer := m.Layers[0] + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + + fullCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + full := layer.Forward(mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16), fullCache, 1, 6, &cfg, 0) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + layer.Forward(mlx.FromValues(xVals[:5*64], 1, 5, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 5, &cfg, 0) + step := layer.Forward(mlx.FromValues(xVals[5*64:], 1, 1, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 1, &cfg, 0) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("layer output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + diff := math.Abs(float64(fullLast[i] - stepLast[i])) + mag := math.Max(math.Abs(float64(fullLast[i])), 1e-6) + if diff/mag > 1e-3 { + t.Fatalf("layer last-token output[%d] = %v, want %v (diff %v, rel %v)", i, stepLast[i], fullLast[i], diff, diff/mag) + } + } +} + +func TestUnembedLastTokenMatchesPrefillStepPath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + + hVals := make([]float32, 6*64) + for i := range hVals { + hVals[i] = float32((i%23)-11) / 11 + } + + full := m.Unembed(mlx.FromValues(hVals, 1, 6, 64).AsType(mlx.DTypeBFloat16)) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + step := m.Unembed(mlx.FromValues(hVals[5*64:], 1, 1, 64).AsType(mlx.DTypeBFloat16)) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("unembed output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + if diff := math.Abs(float64(fullLast[i] - stepLast[i])); diff > 1e-2 { + t.Fatalf("unembed last-token output[%d] = %v, want %v (diff %v)", i, stepLast[i], fullLast[i], diff) + } + } +} + +func TestAttentionLastTokenMatchesPrefillStepPathLoadedLayer(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + prepareSyntheticF32Sinks(t, m) + attn := m.Layers[0].Attention + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + + fullCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + full := attn.Forward(mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16), fullCache, 1, 6, &cfg, 0) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + attn.Forward(mlx.FromValues(xVals[:5*64], 1, 5, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 5, &cfg, 0) + step := attn.Forward(mlx.FromValues(xVals[5*64:], 1, 1, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 1, &cfg, 0) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("attention output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + diff := math.Abs(float64(fullLast[i] - stepLast[i])) + mag := math.Max(math.Abs(float64(fullLast[i])), 1e-6) + if diff/mag > 1e-3 { + t.Fatalf("loaded attention last-token output[%d] = %v, want %v (diff %v, rel %v)", i, stepLast[i], fullLast[i], diff, diff/mag) + } + } +} + +func TestAttentionLastTokenMatchesPrefillStepPathLoadedLayerCausalCache(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + prepareSyntheticF32Sinks(t, m) + attn := m.Layers[0].Attention + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + + fullCache := cache.NewKVCache() + full := attn.Forward(mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16), fullCache, 1, 6, &cfg, 0) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepCache := cache.NewKVCache() + attn.Forward(mlx.FromValues(xVals[:5*64], 1, 5, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 5, &cfg, 0) + step := attn.Forward(mlx.FromValues(xVals[5*64:], 1, 1, 64).AsType(mlx.DTypeBFloat16), stepCache, 1, 1, &cfg, 0) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("causal attention output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + diff := math.Abs(float64(fullLast[i] - stepLast[i])) + mag := math.Max(math.Abs(float64(fullLast[i])), 1e-6) + if diff/mag > 1e-3 { + t.Fatalf("loaded causal attention last-token output[%d] = %v, want %v (diff %v, rel %v)", i, stepLast[i], fullLast[i], diff, diff/mag) + } + } +} + +func TestQuantizedAttentionCachedPrefillParity(t *testing.T) { + skipIfNoMLX(t) + + // Keep MLX graph construction and evaluation on this locked goroutine. + // MLX default streams are thread-local, so nested t.Run subtests can move + // reused lazy arrays/caches onto a different stream. + for _, tc := range quantizedCachedPrefillParityCases() { + m, cfg := quantizedAttentionTestModelWithQuantization(t, tc.groupSize, tc.bits, tc.mode) + attn := m.Layers[0].Attention + seqLen := 48 + xVals := patternedHiddenValues(seqLen, int(cfg.HiddenSize)) + x := mlx.FromValues(xVals, 1, seqLen, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + + assertQuantizedAttentionCachedPrefillParity(t, tc.name+"/rotating-cache", attn, x, cache.NewRotatingKVCache(int(cfg.SlidingWindow)), cache.NewRotatingKVCache(int(cfg.SlidingWindow)), seqLen, &cfg, tc) + assertQuantizedAttentionCachedPrefillParity(t, tc.name+"/causal-cache", attn, x, cache.NewKVCache(), cache.NewKVCache(), seqLen, &cfg, tc) + } +} + +func assertQuantizedAttentionCachedPrefillParity( + t *testing.T, + label string, + attn *Attention, + x *mlx.Array, + fullCache, stepCache cache.Cache, + seqLen int, + cfg *Config, + tc quantizedCachedPrefillParityCase, +) { + t.Helper() + + _, fullTrace := attentionForwardTraceForTest(t, attn, x, fullCache, 1, seqLen, cfg, 0) + + var stepTrace attentionPrefillTrace + for pos := range seqLen { + _, stepTrace = attentionForwardTraceForTest(t, attn, row11SliceSequence(x, pos), stepCache, 1, 1, cfg, 0) + } + + assertFloatSliceClose(t, label+" q_proj last-token", fullTrace.qProjLast, stepTrace.qProjLast, tc.projectionTol, tc.projectionTol) + assertFloatSliceClose(t, label+" k_proj last-token", fullTrace.kProjLast, stepTrace.kProjLast, tc.projectionTol, tc.projectionTol) + assertFloatSliceClose(t, label+" v_proj last-token", fullTrace.vProjLast, stepTrace.vProjLast, tc.projectionTol, tc.projectionTol) + if fullTrace.visibleKeyLen != stepTrace.visibleKeyLen || fullTrace.visibleValueLen != stepTrace.visibleValueLen { + t.Fatalf( + "%s visible cache lens = (%d, %d), want (%d, %d)", + label, + fullTrace.visibleKeyLen, + fullTrace.visibleValueLen, + stepTrace.visibleKeyLen, + stepTrace.visibleValueLen, + ) + } + assertFloatSliceClose(t, label+" visible cache last-key", fullTrace.visibleKeyLast, stepTrace.visibleKeyLast, tc.projectionTol, tc.projectionTol) + assertFloatSliceClose(t, label+" visible cache last-value", fullTrace.visibleValueLast, stepTrace.visibleValueLast, tc.projectionTol, tc.projectionTol) + assertFloatSliceClose(t, label+" pre-o-proj attention", fullTrace.preOProjLast, stepTrace.preOProjLast, tc.attentionTol, tc.attentionTol) + assertFloatSliceClose(t, label+" attention output", fullTrace.outputLast, stepTrace.outputLast, tc.attentionTol, tc.attentionTol) +} + +func TestQuantizedLayerLastTokenMatchesCachedPrefillStepPath(t *testing.T) { + skipIfNoMLX(t) + + m, cfg := quantizedAttentionTestModel(t) + layer := m.Layers[0] + seqLen := 48 + xVals := patternedHiddenValues(seqLen, int(cfg.HiddenSize)) + x := mlx.FromValues(xVals, 1, seqLen, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + + fullCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + full := layer.Forward(x, fullCache, 1, seqLen, &cfg, 0) + fullLast := lastTokenFloats(full.AsType(mlx.DTypeFloat32)) + + stepCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + var step *mlx.Array + for pos := range seqLen { + step = layer.Forward(row11SliceSequence(x, pos), stepCache, 1, 1, &cfg, 0) + } + stepLast := lastTokenFloats(step.AsType(mlx.DTypeFloat32)) + + assertFloatSliceClose(t, "quantized layer output", fullLast, stepLast, 5e-2, 5e-2) +} + +func TestQuantizedModelLastTokenLogitsMatchCachedPrefillStepPath(t *testing.T) { + skipIfNoMLX(t) + + m, cfg := quantizedAttentionTestModel(t) + seqLen := 48 + tokenVals := patternedTokenValues(seqLen, int(cfg.VocabSize)) + tokens := mlx.FromValues(tokenVals, 1, seqLen) + + fullCaches := m.NewCaches() + fullHidden := forwardModel(m, tokens, fullCaches) + fullLogits := m.Unembed(fullHidden) + fullLast := lastTokenFloats(fullLogits.AsType(mlx.DTypeFloat32)) + + stepCaches := m.NewCaches() + var stepHidden *mlx.Array + for pos := range seqLen { + stepHidden = forwardModel(m, mlx.FromValues(tokenVals[pos:pos+1], 1, 1), stepCaches) + } + stepLogits := m.Unembed(stepHidden) + stepLast := lastTokenFloats(stepLogits.AsType(mlx.DTypeFloat32)) + + assertFloatSliceClose(t, "quantized model logits", fullLast, stepLast, 5e-2, 5e-2) +} + +func TestBatchedQuantizedModelCachedPrefillParity(t *testing.T) { + skipIfNoMLX(t) + + // Keep this loop on one locked goroutine; the MLX caches and lazy arrays + // are evaluated after construction and depend on thread-local streams. + for _, tc := range quantizedCachedPrefillParityCases() { + m, cfg := quantizedAttentionTestModelWithQuantization(t, tc.groupSize, tc.bits, tc.mode) + seqLen := 48 + tokenVals := patternedTokenValues(seqLen, int(cfg.VocabSize)) + tokens := mlx.FromValues(tokenVals, 1, seqLen) + + fullCaches := m.NewCaches() + fullHidden := forwardModel(m, tokens, fullCaches) + fullLogits := m.Unembed(fullHidden) + fullLast := lastTokenFloats(fullLogits.AsType(mlx.DTypeFloat32)) + + stepCaches := m.NewCaches() + var stepHidden *mlx.Array + for pos := range seqLen { + stepHidden = forwardModel(m, mlx.FromValues(tokenVals[pos:pos+1], 1, 1), stepCaches) + } + stepLogits := m.Unembed(stepHidden) + stepLast := lastTokenFloats(stepLogits.AsType(mlx.DTypeFloat32)) + + assertFloatSliceClose(t, tc.name+" batched quantized model logits", fullLast, stepLast, tc.modelTol, tc.modelTol) + } +} + +func TestQProjLastTokenMatchesBatchPathLoadedLayer(t *testing.T) { + skipIfNoMLX(t) + t.Skip("diagnostic only: MLX batched affine path diverges on macOS; gpt-oss runtime avoids this path") + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + + full := qproj.Forward(mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16)) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + step := qproj.Forward(mlx.FromValues(xVals[5*64:], 1, 1, 64).AsType(mlx.DTypeBFloat16)) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("q_proj output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + if diff := math.Abs(float64(fullLast[i] - stepLast[i])); diff > 1e-3 { + t.Fatalf("q_proj last-token output[%d] = %v, want %v (diff %v)", i, stepLast[i], fullLast[i], diff) + } + } +} + +func TestQProjLoadedLayerMatchesExplicitAffine(t *testing.T) { + skipIfNoMLX(t) + t.Skip("diagnostic only: MLX batched affine path diverges on macOS; gpt-oss runtime avoids this path") + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj.(*nn.Linear) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + x := mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16) + + got := qproj.Forward(x) + want := x.Matmul(qproj.Weight.Transpose(1, 0)).Add(qproj.Bias) + + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + wantVals := materializedFloats(want.AsType(mlx.DTypeFloat32)) + if len(gotVals) != len(wantVals) { + t.Fatalf("q_proj explicit affine length mismatch: got=%d want=%d", len(gotVals), len(wantVals)) + } + for i := range gotVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 1e-3 { + t.Fatalf("q_proj explicit affine output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestQProjExplicitAffineMatchesCPUReference(t *testing.T) { + skipIfNoMLX(t) + t.Skip("diagnostic only: MLX batched affine path diverges on macOS; gpt-oss runtime avoids this path") + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj.(*nn.Linear) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + lastX := xVals[5*64:] + + weightVals := materializedFloats(qproj.Weight.AsType(mlx.DTypeFloat32)) + biasVals := materializedFloats(qproj.Bias.AsType(mlx.DTypeFloat32)) + w := make([][]float32, 64) + for in := range 64 { + w[in] = make([]float32, 64) + for out := range 64 { + w[in][out] = weightVals[out*64+in] + } + } + want := affineRef(lastX, w, biasVals) + + fullX := mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeBFloat16) + full := fullX.Matmul(qproj.Weight.Transpose(1, 0)).Add(qproj.Bias) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepX := mlx.FromValues(lastX, 1, 1, 64).AsType(mlx.DTypeBFloat16) + step := stepX.Matmul(qproj.Weight.Transpose(1, 0)).Add(qproj.Bias) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + for i := range want { + if diff := math.Abs(float64(fullLast[i] - want[i])); diff > 1e-2 { + t.Fatalf("explicit affine batch output[%d] = %v, want %v (diff %v)", i, fullLast[i], want[i], diff) + } + if diff := math.Abs(float64(stepLast[i] - want[i])); diff > 1e-2 { + t.Fatalf("explicit affine step output[%d] = %v, want %v (diff %v)", i, stepLast[i], want[i], diff) + } + } +} + +func TestQProjExplicitAffineFloat32MatchesCPUReference(t *testing.T) { + skipIfNoMLX(t) + t.Skip("diagnostic only: MLX batched affine path diverges on macOS; gpt-oss runtime avoids this path") + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj.(*nn.Linear) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + lastX := xVals[5*64:] + + weightVals := materializedFloats(qproj.Weight.AsType(mlx.DTypeFloat32)) + biasVals := materializedFloats(qproj.Bias.AsType(mlx.DTypeFloat32)) + w := make([][]float32, 64) + for in := range 64 { + w[in] = make([]float32, 64) + for out := range 64 { + w[in][out] = weightVals[out*64+in] + } + } + want := affineRef(lastX, w, biasVals) + + fullX := mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeFloat32) + fullW := qproj.Weight.Transpose(1, 0).AsType(mlx.DTypeFloat32) + fullB := qproj.Bias.AsType(mlx.DTypeFloat32) + full := fullX.Matmul(fullW).Add(fullB) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1)) + + stepX := mlx.FromValues(lastX, 1, 1, 64).AsType(mlx.DTypeFloat32) + step := stepX.Matmul(fullW).Add(fullB) + stepLast := materializedFloats(step.Squeeze(1)) + + for i := range want { + if diff := math.Abs(float64(fullLast[i] - want[i])); diff > 1e-3 { + t.Fatalf("float32 affine batch output[%d] = %v, want %v (diff %v)", i, fullLast[i], want[i], diff) + } + if diff := math.Abs(float64(stepLast[i] - want[i])); diff > 1e-3 { + t.Fatalf("float32 affine step output[%d] = %v, want %v (diff %v)", i, stepLast[i], want[i], diff) + } + } +} + +func TestQProjExplicitAffine2DAndContiguousMatchCPUReference(t *testing.T) { + skipIfNoMLX(t) + t.Skip("diagnostic only: MLX batched affine path diverges on macOS; gpt-oss runtime avoids this path") + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj.(*nn.Linear) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + lastX := xVals[5*64:] + + weightVals := materializedFloats(qproj.Weight.AsType(mlx.DTypeFloat32)) + biasVals := materializedFloats(qproj.Bias.AsType(mlx.DTypeFloat32)) + w := make([][]float32, 64) + for in := range 64 { + w[in] = make([]float32, 64) + for out := range 64 { + w[in][out] = weightVals[out*64+in] + } + } + want := affineRef(lastX, w, biasVals) + + wT := qproj.Weight.Transpose(1, 0).AsType(mlx.DTypeFloat32) + bias := qproj.Bias.AsType(mlx.DTypeFloat32) + + full3D := mlx.FromValues(xVals, 1, 6, 64).AsType(mlx.DTypeFloat32) + out3D := full3D.Matmul(wT).Add(bias) + last3D := materializedFloats(out3D.Slice(mlx.Slice(), mlx.Slice(out3D.Dim(1)-1), mlx.Slice()).Squeeze(1)) + + full2D := mlx.FromValues(xVals, 6, 64).AsType(mlx.DTypeFloat32) + out2D := full2D.Matmul(wT).Add(bias) + last2D := materializedFloats(out2D.Slice(mlx.Slice(out2D.Dim(0)-1), mlx.Slice()).Squeeze(0)) + + full2DContig := mlx.Contiguous(full2D, false) + wTContig := mlx.Contiguous(wT, false) + out2DContig := full2DContig.Matmul(wTContig).Add(bias) + last2DContig := materializedFloats(out2DContig.Slice(mlx.Slice(out2DContig.Dim(0)-1), mlx.Slice()).Squeeze(0)) + + step2D := mlx.FromValues(lastX, 1, 64).AsType(mlx.DTypeFloat32) + outStep2D := step2D.Matmul(wT).Add(bias) + lastStep2D := materializedFloats(outStep2D.Squeeze(0)) + + check := func(label string, got []float32) { + t.Helper() + for i := range want { + if diff := math.Abs(float64(got[i] - want[i])); diff > 1e-3 { + t.Errorf("%s output[%d] = %v, want %v (diff %v)", label, i, got[i], want[i], diff) + return + } + } + } + + check("3D affine", last3D) + check("2D affine", last2D) + check("2D contiguous affine", last2DContig) + check("2D step affine", lastStep2D) + if t.Failed() { + t.FailNow() + } +} + +func TestQProjExplicitAffineBatch2DRowWiseBackendBehavior(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + if err := m.LoadWeights(denseTestTensors(t, cfg)); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + qproj := m.Layers[0].Attention.QProj.(*nn.Linear) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + + wT := qproj.Weight.Transpose(1, 0).AsType(mlx.DTypeFloat32) + bias := qproj.Bias.AsType(mlx.DTypeFloat32) + + full2D := mlx.FromValues(xVals, 6, 64).AsType(mlx.DTypeFloat32) + batched := full2D.Matmul(wT).Add(bias) + batchedVals := materializedFloats(batched) + + rowWiseParts := make([]*mlx.Array, 0, 6) + for row := range 6 { + xRow := mlx.FromValues(xVals[row*64:(row+1)*64], 1, 64).AsType(mlx.DTypeFloat32) + rowWiseParts = append(rowWiseParts, xRow.Matmul(wT).Add(bias)) + } + rowWise := mlx.Concatenate(rowWiseParts, 0) + rowWiseVals := materializedFloats(rowWise) + + if len(batchedVals) != len(rowWiseVals) { + t.Fatalf("batched vs row-wise length mismatch: got=%d want=%d", len(batchedVals), len(rowWiseVals)) + } + + diffCount := 0 + maxDiff := float64(0) + maxDiffIndex := -1 + for i := range batchedVals { + if diff := math.Abs(float64(batchedVals[i] - rowWiseVals[i])); diff > 1e-3 { + diffCount++ + if diff > maxDiff { + maxDiff = diff + maxDiffIndex = i + } + } + } + if diffCount == 0 { + t.Logf("batched 2D affine matches row-wise 2D affine within tolerance; MLX backend does not reproduce the historical divergence (max_diff=%g)", maxDiff) + return + } + t.Logf("batched 2D affine diverges from row-wise 2D affine; historical MLX backend issue still reproduces (diff_count=%d max_diff=%g max_diff_index=%d)", diffCount, maxDiff, maxDiffIndex) +} + +func TestAttentionLastTokenMatchesPrefillStepPathScaledRoPE(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.HeadDim = 64 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + xVals := make([]float32, 6*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + x := mlx.FromValues(xVals, 1, 6, 64) + + attn := &Attention{ + QProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + KProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + VProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + OProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + Sinks: preparedSyntheticSink(t, []float32{0.3}, mlx.DTypeFloat32), + RoPEFreqs: testGPTOSSRoPEFreqs(&cfg), + RoPEScale: testGPTOSSYarnScale(&cfg), + } + + fullCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + full := attn.Forward(x, fullCache, 1, 6, &cfg, 0) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + stepCache := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + attn.Forward(mlx.FromValues(xVals[:5*64], 1, 5, 64), stepCache, 1, 5, &cfg, 0) + step := attn.Forward(mlx.FromValues(xVals[5*64:], 1, 1, 64), stepCache, 1, 1, &cfg, 0) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("attention output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + if diff := math.Abs(float64(fullLast[i] - stepLast[i])); diff > 3e-2 { + t.Fatalf("attention last-token output[%d] = %v, want %v (diff %v)", i, stepLast[i], fullLast[i], diff) + } + } +} + +func TestExpertsLastTokenMatchesBatchPath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + + xVals := make([]float32, 6*64) + routerVals := make([]float32, 6*int(cfg.NumLocalExperts)) + for i := range xVals { + xVals[i] = float32((i%19)-9) / 9 + } + for i := range routerVals { + routerVals[i] = float32((i%11)-5) / 4 + } + + experts := &Experts{ + GateUp: &ExpertPair{ + Gate: &ExpertProjection{ + Weight: denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize), 1).AsType(mlx.DTypeBFloat16), + Bias: expertBias(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), 0).AsType(mlx.DTypeBFloat16), + }, + Up: &ExpertProjection{ + Weight: denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize), 10).AsType(mlx.DTypeBFloat16), + Bias: expertBias(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), 0.5).AsType(mlx.DTypeBFloat16), + }, + }, + Down: &ExpertProjection{ + Weight: denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize), 20).AsType(mlx.DTypeBFloat16), + Bias: expertBias(int(cfg.NumLocalExperts), int(cfg.HiddenSize), 1).AsType(mlx.DTypeBFloat16), + }, + } + + full := experts.Forward( + mlx.FromValues(xVals, 1, 6, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + mlx.FromValues(routerVals, 1, 6, int(cfg.NumLocalExperts)).AsType(mlx.DTypeBFloat16), + &cfg, + 0, + ) + fullLast := materializedFloats(full.Slice(mlx.Slice(), mlx.Slice(full.Dim(1)-1), mlx.Slice()).Squeeze(1).AsType(mlx.DTypeFloat32)) + + step := experts.Forward( + mlx.FromValues(xVals[5*64:], 1, 1, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + mlx.FromValues(routerVals[5*int(cfg.NumLocalExperts):], 1, 1, int(cfg.NumLocalExperts)).AsType(mlx.DTypeBFloat16), + &cfg, + 0, + ) + stepLast := materializedFloats(step.Squeeze(1).AsType(mlx.DTypeFloat32)) + + if len(fullLast) != len(stepLast) { + t.Fatalf("expert output length mismatch: full=%d step=%d", len(fullLast), len(stepLast)) + } + for i := range fullLast { + if diff := math.Abs(float64(fullLast[i] - stepLast[i])); diff > 5e-2 { + t.Fatalf("expert last-token output[%d] = %v, want %v (diff %v)", i, stepLast[i], fullLast[i], diff) + } + } +} + +func TestExpertsForwardMatchesReference(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.IntermediateSize = 2 + cfg.NumLocalExperts = 2 + cfg.NumExpertsPerTok = 2 + + xVals := []float32{0.5, -1.0} + routerVals := []float32{2.0, 1.0} + x := mlx.FromValues(xVals, 1, 1, int(cfg.HiddenSize)) + router := mlx.FromValues(routerVals, 1, 1, int(cfg.NumLocalExperts)) + + experts := &Experts{ + GateUp: &ExpertPair{ + Gate: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 1.0, 0.0, + 0.0, 1.0, + + -0.5, 1.0, + 1.5, -1.0, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.1, -0.2, + 0.3, 0.4, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + Up: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 2.0, 0.0, + 0.0, 2.0, + + 1.0, -1.0, + 0.5, 1.5, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.3, 0.4, + -0.2, 0.1, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + }, + Down: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 1.0, 0.0, + 0.0, 1.0, + + 0.5, -1.0, + 1.0, 0.5, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.05, -0.05, + -0.1, 0.2, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + }, + } + + got := experts.Forward(x, router, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatal("Experts.Forward() returned invalid tensor") + } + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + + wantVals := referenceExpertsForward( + xVals, + routerVals, + int(cfg.NumExpertsPerTok), + [][][]float32{ + {{1.0, 0.0}, {0.0, 1.0}}, + {{-0.5, 1.0}, {1.5, -1.0}}, + }, + [][]float32{ + {0.1, -0.2}, + {0.3, 0.4}, + }, + [][][]float32{ + {{2.0, 0.0}, {0.0, 2.0}}, + {{1.0, -1.0}, {0.5, 1.5}}, + }, + [][]float32{ + {0.3, 0.4}, + {-0.2, 0.1}, + }, + [][][]float32{ + {{1.0, 0.0}, {0.0, 1.0}}, + {{0.5, -1.0}, {1.0, 0.5}}, + }, + [][]float32{ + {0.05, -0.05}, + {-0.1, 0.2}, + }, + ) + + if len(gotVals) != len(wantVals) { + t.Fatalf("Experts.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 1e-2 { + t.Fatalf("Experts.Forward() output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestExpertsForwardMatchesReferenceSortedPath(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.IntermediateSize = 2 + cfg.NumLocalExperts = 3 + cfg.NumExpertsPerTok = 2 + + const seqLen = 64 + xVals := make([]float32, seqLen*int(cfg.HiddenSize)) + routerVals := make([]float32, seqLen*int(cfg.NumLocalExperts)) + for i := range seqLen { + xVals[i*2+0] = float32(i%7)/7 + 0.1 + xVals[i*2+1] = float32((i*3)%11)/11 - 0.2 + + routerVals[i*3+0] = float32((i%5)-2) * 0.4 + routerVals[i*3+1] = float32(((i+2)%7)-3) * 0.3 + routerVals[i*3+2] = float32(((i*2)%9)-4) * 0.2 + } + + x := mlx.FromValues(xVals, 1, seqLen, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + router := mlx.FromValues(routerVals, 1, seqLen, int(cfg.NumLocalExperts)).AsType(mlx.DTypeBFloat16) + + experts := &Experts{ + GateUp: &ExpertPair{ + Gate: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 1.0, 0.0, + 0.2, 0.8, + + -0.4, 0.9, + 1.1, -0.3, + + 0.7, -0.2, + 0.5, 1.2, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.1, -0.2, + 0.0, 0.05, + -0.1, 0.2, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + Up: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 0.9, 0.1, + -0.3, 1.1, + + 1.4, -0.6, + 0.2, 0.7, + + -0.5, 0.8, + 1.0, 0.4, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.0, 0.15, + -0.05, 0.1, + 0.2, -0.1, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + }, + Down: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 0.8, -0.2, + 0.1, 1.0, + + 1.2, 0.3, + -0.6, 0.9, + + 0.4, 1.1, + 0.7, -0.5, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.05, -0.05, + 0.1, 0.0, + -0.15, 0.2, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + }, + } + + got := experts.Forward(x, router, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatal("Experts.Forward() returned invalid tensor") + } + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + + wantVals := make([]float32, 0, len(gotVals)) + for i := range seqLen { + wantVals = append(wantVals, referenceExpertsForward( + xVals[i*2:(i+1)*2], + routerVals[i*3:(i+1)*3], + int(cfg.NumExpertsPerTok), + [][][]float32{ + {{1.0, 0.0}, {0.2, 0.8}}, + {{-0.4, 0.9}, {1.1, -0.3}}, + {{0.7, -0.2}, {0.5, 1.2}}, + }, + [][]float32{ + {0.1, -0.2}, + {0.0, 0.05}, + {-0.1, 0.2}, + }, + [][][]float32{ + {{0.9, 0.1}, {-0.3, 1.1}}, + {{1.4, -0.6}, {0.2, 0.7}}, + {{-0.5, 0.8}, {1.0, 0.4}}, + }, + [][]float32{ + {0.0, 0.15}, + {-0.05, 0.1}, + {0.2, -0.1}, + }, + [][][]float32{ + {{0.8, -0.2}, {0.1, 1.0}}, + {{1.2, 0.3}, {-0.6, 0.9}}, + {{0.4, 1.1}, {0.7, -0.5}}, + }, + [][]float32{ + {0.05, -0.05}, + {0.1, 0.0}, + {-0.15, 0.2}, + }, + )...) + } + + if len(gotVals) != len(wantVals) { + t.Fatalf("Experts.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 2e-2 { + t.Fatalf("Experts.Forward() sorted output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestExpertsForwardMatchesReferencePromptLength63(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.IntermediateSize = 2 + cfg.NumLocalExperts = 3 + cfg.NumExpertsPerTok = 2 + + const seqLen = 63 + xVals := make([]float32, seqLen*int(cfg.HiddenSize)) + routerVals := make([]float32, seqLen*int(cfg.NumLocalExperts)) + for i := range seqLen { + xVals[i*2+0] = float32(i%7)/7 + 0.1 + xVals[i*2+1] = float32((i*3)%11)/11 - 0.2 + + routerVals[i*3+0] = float32((i%5)-2) * 0.4 + routerVals[i*3+1] = float32(((i+2)%7)-3) * 0.3 + routerVals[i*3+2] = float32(((i*2)%9)-4) * 0.2 + } + + x := mlx.FromValues(xVals, 1, seqLen, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + router := mlx.FromValues(routerVals, 1, seqLen, int(cfg.NumLocalExperts)).AsType(mlx.DTypeBFloat16) + + experts := &Experts{ + GateUp: &ExpertPair{ + Gate: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 1.0, 0.0, + 0.2, 0.8, + + -0.4, 0.9, + 1.1, -0.3, + + 0.7, -0.2, + 0.5, 1.2, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.1, -0.2, + 0.0, 0.05, + -0.1, 0.2, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + Up: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 0.9, 0.1, + -0.3, 1.1, + + 1.4, -0.6, + 0.2, 0.7, + + -0.5, 0.8, + 1.0, 0.4, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.0, 0.15, + -0.05, 0.1, + 0.2, -0.1, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize)).AsType(mlx.DTypeBFloat16), + }, + }, + Down: &ExpertProjection{ + Weight: mlx.FromValues([]float32{ + 0.8, -0.2, + 0.1, 1.0, + + 1.2, 0.3, + -0.6, 0.9, + + 0.4, 1.1, + 0.7, -0.5, + }, int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + Bias: mlx.FromValues([]float32{ + 0.05, -0.05, + 0.1, 0.0, + -0.15, 0.2, + }, int(cfg.NumLocalExperts), int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16), + }, + } + + got := experts.Forward(x, router, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatal("Experts.Forward() returned invalid tensor") + } + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + + wantVals := make([]float32, 0, len(gotVals)) + for i := range seqLen { + wantVals = append(wantVals, referenceExpertsForward( + xVals[i*2:(i+1)*2], + routerVals[i*3:(i+1)*3], + int(cfg.NumExpertsPerTok), + [][][]float32{ + {{1.0, 0.0}, {0.2, 0.8}}, + {{-0.4, 0.9}, {1.1, -0.3}}, + {{0.7, -0.2}, {0.5, 1.2}}, + }, + [][]float32{ + {0.1, -0.2}, + {0.0, 0.05}, + {-0.1, 0.2}, + }, + [][][]float32{ + {{0.9, 0.1}, {-0.3, 1.1}}, + {{1.4, -0.6}, {0.2, 0.7}}, + {{-0.5, 0.8}, {1.0, 0.4}}, + }, + [][]float32{ + {0.0, 0.15}, + {-0.05, 0.1}, + {0.2, -0.1}, + }, + [][][]float32{ + {{0.8, -0.2}, {0.1, 1.0}}, + {{1.2, 0.3}, {-0.6, 0.9}}, + {{0.4, 1.1}, {0.7, -0.5}}, + }, + [][]float32{ + {0.05, -0.05}, + {0.1, 0.0}, + {-0.15, 0.2}, + }, + )...) + } + + if len(gotVals) != len(wantVals) { + t.Fatalf("Experts.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 2e-2 { + t.Fatalf("Experts.Forward() prompt-length-63 output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestAttentionForwardMatchesReference(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.HeadDim = 2 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.RopeTheta = 10000 + cfg.RopeScaling.Factor = 1 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4 + refreshRow11CheckedConfig(t, &cfg) + + xVals := []float32{ + 1.0, 0.5, + -0.25, 0.75, + } + x := mlx.FromValues(xVals, 1, 2, 2) + + attn := &Attention{ + QProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, nil, 2, 2), + KProj: nnLinearFromValues([]float32{ + 0.5, 0, + 0, 1.5, + }, []float32{0.1, -0.2}, 2, 2), + VProj: nnLinearFromValues([]float32{ + 1.2, 0, + 0, 0.8, + }, []float32{-0.05, 0.2}, 2, 2), + OProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, []float32{0.01, -0.02}, 2, 2), + Sinks: preparedSyntheticSink(t, []float32{row11NoSinkSentinel}, mlx.DTypeFloat32), + } + + got := attn.Forward(x, nil, 1, 2, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatal("Attention.Forward() returned invalid tensor") + } + + want := referenceAttentionForward( + xVals, + [][]float32{{1, 0}, {0, 1}}, + nil, + [][]float32{{0.5, 0}, {0, 1.5}}, + []float32{0.1, -0.2}, + [][]float32{{1.2, 0}, {0, 0.8}}, + []float32{-0.05, 0.2}, + [][]float32{{1, 0}, {0, 1}}, + []float32{0.01, -0.02}, + cfg.RopeTheta, + ) + + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + if len(gotVals) != len(want) { + t.Fatalf("Attention.Forward() output length = %d, want %d", len(gotVals), len(want)) + } + for i := range want { + if diff := math.Abs(float64(gotVals[i] - want[i])); diff > 2e-1 { + t.Fatalf("Attention.Forward() output[%d] = %v, want %v (diff %v)", i, gotVals[i], want[i], diff) + } + } +} + +func TestAttentionForwardMatchesReferenceWithSinksAndScaledRoPE(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.HeadDim = 64 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + refreshRow11CheckedConfig(t, &cfg) + + xVals := make([]float32, 2*64) + for i := range xVals { + xVals[i] = float32((i%17)-8) / 8 + } + x := mlx.FromValues(xVals, 1, 2, 64) + + attn := &Attention{ + QProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + KProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + VProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + OProj: nnLinearFromValues(identityFlat(64), nil, 64, 64), + Sinks: preparedSyntheticSink(t, []float32{0.3}, mlx.DTypeFloat32), + RoPEFreqs: testGPTOSSRoPEFreqs(&cfg), + RoPEScale: testGPTOSSYarnScale(&cfg), + } + + got := attn.Forward(x, nil, 1, 2, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatal("Attention.Forward() returned invalid tensor") + } + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + + wantVals := referenceAttentionForwardWithSinksScaledRoPE(xVals, &cfg, 0.3) + if len(gotVals) != len(wantVals) { + t.Fatalf("Attention.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 3e-2 { + t.Fatalf("Attention.Forward() scaled-rope+sinks output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestExpertsForwardMatchesExplicitRealImportedModel(t *testing.T) { + skipIfNoMLX(t) + + if os.Getenv("OLLAMA_MODELS") == "" { + t.Skip("OLLAMA_MODELS not set") + } + + root, err := model.Open("gptoss-mlx-runtime") + if err != nil { + t.Skipf("imported model not available: %v", err) + } + defer root.Close() + + baseModel, err := base.New(root) + if err != nil { + t.Fatalf("base.New() error = %v", err) + } + tensors, err := loadRuntimeTensorsForTest(root) + if err != nil { + t.Fatalf("loadRuntimeTensorsForTest() error = %v", err) + } + if err := base.Weights(baseModel)(tensors); err != nil { + t.Fatalf("load weights error = %v", err) + } + + m := baseModel.(*Model) + layer := m.Layers[0] + cfg := m.Config + + xVals := make([]float32, int(cfg.HiddenSize)) + for i := range xVals { + xVals[i] = float32((i%11)-5) / 5 + } + x := mlx.FromValues(xVals, 1, 1, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + + routerVals := make([]float32, int(cfg.NumLocalExperts)) + for i := range routerVals { + routerVals[i] = float32(((i*7)%17)-8) / 3 + } + router := mlx.FromValues(routerVals, 1, 1, int(cfg.NumLocalExperts)).AsType(mlx.DTypeBFloat16) + + got := layer.Experts.Forward(x, router, cfg, 0) + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + wantVals := explicitExpertsForwardFromLoadedWeights(t, layer.Experts, cfg, x, routerVals) + + if len(gotVals) != len(wantVals) { + t.Fatalf("Experts.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 5e-2 { + t.Fatalf("Experts.Forward() real imported output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestAttentionForwardMatchesExplicitRealImportedModel(t *testing.T) { + skipIfNoMLX(t) + + if os.Getenv("OLLAMA_MODELS") == "" { + t.Skip("OLLAMA_MODELS not set") + } + + root, err := model.Open("gptoss-mlx-runtime") + if err != nil { + t.Skipf("imported model not available: %v", err) + } + defer root.Close() + + baseModel, err := base.New(root) + if err != nil { + t.Fatalf("base.New() error = %v", err) + } + tensors, err := loadRuntimeTensorsForTest(root) + if err != nil { + t.Fatalf("loadRuntimeTensorsForTest() error = %v", err) + } + if err := base.Weights(baseModel)(tensors); err != nil { + t.Fatalf("load weights error = %v", err) + } + + m := baseModel.(*Model) + layer := m.Layers[0] + cfg := m.Config + + xVals := make([]float32, 2*int(cfg.HiddenSize)) + for i := range xVals { + xVals[i] = float32((i%13)-6) / 6 + } + x := mlx.FromValues(xVals, 1, 2, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + + got := layer.Attention.Forward(x, nil, 1, 2, cfg, 0) + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + wantVals := explicitAttentionForwardFromLoadedWeights(t, layer.Attention, cfg, xVals) + + if len(gotVals) != len(wantVals) { + t.Fatalf("Attention.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 1e-1 { + t.Fatalf("Attention.Forward() real imported output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestLayerForwardMatchesExplicitRealImportedModel(t *testing.T) { + skipIfNoMLX(t) + + if os.Getenv("OLLAMA_MODELS") == "" { + t.Skip("OLLAMA_MODELS not set") + } + + root, err := model.Open("gptoss-mlx-runtime") + if err != nil { + t.Skipf("imported model not available: %v", err) + } + defer root.Close() + + baseModel, err := base.New(root) + if err != nil { + t.Fatalf("base.New() error = %v", err) + } + tensors, err := loadRuntimeTensorsForTest(root) + if err != nil { + t.Fatalf("loadRuntimeTensorsForTest() error = %v", err) + } + if err := base.Weights(baseModel)(tensors); err != nil { + t.Fatalf("load weights error = %v", err) + } + + m := baseModel.(*Model) + layer := m.Layers[0] + cfg := m.Config + + xVals := make([]float32, 2*int(cfg.HiddenSize)) + for i := range xVals { + xVals[i] = float32((i%13)-6) / 6 + } + x := mlx.FromValues(xVals, 1, 2, int(cfg.HiddenSize)).AsType(mlx.DTypeBFloat16) + + got := layer.Forward(x, nil, 1, 2, cfg, 0) + gotVals := materializedFloats(got.AsType(mlx.DTypeFloat32)) + wantVals := explicitLayerForwardFromLoadedWeights(t, layer, cfg, xVals) + + if len(gotVals) != len(wantVals) { + t.Fatalf("Layer.Forward() output length = %d, want %d", len(gotVals), len(wantVals)) + } + for i := range wantVals { + if diff := math.Abs(float64(gotVals[i] - wantVals[i])); diff > 1.5e-1 { + t.Fatalf("Layer.Forward() real imported output[%d] = %v, want %v (diff %v)", i, gotVals[i], wantVals[i], diff) + } + } +} + +func TestAttentionForwardMatchesReferenceWithKVCache(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.HeadDim = 2 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.RopeTheta = 10000 + cfg.RopeScaling.Factor = 1 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4 + refreshRow11CheckedConfig(t, &cfg) + + xVals := []float32{ + 1.0, 0.5, + -0.25, 0.75, + 0.4, -0.6, + 0.2, 0.3, + } + + attn := &Attention{ + QProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, nil, 2, 2), + KProj: nnLinearFromValues([]float32{ + 0.5, 0, + 0, 1.5, + }, []float32{0.1, -0.2}, 2, 2), + VProj: nnLinearFromValues([]float32{ + 1.2, 0, + 0, 0.8, + }, []float32{-0.05, 0.2}, 2, 2), + OProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, []float32{0.01, -0.02}, 2, 2), + Sinks: preparedSyntheticSink(t, []float32{row11NoSinkSentinel}, mlx.DTypeFloat32), + } + + c := cache.NewKVCache() + gotVals := make([]float32, 0, len(xVals)) + for pos := range len(xVals) / 2 { + x := mlx.FromValues(xVals[pos*2:(pos+1)*2], 1, 1, 2) + got := attn.Forward(x, c, 1, 1, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatalf("Attention.Forward() returned invalid tensor at token %d", pos) + } + gotVals = append(gotVals, materializedFloats(got.AsType(mlx.DTypeFloat32))...) + } + + want := referenceAttentionForwardWindowed( + xVals, + [][]float32{{1, 0}, {0, 1}}, + nil, + [][]float32{{0.5, 0}, {0, 1.5}}, + []float32{0.1, -0.2}, + [][]float32{{1.2, 0}, {0, 0.8}}, + []float32{-0.05, 0.2}, + [][]float32{{1, 0}, {0, 1}}, + []float32{0.01, -0.02}, + cfg.RopeTheta, + 0, + ) + + if len(gotVals) != len(want) { + t.Fatalf("Attention.Forward() output length = %d, want %d", len(gotVals), len(want)) + } + for i := range want { + if diff := math.Abs(float64(gotVals[i] - want[i])); diff > 2e-1 { + t.Fatalf("cached Attention.Forward() output[%d] = %v, want %v (diff %v)", i, gotVals[i], want[i], diff) + } + } +} + +func TestAttentionForwardMatchesReferenceWithSlidingWindowCache(t *testing.T) { + skipIfNoMLX(t) + + cfg := denseTestConfig(t) + cfg.HiddenSize = 2 + cfg.HeadDim = 2 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.RopeTheta = 10000 + cfg.RopeScaling.Factor = 1 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4 + cfg.SlidingWindow = 2 + refreshRow11CheckedConfig(t, &cfg) + + xVals := []float32{ + 1.0, 0.5, + -0.25, 0.75, + 0.4, -0.6, + 0.2, 0.3, + } + + attn := &Attention{ + QProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, nil, 2, 2), + KProj: nnLinearFromValues([]float32{ + 0.5, 0, + 0, 1.5, + }, []float32{0.1, -0.2}, 2, 2), + VProj: nnLinearFromValues([]float32{ + 1.2, 0, + 0, 0.8, + }, []float32{-0.05, 0.2}, 2, 2), + OProj: nnLinearFromValues([]float32{ + 1, 0, + 0, 1, + }, []float32{0.01, -0.02}, 2, 2), + Sinks: preparedSyntheticSink(t, []float32{row11NoSinkSentinel}, mlx.DTypeFloat32), + } + + c := cache.NewRotatingKVCache(int(cfg.SlidingWindow)) + gotVals := make([]float32, 0, len(xVals)) + for pos := range len(xVals) / 2 { + x := mlx.FromValues(xVals[pos*2:(pos+1)*2], 1, 1, 2) + got := attn.Forward(x, c, 1, 1, &cfg, 0) + if got == nil || !got.Valid() { + t.Fatalf("Attention.Forward() returned invalid tensor at token %d", pos) + } + gotVals = append(gotVals, materializedFloats(got.AsType(mlx.DTypeFloat32))...) + } + + want := referenceAttentionForwardWindowed( + xVals, + [][]float32{{1, 0}, {0, 1}}, + nil, + [][]float32{{0.5, 0}, {0, 1.5}}, + []float32{0.1, -0.2}, + [][]float32{{1.2, 0}, {0, 0.8}}, + []float32{-0.05, 0.2}, + [][]float32{{1, 0}, {0, 1}}, + []float32{0.01, -0.02}, + cfg.RopeTheta, + int(cfg.SlidingWindow), + ) + + if len(gotVals) != len(want) { + t.Fatalf("Attention.Forward() output length = %d, want %d", len(gotVals), len(want)) + } + for i := range want { + if diff := math.Abs(float64(gotVals[i] - want[i])); diff > 2e-1 { + t.Fatalf("sliding Attention.Forward() output[%d] = %v, want %v (diff %v)", i, gotVals[i], want[i], diff) + } + } +} + +func skipIfNoMLX(t *testing.T) { + t.Helper() + mlxTestMu.Lock() + runtime.LockOSThread() + if err := mlx.CheckInit(); err != nil { + runtime.UnlockOSThread() + mlxTestMu.Unlock() + t.Skipf("MLX not available: %v", err) + } + if mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + t.Cleanup(func() { + runtime.UnlockOSThread() + mlxTestMu.Unlock() + }) +} + +func denseTestConfig(t *testing.T) Config { + t.Helper() + + cfg, err := parseConfig([]byte(`{ + "architectures": ["GptOssForCausalLM"], + "model_type": "gpt_oss", + "num_hidden_layers": 2, + "hidden_size": 4, + "intermediate_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 2, + "num_local_experts": 2, + "num_experts_per_tok": 1, + "sliding_window": 4, + "rope_theta": 150000, + "rope_scaling": { + "factor": 2.0, + "original_max_position_embeddings": 4 + }, + "rms_norm_eps": 0.00001, + "vocab_size": 8, + "tie_word_embeddings": false, + "quantization_config": { + "quant_method": "mxfp4" + } + }`)) + if err != nil { + t.Fatalf("parseConfig() error = %v", err) + } + + return cfg +} + +func denseTestTensors(t *testing.T, cfg Config) map[string]*mlx.Array { + t.Helper() + + tensors := map[string]*mlx.Array{ + "embedding.weight": denseMatrix(int(cfg.VocabSize), int(cfg.HiddenSize), 1), + "output_norm.weight": denseVector(int(cfg.HiddenSize), 1), + "output.weight": denseMatrix(int(cfg.VocabSize), int(cfg.HiddenSize), 2), + } + + for i := range cfg.NumHiddenLayers { + prefix := fmt.Sprintf("blocks.%d", i) + tensors[prefix+".attn_norm.weight"] = denseVector(int(cfg.HiddenSize), 3+float32(i)) + tensors[prefix+".q_proj.weight"] = denseMatrix(int(cfg.NumAttentionHeads*cfg.HeadDim), int(cfg.HiddenSize), 4+float32(i)) + tensors[prefix+".q_proj.bias"] = denseVector(int(cfg.NumAttentionHeads*cfg.HeadDim), 5+float32(i)) + tensors[prefix+".k_proj.weight"] = denseMatrix(int(cfg.NumKeyValueHeads*cfg.HeadDim), int(cfg.HiddenSize), 6+float32(i)) + tensors[prefix+".k_proj.bias"] = denseVector(int(cfg.NumKeyValueHeads*cfg.HeadDim), 7+float32(i)) + tensors[prefix+".v_proj.weight"] = denseMatrix(int(cfg.NumKeyValueHeads*cfg.HeadDim), int(cfg.HiddenSize), 8+float32(i)) + tensors[prefix+".v_proj.bias"] = denseVector(int(cfg.NumKeyValueHeads*cfg.HeadDim), 9+float32(i)) + tensors[prefix+".attn_out.weight"] = denseMatrix(int(cfg.HiddenSize), int(cfg.NumAttentionHeads*cfg.HeadDim), 10+float32(i)) + tensors[prefix+".attn_out.bias"] = denseVector(int(cfg.HiddenSize), 11+float32(i)) + tensors[prefix+".attn_sinks"] = denseVector(int(cfg.NumAttentionHeads), 12+float32(i)).AsType(mlx.DTypeBFloat16) + tensors[prefix+".ffn_norm.weight"] = denseVector(int(cfg.HiddenSize), 13+float32(i)) + tensors[prefix+".router.weight"] = denseMatrix(int(cfg.NumLocalExperts), int(cfg.HiddenSize), 14+float32(i)) + tensors[prefix+".router.bias"] = denseVector(int(cfg.NumLocalExperts), 15+float32(i)) + tensors[prefix+".experts.gate_proj.weight"] = denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize), 16+float32(i)).AsType(mlx.DTypeBFloat16) + tensors[prefix+".experts.gate_proj.bias"] = expertBias(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), 0) + tensors[prefix+".experts.up_proj.weight"] = denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), int(cfg.HiddenSize), 20+float32(i)).AsType(mlx.DTypeBFloat16) + tensors[prefix+".experts.up_proj.bias"] = expertBias(int(cfg.NumLocalExperts), int(cfg.IntermediateSize), 0.5) + tensors[prefix+".experts.down_proj.weight"] = denseExpertWeight(int(cfg.NumLocalExperts), int(cfg.HiddenSize), int(cfg.IntermediateSize), 24+float32(i)).AsType(mlx.DTypeBFloat16) + tensors[prefix+".experts.down_proj.bias"] = expertBias(int(cfg.NumLocalExperts), int(cfg.HiddenSize), 0) + } + + return tensors +} + +// denseRuntimeTestTensors models the BF16 checkpoint used by the production +// GPT-OSS path. Keeping denseTestTensors in its original dtype lets loader-only +// tests continue to vary individual tensor contracts independently. +func denseRuntimeTestTensors(t *testing.T, cfg Config) map[string]*mlx.Array { + t.Helper() + tensors := denseTestTensors(t, cfg) + for name, tensor := range tensors { + tensors[name] = tensor.AsType(mlx.DTypeBFloat16) + } + return tensors +} + +// prepareSyntheticF32Sinks adapts direct synthetic layer/SDPA fixtures whose +// deliberately F32 projections bypass the model-load/runtime BF16 contract. +func prepareSyntheticF32Sinks(t *testing.T, m *Model) { + t.Helper() + for i, layer := range m.Layers { + if layer == nil || layer.Attention == nil || layer.Attention.Sinks == nil { + t.Fatalf("layer %d is missing attention sinks", i) + } + values := materializedFloats(layer.Attention.Sinks.AsType(mlx.DTypeFloat32)) + layer.Attention.Sinks = preparedSyntheticSink(t, values, mlx.DTypeFloat32) + } +} + +func preparedSyntheticSink(t *testing.T, values []float32, dtype mlx.DType) *mlx.Array { + t.Helper() + prepared := mlx.Contiguous(mlx.FromValues(values, len(values)).AsType(dtype), false) + mlx.Eval(prepared) + return prepared +} + +func quantizeTensorForTest(t *testing.T, tensors map[string]*mlx.Array, name string, groupSize, bits int, mode string) { + t.Helper() + weight := tensors[name] + if weight == nil { + t.Fatalf("quantizeTensorForTest(%q): missing weight tensor", name) + } + qw, scales, qbiases := mlx.Quantize(weight, groupSize, bits, mode) + if qw == nil || scales == nil { + t.Fatalf("quantizeTensorForTest(%q): quantize returned invalid tensors", name) + } + tensors[name] = qw + tensors[name+"_scale"] = scales + if qbiases != nil { + tensors[name+"_qbias"] = qbiases + } +} + +func loadRuntimeTensorsForTest(root *model.Root) (map[string]*mlx.Array, error) { + rawTensors := make(map[string]*mlx.Array) + seen := make(map[string]bool) + for _, layer := range root.Manifest.GetTensorLayers("") { + if seen[layer.Digest] { + continue + } + seen[layer.Digest] = true + blobPath := root.Manifest.BlobPath(layer.Digest) + for name, arr := range mlx.Load(blobPath) { + rawTensors[name] = arr + } + } + + scaleBaseNames := make(map[string]bool) + allTensors := make(map[string]*mlx.Array, len(rawTensors)) + for name, arr := range rawTensors { + if strings.HasSuffix(name, ".scale") { + baseName := strings.TrimSuffix(name, ".scale") + allTensors[baseName+"_scale"] = arr + scaleBaseNames[baseName] = true + } + } + + for name, arr := range rawTensors { + if strings.HasSuffix(name, ".scale") { + continue + } + if strings.HasSuffix(name, ".bias") && !strings.HasSuffix(name, ".weight_qbias") { + baseName := strings.TrimSuffix(name, ".bias") + if scaleBaseNames[baseName] { + allTensors[baseName+"_qbias"] = arr + } else { + allTensors[name] = arr + } + } else { + allTensors[name] = arr + } + } + return allTensors, nil +} + +func explicitExpertsForwardFromLoadedWeights(t *testing.T, experts *Experts, cfg *Config, x *mlx.Array, routerVals []float32) []float32 { + t.Helper() + + type routedExpert struct { + index int + logit float32 + } + + dims := x.Dims() + if len(dims) != 3 { + t.Fatalf("explicitExpertsForwardFromLoadedWeights() x dims = %v, want [batch seq hidden]", dims) + } + if dims[0] != 1 { + t.Fatalf("explicitExpertsForwardFromLoadedWeights() batch = %d, want 1", dims[0]) + } + + seqLen := dims[1] + hidden := dims[2] + if len(routerVals) != seqLen*int(cfg.NumLocalExperts) { + t.Fatalf("explicitExpertsForwardFromLoadedWeights() router length = %d, want %d", len(routerVals), seqLen*int(cfg.NumLocalExperts)) + } + + xVals := materializedFloats(x.AsType(mlx.DTypeFloat32)) + out := make([]float32, seqLen*hidden) + + gateW := materializedFloats(experts.GateUp.Gate.Weight.AsType(mlx.DTypeFloat32)) + gateB := materializedFloats(experts.GateUp.Gate.Bias.AsType(mlx.DTypeFloat32)) + upW := materializedFloats(experts.GateUp.Up.Weight.AsType(mlx.DTypeFloat32)) + upB := materializedFloats(experts.GateUp.Up.Bias.AsType(mlx.DTypeFloat32)) + downW := materializedFloats(experts.Down.Weight.AsType(mlx.DTypeFloat32)) + downB := materializedFloats(experts.Down.Bias.AsType(mlx.DTypeFloat32)) + + gateWMats := make([][][]float32, int(cfg.NumLocalExperts)) + gateBMats := make([][]float32, int(cfg.NumLocalExperts)) + upWMats := make([][][]float32, int(cfg.NumLocalExperts)) + upBMats := make([][]float32, int(cfg.NumLocalExperts)) + downWMats := make([][][]float32, int(cfg.NumLocalExperts)) + downBMats := make([][]float32, int(cfg.NumLocalExperts)) + + for expert := range int(cfg.NumLocalExperts) { + gateWMats[expert] = make([][]float32, int(cfg.IntermediateSize)) + upWMats[expert] = make([][]float32, int(cfg.IntermediateSize)) + downWMats[expert] = make([][]float32, int(cfg.HiddenSize)) + gateBMats[expert] = append([]float32(nil), gateB[expert*int(cfg.IntermediateSize):(expert+1)*int(cfg.IntermediateSize)]...) + upBMats[expert] = append([]float32(nil), upB[expert*int(cfg.IntermediateSize):(expert+1)*int(cfg.IntermediateSize)]...) + downBMats[expert] = append([]float32(nil), downB[expert*int(cfg.HiddenSize):(expert+1)*int(cfg.HiddenSize)]...) + + for row := range int(cfg.IntermediateSize) { + start := (expert*int(cfg.IntermediateSize) + row) * hidden + gateWMats[expert][row] = append([]float32(nil), gateW[start:start+hidden]...) + upWMats[expert][row] = append([]float32(nil), upW[start:start+hidden]...) + } + for row := range int(cfg.HiddenSize) { + start := (expert*int(cfg.HiddenSize) + row) * int(cfg.IntermediateSize) + downWMats[expert][row] = append([]float32(nil), downW[start:start+int(cfg.IntermediateSize)]...) + } + } + + for pos := range seqLen { + selected := make([]routedExpert, int(cfg.NumLocalExperts)) + for expert := range int(cfg.NumLocalExperts) { + selected[expert] = routedExpert{ + index: expert, + logit: routerVals[pos*int(cfg.NumLocalExperts)+expert], + } + } + slices.SortFunc(selected, func(a, b routedExpert) int { + switch { + case a.logit > b.logit: + return -1 + case a.logit < b.logit: + return 1 + default: + return 0 + } + }) + topK := int(cfg.NumExpertsPerTok) + if topK > len(selected) { + topK = len(selected) + } + selected = selected[:topK] + + xRow := xVals[pos*hidden : (pos+1)*hidden] + routerRow := make([]float32, len(selected)) + selectedIndices := make([]int, len(selected)) + for i, s := range selected { + routerRow[i] = s.logit + selectedIndices[i] = s.index + } + expertOut := referenceExpertsForward( + xRow, + routerRow, + topK, + pickExpertMatrices(gateWMats, selectedIndices), + pickExpertBiases(gateBMats, selectedIndices), + pickExpertMatrices(upWMats, selectedIndices), + pickExpertBiases(upBMats, selectedIndices), + pickExpertMatrices(downWMats, selectedIndices), + pickExpertBiases(downBMats, selectedIndices), + ) + copy(out[pos*hidden:(pos+1)*hidden], expertOut) + } + + return out +} + +func pickExpertMatrices(src [][][]float32, selected []int) [][][]float32 { + out := make([][][]float32, len(selected)) + for i, index := range selected { + out[i] = src[index] + } + return out +} + +func pickExpertBiases(src [][]float32, selected []int) [][]float32 { + out := make([][]float32, len(selected)) + for i, index := range selected { + out[i] = src[index] + } + return out +} + +func explicitAttentionForwardFromLoadedWeights(t *testing.T, attn *Attention, cfg *Config, xVals []float32) []float32 { + t.Helper() + + hidden := int(cfg.HiddenSize) + seqLen := len(xVals) / hidden + qOut := int(cfg.NumAttentionHeads * cfg.HeadDim) + kvOut := int(cfg.NumKeyValueHeads * cfg.HeadDim) + headDim := int(cfg.HeadDim) + numHeads := int(cfg.NumAttentionHeads) + numKVHeads := int(cfg.NumKeyValueHeads) + qMul := numHeads / numKVHeads + + qW := materializedFloats(attn.QProj.(*nn.Linear).Weight.AsType(mlx.DTypeFloat32)) + qB := materializedFloats(attn.QProj.(*nn.Linear).Bias.AsType(mlx.DTypeFloat32)) + kW := materializedFloats(attn.KProj.(*nn.Linear).Weight.AsType(mlx.DTypeFloat32)) + kB := materializedFloats(attn.KProj.(*nn.Linear).Bias.AsType(mlx.DTypeFloat32)) + vW := materializedFloats(attn.VProj.(*nn.Linear).Weight.AsType(mlx.DTypeFloat32)) + vB := materializedFloats(attn.VProj.(*nn.Linear).Bias.AsType(mlx.DTypeFloat32)) + oW := materializedFloats(attn.OProj.(*nn.Linear).Weight.AsType(mlx.DTypeFloat32)) + oB := materializedFloats(attn.OProj.(*nn.Linear).Bias.AsType(mlx.DTypeFloat32)) + sinks := materializedFloats(attn.Sinks.AsType(mlx.DTypeFloat32)) + + denoms := referenceGPTOSSRoPEDenominators(cfg) + concentration := testGPTOSSYarnScale(cfg) + scale := float32(1 / math.Sqrt(float64(headDim))) + + query := make([][][]float32, seqLen) + key := make([][][]float32, seqLen) + value := make([][][]float32, seqLen) + for pos := range seqLen { + x := xVals[pos*hidden : (pos+1)*hidden] + qVec := affineFlatRef(x, qW, qB, qOut, hidden) + kVec := affineFlatRef(x, kW, kB, kvOut, hidden) + vVec := affineFlatRef(x, vW, vB, kvOut, hidden) + + query[pos] = make([][]float32, numHeads) + for h := range numHeads { + query[pos][h] = applyRoPEGeneric(qVec[h*headDim:(h+1)*headDim], pos, denoms, concentration) + } + key[pos] = make([][]float32, numKVHeads) + value[pos] = make([][]float32, numKVHeads) + for h := range numKVHeads { + key[pos][h] = applyRoPEGeneric(kVec[h*headDim:(h+1)*headDim], pos, denoms, concentration) + value[pos][h] = append([]float32(nil), vVec[h*headDim:(h+1)*headDim]...) + } + } + + out := make([]float32, seqLen*hidden) + for pos := range seqLen { + attnHidden := make([]float32, numHeads*headDim) + for h := range numHeads { + kvHead := h / qMul + sink := sinks[h] + scores := make([]float32, pos+2) + scores[0] = sink + maxScore := sink + for j := 0; j <= pos; j++ { + score := dotRef(query[pos][h], key[j][kvHead]) * scale + scores[j+1] = score + if score > maxScore { + maxScore = score + } + } + sum := float32(0) + for i := range scores { + scores[i] = float32(math.Exp(float64(scores[i] - maxScore))) + sum += scores[i] + } + for i := range scores { + scores[i] /= sum + } + for j := 0; j <= pos; j++ { + w := scores[j+1] + for d := range headDim { + attnHidden[h*headDim+d] += w * value[j][kvHead][d] + } + } + } + projected := affineFlatRef(attnHidden, oW, oB, hidden, numHeads*headDim) + copy(out[pos*hidden:(pos+1)*hidden], projected) + } + return out +} + +func explicitLayerForwardFromLoadedWeights(t *testing.T, layer *Layer, cfg *Config, xVals []float32) []float32 { + t.Helper() + + hidden := int(cfg.HiddenSize) + seqLen := len(xVals) / hidden + attnNormWeight := materializedFloats(layer.AttentionNorm.Weight.AsType(mlx.DTypeFloat32)) + ffnNormWeight := materializedFloats(layer.FFNNorm.Weight.AsType(mlx.DTypeFloat32)) + routerLinear := layer.Router.(*nn.Linear) + routerW := materializedFloats(routerLinear.Weight.AsType(mlx.DTypeFloat32)) + routerB := materializedFloats(routerLinear.Bias.AsType(mlx.DTypeFloat32)) + + attnIn := make([]float32, len(xVals)) + for pos := range seqLen { + copy(attnIn[pos*hidden:(pos+1)*hidden], referenceRMSNorm(xVals[pos*hidden:(pos+1)*hidden], attnNormWeight, cfg.RMSNormEps)) + } + attnOut := explicitAttentionForwardFromLoadedWeights(t, layer.Attention, cfg, attnIn) + + postAttn := make([]float32, len(xVals)) + for i := range xVals { + postAttn[i] = xVals[i] + attnOut[i] + } + + ffnIn := make([]float32, len(postAttn)) + for pos := range seqLen { + copy(ffnIn[pos*hidden:(pos+1)*hidden], referenceRMSNorm(postAttn[pos*hidden:(pos+1)*hidden], ffnNormWeight, cfg.RMSNormEps)) + } + + routerVals := make([]float32, seqLen*int(cfg.NumLocalExperts)) + for pos := range seqLen { + row := affineFlatRef( + ffnIn[pos*hidden:(pos+1)*hidden], + routerW, + routerB, + int(cfg.NumLocalExperts), + hidden, + ) + copy(routerVals[pos*int(cfg.NumLocalExperts):(pos+1)*int(cfg.NumLocalExperts)], row) + } + + xArr := mlx.FromValues(ffnIn, 1, seqLen, hidden).AsType(mlx.DTypeBFloat16) + expertOut := explicitExpertsForwardFromLoadedWeights(t, layer.Experts, cfg, xArr, routerVals) + + out := make([]float32, len(postAttn)) + for i := range out { + out[i] = postAttn[i] + expertOut[i] + } + return out +} + +func referenceRMSNorm(x, weight []float32, eps float32) []float32 { + ss := float32(0) + for _, v := range x { + ss += v * v + } + inv := float32(1 / math.Sqrt(float64(ss/float32(len(x))+eps))) + out := make([]float32, len(x)) + for i := range x { + out[i] = x[i] * inv * weight[i] + } + return out +} + +func affineFlatRef(x, w, b []float32, outDim, inDim int) []float32 { + out := make([]float32, outDim) + for o := range outDim { + sum := float32(0) + row := w[o*inDim : (o+1)*inDim] + for i := range inDim { + sum += row[i] * x[i] + } + if len(b) > 0 { + sum += b[o] + } + out[o] = sum + } + return out +} + +type attentionPrefillTrace struct { + qProjLast []float32 + kProjLast []float32 + vProjLast []float32 + visibleKeyLen int + visibleValueLen int + visibleKeyLast []float32 + visibleValueLast []float32 + preOProjLast []float32 + outputLast []float32 +} + +type quantizedCachedPrefillParityCase struct { + name string + groupSize int + bits int + mode string + projectionTol float64 + attentionTol float64 + modelTol float64 +} + +func quantizedCachedPrefillParityCases() []quantizedCachedPrefillParityCase { + return []quantizedCachedPrefillParityCase{ + { + name: "affine-int8", + groupSize: 64, + bits: 8, + mode: "affine", + projectionTol: 1e-2, + attentionTol: 5e-2, + modelTol: 5e-2, + }, + { + name: "affine-int4", + groupSize: 64, + bits: 4, + mode: "affine", + projectionTol: 3e-2, + attentionTol: 1e-1, + modelTol: 1e-1, + }, + } +} + +func quantizedAttentionTestModel(t *testing.T) (*Model, Config) { + return quantizedAttentionTestModelWithQuantization(t, 64, 8, "affine") +} + +func quantizedAttentionTestModelWithQuantization(t *testing.T, groupSize, bits int, mode string) (*Model, Config) { + t.Helper() + + cfg := denseTestConfig(t) + cfg.HiddenSize = 64 + cfg.IntermediateSize = 128 + cfg.NumAttentionHeads = 1 + cfg.NumKeyValueHeads = 1 + cfg.HeadDim = 64 + cfg.NumLocalExperts = 4 + cfg.NumExpertsPerTok = 2 + cfg.SlidingWindow = 16 + cfg.RopeTheta = 150000 + cfg.RopeScaling.Factor = 32 + cfg.RopeScaling.OriginalMaxPositionEmbeddings = 4096 + cfg.RopeScaling.BetaFast = 32 + cfg.RopeScaling.BetaSlow = 1 + cfg.QuantGroupSize = groupSize + cfg.QuantBits = bits + cfg.QuantMode = mode + refreshRow11CheckedConfig(t, &cfg) + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + } + + tensors := denseTestTensors(t, cfg) + for i := range int(cfg.NumHiddenLayers) { + prefix := fmt.Sprintf("blocks.%d", i) + quantizeTensorForTest(t, tensors, prefix+".q_proj.weight", groupSize, bits, mode) + quantizeTensorForTest(t, tensors, prefix+".k_proj.weight", groupSize, bits, mode) + quantizeTensorForTest(t, tensors, prefix+".v_proj.weight", groupSize, bits, mode) + quantizeTensorForTest(t, tensors, prefix+".attn_out.weight", groupSize, bits, mode) + } + + if err := m.LoadWeights(tensors); err != nil { + t.Fatalf("LoadWeights() error = %v", err) + } + prepareSyntheticF32Sinks(t, m) + + return m, cfg +} + +func patternedHiddenValues(seqLen, hidden int) []float32 { + values := make([]float32, seqLen*hidden) + for i := range values { + values[i] = float32((i%29)-14) / 11 + } + return values +} + +func patternedTokenValues(seqLen, vocab int) []int32 { + values := make([]int32, seqLen) + for i := range values { + values[i] = int32((i*3 + 1) % vocab) + } + return values +} + +func attentionForwardTraceForTest( + t *testing.T, + a *Attention, + x *mlx.Array, + c cache.Cache, + batchSize, seqLen int, + cfg *Config, + layerIndex int, +) (*mlx.Array, attentionPrefillTrace) { + t.Helper() + + var trace attentionPrefillTrace + + query := a.QProj.Forward(x) + key := a.KProj.Forward(x) + value := a.VProj.Forward(x) + trace.qProjLast = lastTokenFloats(query.AsType(mlx.DTypeFloat32)) + trace.kProjLast = lastTokenFloats(key.AsType(mlx.DTypeFloat32)) + trace.vProjLast = lastTokenFloats(value.AsType(mlx.DTypeFloat32)) + + batchDim := int32(batchSize) + seq := int32(seqLen) + numHeads := cfg.NumAttentionHeads + numKVHeads := cfg.NumKeyValueHeads + headDim := cfg.HeadDim + + query = mlx.Reshape(query, batchDim, seq, numHeads, headDim) + key = mlx.Reshape(key, batchDim, seq, numKVHeads, headDim) + value = mlx.Reshape(value, batchDim, seq, numKVHeads, headDim) + query = mlx.Transpose(query, 0, 2, 1, 3) + key = mlx.Transpose(key, 0, 2, 1, 3) + value = mlx.Transpose(value, 0, 2, 1, 3) + + b := testForwardBatch(mlx.Zeros(mlx.DTypeInt32, batchSize, seqLen), c) + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + attentionScale := float32(1.0 / math.Sqrt(float64(cfg.HeadDim))) + if a.RoPEFreqs != nil && a.RoPEFreqs.Valid() { + query = mlx.RoPEWithFreqs(query, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, a.RoPEFreqs) + key = mlx.RoPEWithFreqs(key, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, a.RoPEFreqs) + attentionScale *= testGPTOSSYarnScale(cfg) * testGPTOSSYarnScale(cfg) + } else { + ropeBase, ropeScale, _ := cfg.RopeParameters() + query = mlx.RoPEWithBase(query, int(cfg.HeadDim), false, ropeBase, ropeScale, positions) + key = mlx.RoPEWithBase(key, int(cfg.HeadDim), false, ropeBase, ropeScale, positions) + } + + var kv nn.SDPAOption + if c != nil { + attnCache := c.(cache.Attention) + history := attnCache.Update(b, key, value) + key, value = history.K(), history.V() + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(key, value, b.SeqQueryLens) + } + + visibleKey, visibleValue := visibleKVForLastQueryForTest(key, value, c, cfg) + trace.visibleKeyLen = visibleKey.Dim(2) + trace.visibleValueLen = visibleValue.Dim(2) + trace.visibleKeyLast = lastCacheTokenFloats(visibleKey.AsType(mlx.DTypeFloat32)) + trace.visibleValueLast = lastCacheTokenFloats(visibleValue.AsType(mlx.DTypeFloat32)) + + attention := nn.ScaledDotProductAttention( + b, + query, + attentionScale, + kv, + nn.WithMask(nn.CausalMask()), + nn.WithSinks(a.Sinks), + ) + if attention == nil || !attention.Valid() { + t.Fatalf("layer %d trace attention is invalid", layerIndex) + } + attention = mlx.Transpose(attention, 0, 2, 1, 3) + attention = mlx.Reshape(attention, batchDim, seq, numHeads*headDim) + trace.preOProjLast = lastTokenFloats(attention.AsType(mlx.DTypeFloat32)) + + out := a.OProj.Forward(attention) + trace.outputLast = lastTokenFloats(out.AsType(mlx.DTypeFloat32)) + return out, trace +} + +func visibleKVForLastQueryForTest(key, value *mlx.Array, c cache.Cache, cfg *Config) (*mlx.Array, *mlx.Array) { + if _, ok := c.(*cache.RotatingKVCache); ok && cfg != nil && cfg.SlidingWindow > 0 { + start := max(key.Dim(2)-int(cfg.SlidingWindow), 0) + return key.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, key.Dim(2)), mlx.Slice()), + value.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, value.Dim(2)), mlx.Slice()) + } + return key, value +} + +func lastTokenFloats(a *mlx.Array) []float32 { + if a == nil { + return nil + } + + switch dims := a.Dims(); len(dims) { + case 2: + return materializedFloats(a.Slice(mlx.Slice(dims[0]-1), mlx.Slice()).Squeeze(0)) + case 3: + return materializedFloats(a.Slice(mlx.Slice(), mlx.Slice(dims[1]-1), mlx.Slice()).Squeeze(1)) + default: + return materializedFloats(a) + } +} + +func lastCacheTokenFloats(a *mlx.Array) []float32 { + if a == nil { + return nil + } + if dims := a.Dims(); len(dims) == 4 { + return materializedFloats(a.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(dims[2]-1), mlx.Slice()).Squeeze(2)) + } + return materializedFloats(a) +} + +func assertFloatSliceClose(t *testing.T, label string, got, want []float32, relTolerance, absTolerance float64) { + t.Helper() + + if len(got) != len(want) { + t.Fatalf("%s length mismatch: got=%d want=%d", label, len(got), len(want)) + } + + for i := range got { + diff := math.Abs(float64(got[i] - want[i])) + mag := math.Max(math.Abs(float64(want[i])), 1e-6) + if diff > absTolerance && diff/mag > relTolerance { + t.Fatalf("%s[%d] = %v, want %v (diff %v, rel %v)", label, i, got[i], want[i], diff, diff/mag) + } + } +} + +func denseMatrix(rows, cols int, start float32) *mlx.Array { + values := make([]float32, rows*cols) + for i := range values { + values[i] = start + float32(i) + } + return mlx.FromValues(values, rows, cols) +} + +func denseVector(length int, start float32) *mlx.Array { + values := make([]float32, length) + for i := range values { + values[i] = start + float32(i) + } + return mlx.FromValues(values, length) +} + +func denseExpertWeight(experts, out, in int, start float32) *mlx.Array { + values := make([]float32, experts*out*in) + for i := range values { + values[i] = start + float32(i)/100 + } + return mlx.FromValues(values, experts, out, in) +} + +func referenceExpertsForward( + x []float32, + router []float32, + topK int, + gateW [][][]float32, + gateB [][]float32, + upW [][][]float32, + upB [][]float32, + downW [][][]float32, + downB [][]float32, +) []float32 { + type routedExpert struct { + index int + logit float32 + } + + selected := make([]routedExpert, len(router)) + for i, v := range router { + selected[i] = routedExpert{index: i, logit: v} + } + slices.SortFunc(selected, func(a, b routedExpert) int { + switch { + case a.logit > b.logit: + return -1 + case a.logit < b.logit: + return 1 + default: + return 0 + } + }) + if topK > 0 && topK < len(selected) { + selected = selected[:topK] + } + + maxLogit := float32(math.Inf(-1)) + for _, s := range selected { + if s.logit > maxLogit { + maxLogit = s.logit + } + } + + scores := make([]float32, len(selected)) + sum := float32(0) + for i, s := range selected { + scores[i] = float32(math.Exp(float64(s.logit - maxLogit))) + sum += scores[i] + } + for i := range scores { + scores[i] /= sum + } + + out := make([]float32, len(x)) + for i, score := range scores { + e := selected[i].index + gate := affineRef(x, gateW[e], gateB[e]) + up := affineRef(x, upW[e], upB[e]) + + hidden := make([]float32, len(gate)) + for i := range gate { + clippedGate := gate[i] + if clippedGate > 7 { + clippedGate = 7 + } + + clippedUp := up[i] + if clippedUp < -7 { + clippedUp = -7 + } + if clippedUp > 7 { + clippedUp = 7 + } + + gated := clippedGate / (1 + float32(math.Exp(float64(-1.702*clippedGate)))) + hidden[i] = gated * (clippedUp + 1) + } + + down := affineRef(hidden, downW[e], downB[e]) + for i := range out { + out[i] += score * down[i] + } + } + + return out +} + +func TestSwiGLUAlphaLimitMatchesReference(t *testing.T) { + skipIfNoMLX(t) + + gate := mlx.FromValues([]float32{-8, -1, 0.5, 9}, 1, 1, 1, 4) + up := mlx.FromValues([]float32{-9, -0.5, 2, 8}, 1, 1, 1, 4) + + got := materializedFloats(swiGLUAlphaLimit(gate, up).AsType(mlx.DTypeFloat32)) + want := []float32{ + referenceSwiGLUAlphaLimit(-8, -9), + referenceSwiGLUAlphaLimit(-1, -0.5), + referenceSwiGLUAlphaLimit(0.5, 2), + referenceSwiGLUAlphaLimit(9, 8), + } + + for i := range want { + if diff := math.Abs(float64(got[i] - want[i])); diff > 5e-4 { + t.Fatalf("value %d = %v, want %v (diff=%v)", i, got[i], want[i], diff) + } + } +} + +func referenceSwiGLUAlphaLimit(gate, up float32) float32 { + clippedGate := gate + if clippedGate > 7 { + clippedGate = 7 + } + + clippedUp := up + if clippedUp > 7 { + clippedUp = 7 + } + if clippedUp < -7 { + clippedUp = -7 + } + + gated := clippedGate / (1 + float32(math.Exp(float64(-1.702*clippedGate)))) + return gated * (clippedUp + 1) +} + +func materializedFloats(a *mlx.Array) []float32 { + if a == nil { + return nil + } + cloned := a.Clone() + mlx.Eval(cloned) + return cloned.Floats() +} + +func identityFlat(n int) []float32 { + out := make([]float32, n*n) + for i := range n { + out[i*n+i] = 1 + } + return out +} + +func referenceGPTOSSRoPEDenominators(cfg *Config) []float32 { + dims := int(cfg.HeadDim) + dHalf := float64(dims) / 2 + base := float64(cfg.RopeTheta) + factor := float64(cfg.RopeScaling.Factor) + origCtx := float64(cfg.RopeScaling.OriginalMaxPositionEmbeddings) + betaFast := float64(cfg.RopeScaling.BetaFast) + betaSlow := float64(cfg.RopeScaling.BetaSlow) + if betaFast == 0 { + betaFast = 32 + } + if betaSlow == 0 { + betaSlow = 1 + } + + low := math.Floor(dHalf * math.Log(origCtx/(betaFast*2*math.Pi)) / math.Log(base)) + high := math.Ceil(dHalf * math.Log(origCtx/(betaSlow*2*math.Pi)) / math.Log(base)) + out := make([]float32, 0, dims/2) + for j := range dims / 2 { + divisor := math.Pow(base, float64(2*j)/float64(dims)) + ramp := (float64(j) - low) / (high - low) + if ramp < 0 { + ramp = 0 + } + if ramp > 1 { + ramp = 1 + } + mask := 1 - ramp + invFreq := (1/(factor*divisor))*(1-mask) + (1/divisor)*mask + out = append(out, float32(1/invFreq)) + } + return out +} + +func referenceAttentionForward( + x []float32, + qW [][]float32, qB []float32, + kW [][]float32, kB []float32, + vW [][]float32, vB []float32, + oW [][]float32, oB []float32, + ropeBase float32, +) []float32 { + seqLen := len(x) / 2 + query := make([][]float32, seqLen) + key := make([][]float32, seqLen) + value := make([][]float32, seqLen) + for pos := range seqLen { + query[pos] = applyRoPE2D(affineRef(x[pos*2:(pos+1)*2], qW, qB), pos, ropeBase) + key[pos] = applyRoPE2D(affineRef(x[pos*2:(pos+1)*2], kW, kB), pos, ropeBase) + value[pos] = affineRef(x[pos*2:(pos+1)*2], vW, vB) + } + + scale := float32(1 / math.Sqrt(2)) + out := make([]float32, len(x)) + for pos := range seqLen { + scores := make([]float32, pos+1) + maxScore := float32(math.Inf(-1)) + for j := 0; j <= pos; j++ { + score := dotRef(query[pos], key[j]) * scale + scores[j] = score + if score > maxScore { + maxScore = score + } + } + sum := float32(0) + for j := range scores { + scores[j] = float32(math.Exp(float64(scores[j] - maxScore))) + sum += scores[j] + } + for j := range scores { + scores[j] /= sum + } + + hidden := make([]float32, 2) + for j, score := range scores { + for d := range hidden { + hidden[d] += score * value[j][d] + } + } + projected := affineRef(hidden, oW, oB) + copy(out[pos*2:(pos+1)*2], projected) + } + return out +} + +func referenceAttentionForwardWithSinksScaledRoPE(x []float32, cfg *Config, sink float32) []float32 { + seqLen := len(x) / int(cfg.HiddenSize) + headDim := int(cfg.HeadDim) + denoms := referenceGPTOSSRoPEDenominators(cfg) + concentration := testGPTOSSYarnScale(cfg) + + query := make([][]float32, seqLen) + key := make([][]float32, seqLen) + value := make([][]float32, seqLen) + for pos := range seqLen { + base := append([]float32(nil), x[pos*headDim:(pos+1)*headDim]...) + query[pos] = applyRoPEGeneric(base, pos, denoms, concentration) + key[pos] = applyRoPEGeneric(base, pos, denoms, concentration) + value[pos] = append([]float32(nil), base...) + } + + scale := float32(1 / math.Sqrt(float64(headDim))) + out := make([]float32, len(x)) + for pos := range seqLen { + scores := make([]float32, pos+2) + maxScore := sink + scores[0] = sink + for j := 0; j <= pos; j++ { + score := dotRef(query[pos], key[j]) * scale + scores[j+1] = score + if score > maxScore { + maxScore = score + } + } + sum := float32(0) + for j := range scores { + scores[j] = float32(math.Exp(float64(scores[j] - maxScore))) + sum += scores[j] + } + for j := range scores { + scores[j] /= sum + } + + hidden := make([]float32, headDim) + for j := 0; j <= pos; j++ { + score := scores[j+1] + for d := range headDim { + hidden[d] += score * value[j][d] + } + } + copy(out[pos*headDim:(pos+1)*headDim], hidden) + } + return out +} + +func referenceAttentionForwardWindowed( + x []float32, + qW [][]float32, qB []float32, + kW [][]float32, kB []float32, + vW [][]float32, vB []float32, + oW [][]float32, oB []float32, + ropeBase float32, + window int, +) []float32 { + seqLen := len(x) / 2 + query := make([][]float32, seqLen) + key := make([][]float32, seqLen) + value := make([][]float32, seqLen) + for pos := range seqLen { + query[pos] = applyRoPE2D(affineRef(x[pos*2:(pos+1)*2], qW, qB), pos, ropeBase) + key[pos] = applyRoPE2D(affineRef(x[pos*2:(pos+1)*2], kW, kB), pos, ropeBase) + value[pos] = affineRef(x[pos*2:(pos+1)*2], vW, vB) + } + + scale := float32(1 / math.Sqrt(2)) + out := make([]float32, len(x)) + for pos := range seqLen { + start := 0 + if window > 0 && pos+1 > window { + start = pos + 1 - window + } + scores := make([]float32, pos-start+1) + maxScore := float32(math.Inf(-1)) + for j := start; j <= pos; j++ { + score := dotRef(query[pos], key[j]) * scale + scores[j-start] = score + if score > maxScore { + maxScore = score + } + } + sum := float32(0) + for j := range scores { + scores[j] = float32(math.Exp(float64(scores[j] - maxScore))) + sum += scores[j] + } + for j := range scores { + scores[j] /= sum + } + + hidden := make([]float32, 2) + for j, score := range scores { + val := value[start+j] + for d := range hidden { + hidden[d] += score * val[d] + } + } + projected := affineRef(hidden, oW, oB) + copy(out[pos*2:(pos+1)*2], projected) + } + return out +} + +func applyRoPE2D(v []float32, position int, _ float32) []float32 { + if len(v) != 2 { + panic("applyRoPE2D expects length-2 vector") + } + theta := float64(position) + c, s := float32(math.Cos(theta)), float32(math.Sin(theta)) + return []float32{ + v[0]*c - v[1]*s, + v[0]*s + v[1]*c, + } +} + +func applyRoPEGeneric(v []float32, position int, denoms []float32, concentration float32) []float32 { + if len(v)%2 != 0 { + panic("applyRoPEGeneric expects even-length vector") + } + out := make([]float32, len(v)) + for i := 0; i < len(v); i += 2 { + theta := float64(position) / float64(denoms[i/2]) + c := concentration * float32(math.Cos(theta)) + s := concentration * float32(math.Sin(theta)) + out[i] = v[i]*c - v[i+1]*s + out[i+1] = v[i]*s + v[i+1]*c + } + return out +} + +func dotRef(a, b []float32) float32 { + sum := float32(0) + for i := range a { + sum += a[i] * b[i] + } + return sum +} + +func nnLinearFromValues(weightVals, biasVals []float32, out, in int) *nn.Linear { + weight := mlx.FromValues(weightVals, out, in).AsType(mlx.DTypeBFloat16) + var bias *mlx.Array + if biasVals != nil { + bias = mlx.FromValues(biasVals, out).AsType(mlx.DTypeBFloat16) + } + return nn.NewLinear(weight, bias) +} + +func affineRef(x []float32, w [][]float32, b []float32) []float32 { + width := 0 + if len(b) > 0 { + width = len(b) + } else if len(w) > 0 { + width = len(w[0]) + } + out := make([]float32, width) + copy(out, b) + for i := range x { + for j := range out { + out[j] += x[i] * w[i][j] + } + } + return out +} + +func expertBias(experts, out int, start float32) *mlx.Array { + values := make([]float32, experts*out) + for i := range values { + values[i] = start + float32(i)/100 + } + return mlx.FromValues(values, experts, out).AsType(mlx.DTypeBFloat16) +} + +func testRoot(t *testing.T, configJSON, tokenizerJSON []byte) *model.Root { + t.Helper() + + dir := t.TempDir() + blobDir := filepath.Join(dir, "blobs") + if err := os.MkdirAll(blobDir, 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + + writeBlob := func(digest string, content []byte) { + t.Helper() + path := filepath.Join(blobDir, strings.Replace(digest, ":", "-", 1)) + if err := os.WriteFile(path, content, 0o644); err != nil { + t.Fatalf("WriteFile(%s) error = %v", path, err) + } + } + + writeBlob("sha256:config", configJSON) + writeBlob("sha256:tokenizer", tokenizerJSON) + + mf := &manifest.Manifest{ + SchemaVersion: 2, + MediaType: "application/vnd.ollama.image.model", + Layers: []manifest.ManifestLayer{ + {MediaType: "application/vnd.ollama.image.json", Digest: "sha256:config", Name: "config.json"}, + {MediaType: "application/vnd.ollama.image.json", Digest: "sha256:tokenizer", Name: "tokenizer.json"}, + }, + } + + return &model.Root{ + Manifest: &manifest.ModelManifest{ + Manifest: mf, + BlobDir: blobDir, + }, + } +} + +func testRootWithExtraConfigs(t *testing.T, configJSON, tokenizerJSON []byte, extraConfigs map[string][]byte) *model.Root { + t.Helper() + + dir := t.TempDir() + blobDir := filepath.Join(dir, "blobs") + if err := os.MkdirAll(blobDir, 0o755); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + + writeBlob := func(digest string, content []byte) { + t.Helper() + path := filepath.Join(blobDir, strings.Replace(digest, ":", "-", 1)) + if err := os.WriteFile(path, content, 0o644); err != nil { + t.Fatalf("WriteFile(%s) error = %v", path, err) + } + } + + writeBlob("sha256:config", configJSON) + writeBlob("sha256:tokenizer", tokenizerJSON) + + layers := []manifest.ManifestLayer{ + {MediaType: "application/vnd.ollama.image.json", Digest: "sha256:config", Name: "config.json"}, + {MediaType: "application/vnd.ollama.image.json", Digest: "sha256:tokenizer", Name: "tokenizer.json"}, + } + + extraIndex := 0 + for name, content := range extraConfigs { + digest := fmt.Sprintf("sha256:extra-%d", extraIndex) + extraIndex++ + writeBlob(digest, content) + layers = append(layers, manifest.ManifestLayer{ + MediaType: "application/vnd.ollama.image.json", + Digest: digest, + Name: name, + }) + } + + mf := &manifest.Manifest{ + SchemaVersion: 2, + MediaType: "application/vnd.ollama.image.model", + Layers: layers, + } + + return &model.Root{ + Manifest: &manifest.ModelManifest{ + Manifest: mf, + BlobDir: blobDir, + }, + } +} From 950a0edb4237cf3e96d83823b0c10043be67a7f1 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 18:12:06 +0000 Subject: [PATCH 23/58] create: bound GPT-OSS expert conversion memory Co-authored-by: Codex --- x/create/classify.go | 62 ++++++- x/create/gptoss_test.go | 372 ++++++++++++++++++++++++++++++++++++++++ x/create/plan.go | 8 + 3 files changed, 441 insertions(+), 1 deletion(-) create mode 100644 x/create/gptoss_test.go diff --git a/x/create/classify.go b/x/create/classify.go index e8fc28e2c2c..a11b453e710 100644 --- a/x/create/classify.go +++ b/x/create/classify.go @@ -55,12 +55,19 @@ func Classify(inv Inventory, requested string) (Classification, error) { return Classification{Kind: SourceFloat, Quantize: requested}, nil case SourcePrequantized: + effective := detectPrequantizedQuantization(inv) + if gptossNativeMXFP4Evidence(inv) { + effective = "mxfp4" + } if requested != "" { + if effective == requested && gptossNativeMXFP4Evidence(inv) { + return Classification{Kind: SourcePrequantized, Quantize: effective}, nil + } return Classification{}, fmt.Errorf("cannot requantize an already-quantized source model (requested %q): only bf16/fp16/fp32 sources can be quantized", requested) } return Classification{ Kind: SourcePrequantized, - Quantize: detectPrequantizedQuantization(inv), + Quantize: effective, }, nil case SourceBlockFP8: @@ -126,6 +133,10 @@ func detectKind(inv Inventory) SourceKind { switch { case strings.HasSuffix(name, ".scales"): hasMLXScales = true + case strings.HasSuffix(name, "_scales"): + hasMLXScales = true + case strings.HasSuffix(name, "_blocks"): + hasPacked = true case strings.HasSuffix(name, ".weight_packed"): hasPacked = true case strings.HasSuffix(name, ".weight_scale"): @@ -151,6 +162,55 @@ func detectKind(inv Inventory) SourceKind { } } +func gptossNativeMXFP4Evidence(inv Inventory) bool { + if inv.Config.Architecture() != "GptOssForCausalLM" { + return false + } + var selected sourceQuantization + for _, candidate := range inv.Config.quantizationConfigs() { + if candidate.Bits != 0 { + selected = candidate + break + } + } + mode := selected.Mode + if mode == "" { + mode = selected.QuantMethod + } + if selected.Bits != 4 || selected.GroupSize != 32 || !strings.EqualFold(mode, "mxfp4") { + return false + } + + found := false + for _, name := range sortedTensorNames(inv) { + tensor := inv.Tensors[name] + switch { + case strings.HasSuffix(name, "_blocks"): + found = true + scales, ok := inv.Tensors[strings.TrimSuffix(name, "_blocks")+"_scales"] + bias, hasBias := inv.Tensors[strings.TrimSuffix(name, "_blocks")+"_bias"] + if !ok || tensor.Dtype != "U8" || scales.Dtype != "U8" || len(tensor.Shape) != 4 || len(scales.Shape) != 3 || + tensor.Shape[3] != 16 || tensor.Shape[0] <= 0 || tensor.Shape[1] <= 0 || tensor.Shape[2] <= 0 || + tensor.Shape[0] != scales.Shape[0] || tensor.Shape[1] != scales.Shape[1] || tensor.Shape[2] != scales.Shape[2] || + !hasBias || bias.Dtype != "BF16" || len(bias.Shape) != 2 || bias.Shape[0] != tensor.Shape[0] || bias.Shape[1] != tensor.Shape[1] { + return false + } + if strings.HasSuffix(name, "gate_up_proj_blocks") && tensor.Shape[1]%2 != 0 { + return false + } + case strings.HasSuffix(name, "_scales"): + if _, ok := inv.Tensors[strings.TrimSuffix(name, "_scales")+"_blocks"]; !ok { + return false + } + case strings.HasSuffix(name, "_bias") && strings.Contains(name, ".experts."): + if _, ok := inv.Tensors[strings.TrimSuffix(name, "_bias")+"_blocks"]; !ok { + return false + } + } + } + return found +} + // firstUnsupportedFP8 returns the name of the first F8_E5M2 weight in the // source, if any. We decode only E4M3, so an E5M2 source must be rejected // explicitly rather than silently mishandled. diff --git a/x/create/gptoss_test.go b/x/create/gptoss_test.go new file mode 100644 index 00000000000..51badb5e54f --- /dev/null +++ b/x/create/gptoss_test.go @@ -0,0 +1,372 @@ +package create + +import ( + "io" + "maps" + "path/filepath" + "slices" + "testing" + + st "github.com/ollama/ollama/x/safetensors" +) + +func TestNewTensorImportTransform_GptOSSRegistered(t *testing.T) { + inv := gptossInventory(nil) + transform, err := newTensorImportTransform(inv) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + + if _, ok := transform.(*gptossImportTransform); !ok { + t.Fatalf("newTensorImportTransform() type = %T, want *gptossImportTransform", transform) + } +} + +func TestGPTOSSImportTransformRenamesTensors(t *testing.T) { + transform := &gptossImportTransform{} + + tests := []struct { + name string + want string + }{ + {name: "model.embed_tokens.weight", want: "embedding.weight"}, + {name: "model.embed_tokens.scales", want: "embedding.weight.scale"}, + {name: "model.embed_tokens.biases", want: "embedding.weight.bias"}, + {name: "model.norm.weight", want: "output_norm.weight"}, + {name: "lm_head.weight", want: "output.weight"}, + {name: "lm_head.scales", want: "output.weight.scale"}, + {name: "model.layers.2.input_layernorm.weight", want: "blocks.2.attn_norm.weight"}, + {name: "model.layers.2.self_attn.q_proj.weight", want: "blocks.2.q_proj.weight"}, + {name: "model.layers.2.self_attn.k_proj.bias", want: "blocks.2.k_proj.bias"}, + {name: "model.layers.2.self_attn.sinks", want: "blocks.2.attn_sinks"}, + {name: "model.layers.2.post_attention_layernorm.weight", want: "blocks.2.ffn_norm.weight"}, + {name: "model.layers.2.mlp.router.weight", want: "blocks.2.router.weight"}, + {name: "model.layers.2.mlp.experts.gate_up_proj_blocks", want: "blocks.2.experts.gate_up_proj.weight"}, + {name: "model.layers.2.mlp.experts.gate_up_proj_scales", want: "blocks.2.experts.gate_up_proj.weight"}, + {name: "model.layers.2.mlp.experts.gate_up_proj_bias", want: "blocks.2.experts.gate_up_proj.bias"}, + {name: "model.layers.2.mlp.experts.down_proj_blocks", want: "blocks.2.experts.down_proj.weight"}, + {name: "model.layers.2.mlp.experts.down_proj_scales", want: "blocks.2.experts.down_proj.weight"}, + {name: "model.layers.2.mlp.experts.down_proj_bias", want: "blocks.2.experts.down_proj.bias"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := transform.canonicalTensorName(tt.name); got != tt.want { + t.Fatalf("canonicalTensorName(%q) = %q, want %q", tt.name, got, tt.want) + } + }) + } +} + +func TestGPTOSSImportTransformQuantizationType(t *testing.T) { + transform := &gptossImportTransform{} + if got := transform.quantizationType("blocks.0.experts.gate_proj.weight", []int32{2, 64, 128}, "mxfp4"); got != "" { + t.Fatalf("expert quantizationType = %q, want empty for native expert weights", got) + } + if got := transform.quantizationType("blocks.0.router.weight", []int32{64, 128}, "mxfp4"); got != "" { + t.Fatalf("router quantizationType = %q, want empty", got) + } + if got := transform.quantizationType("blocks.0.q_proj.weight", []int32{2048, 2048}, "mxfp4"); got != "mxfp4" { + t.Fatalf("q_proj quantizationType = %q, want mxfp4", got) + } +} + +func TestClassify_GptOSSMatchingQuantizePreservesNative(t *testing.T) { + inv := gptossInventory(map[string]SourceTensor{ + "model.layers.0.mlp.experts.down_proj_blocks": {Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "U8", Shape: []int32{2, 16, 1, 16}}, + "model.layers.0.mlp.experts.down_proj_scales": {Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 16, 1}}, + "model.layers.0.mlp.experts.down_proj_bias": {Name: "model.layers.0.mlp.experts.down_proj_bias", Dtype: "BF16", Shape: []int32{2, 16}}, + }) + + class, err := Classify(inv, "mxfp4") + if err != nil { + t.Fatalf("Classify() error = %v", err) + } + if class.Kind != SourcePrequantized || class.Quantize != "mxfp4" { + t.Fatalf("Classify() = %+v, want prequantized native preservation", class) + } +} + +func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { + valid := map[string]SourceTensor{ + "model.layers.0.mlp.experts.down_proj_blocks": {Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "U8", Shape: []int32{2, 16, 1, 16}}, + "model.layers.0.mlp.experts.down_proj_scales": {Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 16, 1}}, + "model.layers.0.mlp.experts.down_proj_bias": {Name: "model.layers.0.mlp.experts.down_proj_bias", Dtype: "BF16", Shape: []int32{2, 16}}, + } + for _, requested := range []string{"", "mxfp4"} { + class, err := Classify(gptossInventory(valid), requested) + if err != nil { + t.Fatalf("Classify(requested=%q) error = %v", requested, err) + } + if class.Kind != SourcePrequantized || class.Quantize != "mxfp4" { + t.Fatalf("Classify(requested=%q) = %+v", requested, class) + } + } + + for _, tt := range []struct { + name string + mutate func(*Inventory) + checkNoRequest bool + }{ + {"missing scales", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_scales") }, false}, + {"missing bias", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_bias") }, false}, + {"mismatched scales", func(inv *Inventory) { + inv.Tensors["model.layers.0.mlp.experts.down_proj_scales"] = SourceTensor{Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 15, 1}} + }, false}, + {"wrong blocks dtype", func(inv *Inventory) { + inv.Tensors["model.layers.0.mlp.experts.down_proj_blocks"] = SourceTensor{Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "BF16", Shape: []int32{2, 16, 1, 16}} + }, false}, + {"wrong quant metadata", func(inv *Inventory) { inv.Config.Quantization.Mode = "affine" }, false}, + {"wrong bits", func(inv *Inventory) { inv.Config.Quantization.Bits = 8 }, true}, + {"wrong group", func(inv *Inventory) { inv.Config.Quantization.GroupSize = 64 }, true}, + {"architecture only", func(inv *Inventory) { + clear(inv.Tensors) + inv.Tensors["model.layers.0.weight"] = SourceTensor{Name: "model.layers.0.weight", Dtype: "U32", Shape: []int32{16, 4}} + inv.Tensors["model.layers.0.scales"] = SourceTensor{Name: "model.layers.0.scales", Dtype: "U8", Shape: []int32{16, 1}} + }, false}, + } { + t.Run(tt.name, func(t *testing.T) { + inv := gptossInventory(valid) + inv.Tensors = maps.Clone(inv.Tensors) + tt.mutate(&inv) + if _, err := Classify(inv, "mxfp4"); err == nil { + t.Fatal("unproven native preservation accepted") + } + if tt.checkNoRequest { + class, err := Classify(inv, "") + if err != nil { + t.Fatalf("no-request classification error = %v", err) + } + if class.Kind != SourcePrequantized || class.Quantize != "" { + t.Fatalf("no-request classification = %+v, want unlabelled prequantized source", class) + } + } + }) + } +} + +func TestCreatePipelineReportsGPTOSSMXFP4FileType(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{ + "architectures":["GptOssForCausalLM"], + "quantization":{"quant_method":"mxfp4","mode":"mxfp4","bits":4,"group_size":32} + }`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.layers.0.mlp.experts.down_proj_blocks", "U8", []int32{1, 2, 1, 16}, make([]byte, 32)), + st.NewTensorDataFromBytes("model.layers.0.mlp.experts.down_proj_scales", "U8", []int32{1, 2, 1}, make([]byte, 2)), + st.NewTensorDataFromBytes("model.layers.0.mlp.experts.down_proj_bias", "BF16", []int32{1, 2}, make([]byte, 4)), + }) + + var got Classification + err := Create("gptoss", dir, "mxfp4", newCaptureStore(), func(_ string, _ LayerInfo, _ []LayerInfo, class Classification) error { + got = class + return nil + }, func(string) {}) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + if got.Kind != SourcePrequantized || got.Quantize != "mxfp4" { + t.Fatalf("manifest classification = {%s %q}, want {prequantized mxfp4}", got.Kind, got.Quantize) + } +} + +func TestPlanGPTOSSNativeExperts(t *testing.T) { + inv := gptossInventory(map[string]SourceTensor{ + "model.layers.0.mlp.experts.gate_up_proj_blocks": {Name: "model.layers.0.mlp.experts.gate_up_proj_blocks", Dtype: "U8", Shape: []int32{2, 32, 1, 16}}, + "model.layers.0.mlp.experts.gate_up_proj_scales": {Name: "model.layers.0.mlp.experts.gate_up_proj_scales", Dtype: "U8", Shape: []int32{2, 32, 1}}, + "model.layers.0.mlp.experts.gate_up_proj_bias": {Name: "model.layers.0.mlp.experts.gate_up_proj_bias", Dtype: "BF16", Shape: []int32{2, 32}}, + "model.layers.0.mlp.experts.down_proj_blocks": {Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "U8", Shape: []int32{2, 16, 1, 16}}, + "model.layers.0.mlp.experts.down_proj_scales": {Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 16, 1}}, + "model.layers.0.mlp.experts.down_proj_bias": {Name: "model.layers.0.mlp.experts.down_proj_bias", Dtype: "BF16", Shape: []int32{2, 16}}, + "model.layers.0.mlp.router.weight": {Name: "model.layers.0.mlp.router.weight", Dtype: "BF16", Shape: []int32{2, 2}}, + }) + policy, err := newTensorImportTransform(inv) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, policy) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + experts, ok := specByName(specs, "blocks.0.experts") + if !ok { + t.Fatalf("missing blocks.0.experts spec; got %v", specNames(specs)) + } + if experts.Metadata["quant_type"] != "mxfp4" || experts.Metadata["group_size"] != "32" { + t.Fatalf("expert metadata = %v, want mxfp4 group_size=32", experts.Metadata) + } + + got := tensorNames(experts.Tensors) + want := []string{ + "blocks.0.experts.down_proj.bias", + "blocks.0.experts.down_proj.weight", + "blocks.0.experts.down_proj.weight.scale", + "blocks.0.experts.gate_proj.bias", + "blocks.0.experts.gate_proj.weight", + "blocks.0.experts.gate_proj.weight.scale", + "blocks.0.experts.up_proj.bias", + "blocks.0.experts.up_proj.weight", + "blocks.0.experts.up_proj.weight.scale", + } + slices.Sort(got) + if !slices.Equal(got, want) { + t.Fatalf("expert tensor names = %v, want %v", got, want) + } + + router, ok := specByName(specs, "blocks.0.router.weight") + if !ok { + t.Fatalf("missing router spec; got %v", specNames(specs)) + } + if router.Tensors[0].Quantize != "" { + t.Fatalf("router Quantize = %q, want empty", router.Tensors[0].Quantize) + } +} + +func TestPlanGPTOSSNativeDenseQuantizedBlob(t *testing.T) { + inv := gptossInventory(map[string]SourceTensor{ + "model.layers.0.self_attn.q_proj.weight": {Name: "model.layers.0.self_attn.q_proj.weight", Dtype: "U32", Shape: []int32{64, 8}}, + "model.layers.0.self_attn.q_proj.scales": {Name: "model.layers.0.self_attn.q_proj.scales", Dtype: "U8", Shape: []int32{64, 2}}, + "model.layers.0.self_attn.q_proj.biases": {Name: "model.layers.0.self_attn.q_proj.biases", Dtype: "BF16", Shape: []int32{64, 2}}, + }) + policy, err := newTensorImportTransform(inv) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, policy) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + spec, ok := specByName(specs, "blocks.0.q_proj.weight") + if !ok { + t.Fatalf("missing q_proj spec; got %v", specNames(specs)) + } + got := tensorNames(spec.Tensors) + want := []string{"blocks.0.q_proj.weight", "blocks.0.q_proj.weight.bias", "blocks.0.q_proj.weight.scale"} + slices.Sort(got) + if !slices.Equal(got, want) { + t.Fatalf("q_proj tensor names = %v, want %v", got, want) + } + if spec.Metadata["quant_type"] != "mxfp4" || spec.Metadata["group_size"] != "32" { + t.Fatalf("q_proj metadata = %v, want mxfp4 group_size=32", spec.Metadata) + } +} + +func TestGPTOSSNativeGateUpSplitDequantizesLikeOriginal(t *testing.T) { + raw := []byte{ + 0x10, 0x32, 0x54, 0x76, 0x98, 0xba, 0xdc, 0xfe, 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, + 0x11, 0x33, 0x55, 0x77, 0x99, 0xbb, 0xdd, 0xff, 0x00, 0x22, 0x44, 0x66, 0x88, 0xaa, 0xcc, 0xee, + 0x89, 0x67, 0x45, 0x23, 0x01, 0xef, 0xcd, 0xab, 0x98, 0x76, 0x54, 0x32, 0x10, 0xfe, 0xdc, 0xba, + 0xff, 0xee, 0xdd, 0xcc, 0xbb, 0xaa, 0x99, 0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11, 0x00, + } + scales := []byte{0x10, 0x11, 0x12, 0x13} + + whole, err := dequantizeGPTOSSMXFP4Tensor( + "blocks.0.experts.gate_up_proj.weight", + st.NewTensorDataFromBytes("blocks.0.experts.gate_up_proj.weight", "U8", []int32{1, 4, 1, 16}, raw), + st.NewTensorDataFromBytes("blocks.0.experts.gate_up_proj.weight", "U8", []int32{1, 4, 1}, scales), + ) + if err != nil { + t.Fatalf("dequantize whole gate_up tensor: %v", err) + } + wholeVals := mustDecodeBF16Tensor(t, whole) + + out, err := preserveAndSplitGateUpTensor( + "blocks.0.experts.gate_up_proj.weight", + st.NewTensorDataFromBytes("blocks.0.experts.gate_up_proj.weight", "U8", []int32{1, 4, 1, 16}, raw), + st.NewTensorDataFromBytes("blocks.0.experts.gate_up_proj.weight", "U8", []int32{1, 4, 1}, scales), + ) + if err != nil { + t.Fatalf("preserveAndSplitGateUpTensor() error = %v", err) + } + + gateRaw, err := io.ReadAll(out[0].Reader()) + if err != nil { + t.Fatalf("read gate rows: %v", err) + } + gateScaleRaw, err := io.ReadAll(out[1].Reader()) + if err != nil { + t.Fatalf("read gate scales: %v", err) + } + upRaw, err := io.ReadAll(out[2].Reader()) + if err != nil { + t.Fatalf("read up rows: %v", err) + } + upScaleRaw, err := io.ReadAll(out[3].Reader()) + if err != nil { + t.Fatalf("read up scales: %v", err) + } + + evenVals := decodeNativePackedMXFP4Values(t, gateRaw, gateScaleRaw) + oddVals := decodeNativePackedMXFP4Values(t, upRaw, upScaleRaw) + for i := range 2 { + for j := range 32 { + if wholeVals[(i*2)*32+j] != evenVals[i*32+j] { + t.Fatalf("even row %d col %d mismatch: whole=%v split=%v", i, j, wholeVals[(i*2)*32+j], evenVals[i*32+j]) + } + if wholeVals[(i*2+1)*32+j] != oddVals[i*32+j] { + t.Fatalf("odd row %d col %d mismatch: whole=%v split=%v", i, j, wholeVals[(i*2+1)*32+j], oddVals[i*32+j]) + } + } + } +} + +func gptossInventory(tensors map[string]SourceTensor) Inventory { + if tensors == nil { + tensors = map[string]SourceTensor{} + } + return Inventory{ + Config: sourceModelConfig{ + Architectures: []string{"GptOssForCausalLM"}, + Quantization: sourceQuantization{QuantMethod: "mxfp4", Mode: "mxfp4", Bits: 4, GroupSize: 32}, + }, + RawConfig: []byte(`{ + "architectures": ["GptOssForCausalLM"], + "quantization": {"quant_method": "mxfp4", "bits": 4, "group_size": 32} + }`), + Tensors: tensors, + } +} + +func tensorNames(tensors []TensorSpec) []string { + names := make([]string, 0, len(tensors)) + for _, tensor := range tensors { + names = append(names, tensor.Name) + } + return names +} + +func decodeNativePackedMXFP4Values(t *testing.T, blocks, scales []byte) []float32 { + t.Helper() + if len(blocks)%16 != 0 { + t.Fatalf("native block byte length = %d, want multiple of 16", len(blocks)) + } + groups := len(blocks) / 16 + if len(scales) != groups { + t.Fatalf("native scale byte length = %d, want %d", len(scales), groups) + } + values := make([]float32, groups*32) + for i := range groups { + scale := decodeGPTOSSMXFP4Scale(scales[i]) + for j, packed := range blocks[i*16 : (i+1)*16] { + values[i*32+2*j] = gptossMXFP4Values[packed&0x0F] * scale + values[i*32+2*j+1] = gptossMXFP4Values[packed>>4] * scale + } + } + return values +} + +func mustDecodeBF16Tensor(t *testing.T, td *st.TensorData) []float32 { + t.Helper() + raw, err := io.ReadAll(td.Reader()) + if err != nil { + t.Fatalf("read tensor: %v", err) + } + values, err := DecodeFloatTensor(td.Dtype, raw) + if err != nil { + t.Fatalf("decode tensor: %v", err) + } + return values +} diff --git a/x/create/plan.go b/x/create/plan.go index 35379bbd4bf..c7bf2c36b5b 100644 --- a/x/create/plan.go +++ b/x/create/plan.go @@ -104,6 +104,14 @@ type quantizePolicy interface { // weights are quantized and to what; pass defaultQuantPolicy{} for the generic // policy. func Plan(inv Inventory, class Classification, policy quantizePolicy) ([]BlobSpec, error) { + if inv.Config.Architecture() == "GptOssForCausalLM" { + specs, err := planGPTOSS(inv, class, policy) + if err != nil { + return nil, err + } + return specs, checkOutputCollisions(specs) + } + var ( specs []BlobSpec err error From f7029e8d284c5b518152a9d3c4143fbc71623592 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 18:55:39 +0000 Subject: [PATCH 24/58] create: preserve native GPT-OSS MXFP4 snapshots Co-authored-by: Codex --- x/create/classify.go | 14 -------------- x/create/gptoss_test.go | 17 ++++++++++++++--- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/x/create/classify.go b/x/create/classify.go index a11b453e710..93ebf0da714 100644 --- a/x/create/classify.go +++ b/x/create/classify.go @@ -166,20 +166,6 @@ func gptossNativeMXFP4Evidence(inv Inventory) bool { if inv.Config.Architecture() != "GptOssForCausalLM" { return false } - var selected sourceQuantization - for _, candidate := range inv.Config.quantizationConfigs() { - if candidate.Bits != 0 { - selected = candidate - break - } - } - mode := selected.Mode - if mode == "" { - mode = selected.QuantMethod - } - if selected.Bits != 4 || selected.GroupSize != 32 || !strings.EqualFold(mode, "mxfp4") { - return false - } found := false for _, name := range sortedTensorNames(inv) { diff --git a/x/create/gptoss_test.go b/x/create/gptoss_test.go index 51badb5e54f..d4e4f17ed9d 100644 --- a/x/create/gptoss_test.go +++ b/x/create/gptoss_test.go @@ -103,6 +103,20 @@ func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { } } + // Native GPT-OSS checkpoints have shipped with quantization metadata in + // more than one shape. The packed expert layout, rather than an optional + // metadata spelling, is the import contract. + withoutMetadata := gptossInventory(valid) + withoutMetadata.Config.Quantization = sourceQuantization{} + withoutMetadata.RawConfig = []byte(`{"architectures":["GptOssForCausalLM"]}`) + class, err := Classify(withoutMetadata, "mxfp4") + if err != nil { + t.Fatalf("Classify(native layout without quantization metadata) error = %v", err) + } + if class.Kind != SourcePrequantized || class.Quantize != "mxfp4" { + t.Fatalf("Classify(native layout without quantization metadata) = %+v", class) + } + for _, tt := range []struct { name string mutate func(*Inventory) @@ -116,9 +130,6 @@ func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { {"wrong blocks dtype", func(inv *Inventory) { inv.Tensors["model.layers.0.mlp.experts.down_proj_blocks"] = SourceTensor{Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "BF16", Shape: []int32{2, 16, 1, 16}} }, false}, - {"wrong quant metadata", func(inv *Inventory) { inv.Config.Quantization.Mode = "affine" }, false}, - {"wrong bits", func(inv *Inventory) { inv.Config.Quantization.Bits = 8 }, true}, - {"wrong group", func(inv *Inventory) { inv.Config.Quantization.GroupSize = 64 }, true}, {"architecture only", func(inv *Inventory) { clear(inv.Tensors) inv.Tensors["model.layers.0.weight"] = SourceTensor{Name: "model.layers.0.weight", Dtype: "U32", Shape: []int32{16, 4}} From 64f405f798ff1e2324af159bd2519e5dead41a0d Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 14:14:24 +0000 Subject: [PATCH 25/58] create: restore GPT-OSS family metadata Co-authored-by: Codex --- x/create/client/create.go | 32 +++++++++++++++++++++++++ x/create/client/create_test.go | 43 ++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/x/create/client/create.go b/x/create/client/create.go index e3b87db74d9..64400b9c561 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -407,6 +407,12 @@ func newManifestWriter(opts CreateOptions, capabilities []string, parserName, re } configData.Draft = draft } + if configData.ModelFamily == "" { + configData.ModelFamily = inferModelFamily(opts.ModelDir) + } + if configData.ModelFamily != "" && len(configData.ModelFamilies) == 0 { + configData.ModelFamilies = []string{configData.ModelFamily} + } configJSON, err := json.Marshal(configData) if err != nil { return fmt.Errorf("failed to marshal config: %w", err) @@ -442,6 +448,32 @@ func newManifestWriter(opts CreateOptions, capabilities []string, parserName, re } } +func inferModelFamily(modelDir string) string { + data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) + if err != nil { + return "" + } + + var cfg struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + LLMConfig struct { + ModelType string `json:"model_type"` + } `json:"llm_config"` + } + if err := json.Unmarshal(data, &cfg); err != nil { + return "" + } + + for _, identifier := range append(cfg.Architectures, cfg.ModelType, cfg.LLMConfig.ModelType) { + if isGPTOSSFamily(identifier) { + return "gptoss" + } + } + + return "" +} + func resolveParserName(mf *ModelfileConfig, inferred string) string { if mf != nil && mf.Parser != "" { return mf.Parser diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index e48950f63e5..c3f07297439 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -554,6 +554,49 @@ func TestNewManifestWriter_PopulatesFileTypeFromEffectiveQuantize(t *testing.T) } } +func TestNewManifestWriter_PopulatesGPTOSSFamily(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + modelDir := t.TempDir() + if err := os.WriteFile(filepath.Join(modelDir, "config.json"), []byte(`{ + "architectures": ["GptOssForCausalLM"], + "model_type": "gpt_oss" + }`), 0o644); err != nil { + t.Fatal(err) + } + + opts := CreateOptions{ModelName: "gptoss-family-test", ModelDir: modelDir} + writer := newManifestWriter(opts, []string{"completion", "tools", "thinking"}, "harmony", "") + if err := writer(opts.ModelName, create.LayerInfo{}, nil, create.Classification{Kind: create.SourcePrequantized, Quantize: "mxfp4"}); err != nil { + t.Fatalf("newManifestWriter() error = %v", err) + } + + name := model.ParseName(opts.ModelName) + mf, err := manifest.ParseNamedManifest(name) + if err != nil { + t.Fatal(err) + } + configPath, err := manifest.BlobsPath(mf.Config.Digest) + if err != nil { + t.Fatal(err) + } + data, err := os.ReadFile(configPath) + if err != nil { + t.Fatal(err) + } + + var cfg model.ConfigV2 + if err := json.Unmarshal(data, &cfg); err != nil { + t.Fatal(err) + } + if cfg.ModelFamily != "gptoss" { + t.Fatalf("ModelFamily = %q, want gptoss", cfg.ModelFamily) + } + if !slices.Equal(cfg.ModelFamilies, []string{"gptoss"}) { + t.Fatalf("ModelFamilies = %v, want [gptoss]", cfg.ModelFamilies) + } +} + func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) From 111d0e05396308fc6742d7d4874043f816a79d8b Mon Sep 17 00:00:00 2001 From: Philipp Date: Wed, 26 Aug 2026 17:10:00 +0000 Subject: [PATCH 26/58] mlx: avoid GPT-OSS MoE selector synchronization Guard dynamic expert selectors inside the fused Metal kernels so decode no longer reads router output back to the CPU on every layer. Preserve static validation and cover mixed valid and invalid selectors across multiple threadgroups. Co-authored-by: Codex --- x/mlxrunner/mlx/gptoss_moe.go | 56 +++++++++--------- x/mlxrunner/mlx/gptoss_moe_test.go | 91 +++++++++++++++++++----------- 2 files changed, 89 insertions(+), 58 deletions(-) diff --git a/x/mlxrunner/mlx/gptoss_moe.go b/x/mlxrunner/mlx/gptoss_moe.go index 46aee9142aa..fac52b3ef63 100644 --- a/x/mlxrunner/mlx/gptoss_moe.go +++ b/x/mlxrunner/mlx/gptoss_moe.go @@ -63,23 +63,6 @@ func validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK int, down b return ok } -func validateMoEExpertIDs(expertIds *Array, experts int) bool { - if expertIds == nil || expertIds.DType() != DTypeUint32 || expertIds.Size() == 0 || !validMoEPositiveInt(experts) { - return false - } - Eval(expertIds) - data := C.mlx_array_data_uint32(expertIds.ctx) - if data == nil { - return false - } - for _, expertID := range unsafe.Slice(data, expertIds.Size()) { - if uint64(expertID) >= uint64(experts) { - return false - } - } - return true -} - func validMoEArray(a *Array, dtype DType, dims ...int) bool { if a == nil || a.DType() != dtype || len(dims) != a.NumDims() { return false @@ -139,8 +122,7 @@ func validateMoEGateUpInputs( !validMoEArray(expertIds, DTypeUint32, batch, topK) { return 0, false } - if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, false) || - !validateMoEExpertIDs(expertIds, experts) { + if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, false) { return 0, false } return batch, true @@ -163,8 +145,7 @@ func validateMoEDownInputs(input, weight, scales, bias, expertIds *Array, numRow !validMoEArray(expertIds, DTypeUint32, batch, topK) { return 0, false } - if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, true) || - !validateMoEExpertIDs(expertIds, experts) { + if !validateMoEIndexSpans(experts, batch, numRows, numColVecs, topK, true) { return 0, false } return batch, true @@ -187,7 +168,7 @@ var ( // odd compute up). After a barrier, simdgroup pairs apply SwiGLU to produce // nSg/2 output channels per threadgroup. // -// Template args: NumColVecs (int), NumRows (int), NumTopK (int) +// Template args: NumExperts (int), NumColVecs (int), NumRows (int), NumTopK (int) // Inputs (named): input(float), gate_w(uint), gate_s(bfloat), gate_b(bfloat), // // up_w(uint), up_s(bfloat), up_b(bfloat), @@ -205,6 +186,15 @@ threadgroup float tg_buf[32]; uint outCh = gid.x * (nSg / 2) + (sgIdx / 2); uint expertId = expert_ids[gid.y * NumTopK + gid.z]; +uint tid = sgIdx * sgSize + sgTid; + +if (expertId >= (uint)NumExperts) { + if (tid * 2 < nSg) { + uint ch = gid.x * (nSg / 2) + tid; + output[gid.y * (uint)NumTopK * (uint)NumRows + gid.z * (uint)NumRows + ch] = 0.0f; + } + return; +} uint rowStride = (uint)NumColVecs; uint expertStride = (uint)NumRows * rowStride; @@ -298,7 +288,6 @@ if (simd_is_first()) { } threadgroup_barrier(mem_flags::mem_threadgroup); -uint tid = sgIdx * sgSize + sgTid; if (tid * 2 < nSg) { const float2 gu = reinterpret_cast(tg_buf)[tid]; const float smin = swiglu_params[0]; @@ -355,7 +344,10 @@ func initMoESwiGLUKernel() { // MoEFusedGateUpSwiGLU runs a fused gate+up+SwiGLU kernel for MXFP4 MoE experts. // It returns (result, true) on success, or (nil, false) if the kernel is unavailable -// or the inputs are incompatible. +// or the static input metadata is incompatible. Dynamic expert IDs are expected +// to come from the model's bounded routing operation. An out-of-range ID is +// guarded on the device and produces zeros for that selector's output rows; +// the call still returns (result, true). // // Parameters: // - input: float32 [batch, hiddenSize] — flattened input hidden states @@ -410,6 +402,7 @@ func MoEFusedGateUpSwiGLU( name string value int }{ + {"NumExperts", gateWeight.Dim(0)}, {"NumColVecs", numColVecs}, {"NumRows", numRows}, {"NumTopK", topK}, @@ -489,7 +482,7 @@ func MoEFusedGateUpSwiGLU( // The input is per-expert: input[gid.y * NumTopK * NumInCols + gid.z * NumInCols + ...] // where NumInCols is the intermediate_size (in float32 elements). // -// Template args: NumColVecs (int), NumRows (int), NumTopK (int) +// Template args: NumExperts (int), NumColVecs (int), NumRows (int), NumTopK (int) // Inputs (named): input(float), down_w(uint), down_s(bfloat), down_b(bfloat), // // expert_ids(uint) @@ -505,6 +498,13 @@ constexpr uint sgSize = 32; uint outCh = gid.x * nSg + sgIdx; uint expertId = expert_ids[gid.y * NumTopK + gid.z]; +if (expertId >= (uint)NumExperts) { + if (simd_is_first()) { + output[gid.y * (uint)NumTopK * (uint)NumRows + gid.z * (uint)NumRows + outCh] = 0.0f; + } + return; +} + uint rowStride = (uint)NumColVecs; uint expertStride = (uint)NumRows * rowStride; @@ -621,7 +621,10 @@ func initMoEDownKernel() { // MoEFusedDown runs a fused down-projection kernel for MXFP4 MoE experts. // It returns (result, true) on success, or (nil, false) if the kernel is unavailable -// or the inputs are incompatible. +// or the static input metadata is incompatible. Dynamic expert IDs are expected +// to come from the model's bounded routing operation. An out-of-range ID is +// guarded on the device and produces zeros for that selector's output rows; +// the call still returns (result, true). // // Parameters: // - input: float32 [batch, topK, intermediateSize] — per-expert SwiGLU output @@ -670,6 +673,7 @@ func MoEFusedDown( name string value int }{ + {"NumExperts", downWeight.Dim(0)}, {"NumColVecs", numColVecs}, {"NumRows", numRows}, {"NumTopK", topK}, diff --git a/x/mlxrunner/mlx/gptoss_moe_test.go b/x/mlxrunner/mlx/gptoss_moe_test.go index 75c2cc59f57..3b6a4e3f769 100644 --- a/x/mlxrunner/mlx/gptoss_moe_test.go +++ b/x/mlxrunner/mlx/gptoss_moe_test.go @@ -60,7 +60,11 @@ type moeTestInputs struct { } func newMoETestInputs(expertIDs []uint32) moeTestInputs { - const experts, rows, colVecs, topK = 2, 4, 32, 1 + return newMoETestInputsWithRows(expertIDs, 4) +} + +func newMoETestInputsWithRows(expertIDs []uint32, rows int) moeTestInputs { + const experts, colVecs, topK = 2, 32, 1 batch := len(expertIDs) / topK return moeTestInputs{ input: Zeros(DTypeFloat32, batch, colVecs*32), @@ -74,6 +78,23 @@ func newMoETestInputs(expertIDs []uint32) moeTestInputs { } } +// validateMoEExpertIDsForTest deliberately synchronizes and reads selector +// values on the CPU. Keep it test-only so it cannot be reintroduced into the +// lazy decode graph. +func validateMoEExpertIDsForTest(expertIDs *Array, experts int) bool { + if expertIDs == nil || expertIDs.DType() != DTypeUint32 || expertIDs.Size() == 0 || !validMoEPositiveInt(experts) { + return false + } + ids := expertIDs.AsType(DTypeInt32) + Eval(ids) + for _, expertID := range ids.Ints() { + if expertID < 0 || int64(expertID) >= int64(experts) { + return false + } + } + return true +} + func TestValidateMoEInputs(t *testing.T) { skipIfNoMLX(t) withMLXThread(t, func() { @@ -86,21 +107,17 @@ func TestValidateMoEInputs(t *testing.T) { if _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1); !ok { t.Fatal("valid down inputs rejected") } + if validateMoEExpertIDsForTest(newMoETestInputs([]uint32{2}).expertIDs, 2) { + t.Fatal("expert equal to count accepted by explicit value validation") + } + if validateMoEExpertIDsForTest(newMoETestInputs([]uint32{math.MaxUint32}).expertIDs, 2) { + t.Fatal("maximum expert accepted by explicit value validation") + } for _, tt := range []struct { name string fn func() bool }{ - {name: "expert equal to count", fn: func() bool { - bad := newMoETestInputs([]uint32{2}) - _, ok := validateMoEGateUpInputs(bad.input, bad.weight, bad.scales, bad.bias, bad.upWeight, bad.upScales, bad.upBias, bad.expertIDs, 4, 32, 1) - return ok - }}, - {name: "maximum expert", fn: func() bool { - bad := newMoETestInputs([]uint32{math.MaxUint32}) - _, ok := validateMoEDownInputs(downInput, bad.weight, bad.scales, bad.bias, bad.expertIDs, 4, 32, 1) - return ok - }}, {name: "wrong selector dtype", fn: func() bool { ids := FromValues([]int32{0}, 1, 1) _, ok := validateMoEDownInputs(downInput, valid.weight, valid.scales, valid.bias, ids, 4, 32, 1) @@ -172,41 +189,43 @@ func TestValidateMoEInputs(t *testing.T) { }) } -func TestMoEFusedZeroReferenceAndValidationIsolation(t *testing.T) { +func TestMoEFusedBoundsGuardMixedSelectors(t *testing.T) { skipIfNoMLX(t) withMLXThread(t, func() { if !MetalIsAvailable() { t.Skip("Metal is not available") } - bad := newMoETestInputs([]uint32{2}) - if out, ok := MoEFusedGateUpSwiGLU(bad.input, bad.weight, bad.scales, bad.bias, bad.upWeight, bad.upScales, bad.upBias, bad.expertIDs, 4, 32, 1, -7, 7); ok || out != nil { - t.Fatal("out-of-range gate/up selector accepted") + const rows = 16 + selectors := []uint32{0, 1, 2, math.MaxUint32} + inputs := newMoETestInputsWithRows(selectors, rows) + biasValues := make([]float32, 2*rows) + for i := range biasValues { + biasValues[i] = 1 } + inputs.bias = FromValues(biasValues, 2, rows).AsType(DTypeBFloat16) + inputs.upBias = FromValues(biasValues, 2, rows).AsType(DTypeBFloat16) - valid := newMoETestInputs([]uint32{1}) - gateOut, ok := MoEFusedGateUpSwiGLU(valid.input, valid.weight, valid.scales, valid.bias, valid.upWeight, valid.upScales, valid.upBias, valid.expertIDs, 4, 32, 1, -7, 7) - if !ok { - t.Fatal("valid gate/up call failed after malformed call") + gateOut, ok := MoEFusedGateUpSwiGLU(inputs.input, inputs.weight, inputs.scales, inputs.bias, inputs.upWeight, inputs.upScales, inputs.upBias, inputs.expertIDs, rows, 32, 1, -7, 7) + if !ok || gateOut == nil { + t.Fatal("mixed-selector gate/up call failed") } Eval(gateOut) - assertMoEZeroReference(t, gateOut, []int{1, 1, 4}) + assertMoEGuardedRows(t, gateOut, rows, []bool{true, true, false, false}) - downInput := Zeros(DTypeFloat32, 1, 1, 32*32) - if out, ok := MoEFusedDown(downInput, valid.weight, valid.scales, valid.bias, bad.expertIDs, 4, 32, 1); ok || out != nil { - t.Fatal("out-of-range down selector accepted") - } - downOut, ok := MoEFusedDown(downInput, valid.weight, valid.scales, valid.bias, valid.expertIDs, 4, 32, 1) - if !ok { - t.Fatal("valid down call failed after malformed call") + downInput := Zeros(DTypeFloat32, len(selectors), 1, 32*32) + downOut, ok := MoEFusedDown(downInput, inputs.weight, inputs.scales, inputs.bias, inputs.expertIDs, rows, 32, 1) + if !ok || downOut == nil { + t.Fatal("mixed-selector down call failed") } Eval(downOut) - assertMoEZeroReference(t, downOut, []int{1, 1, 4}) + assertMoEGuardedRows(t, downOut, rows, []bool{true, true, false, false}) }) } -func assertMoEZeroReference(t *testing.T, got *Array, wantShape []int) { +func assertMoEGuardedRows(t *testing.T, got *Array, rows int, wantNonzero []bool) { t.Helper() + wantShape := []int{len(wantNonzero), 1, rows} if dims := got.Dims(); len(dims) != len(wantShape) { t.Fatalf("shape = %v, want %v", dims, wantShape) } else { @@ -216,9 +235,17 @@ func assertMoEZeroReference(t *testing.T, got *Array, wantShape []int) { } } } - for i, value := range got.Floats() { - if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) || math.Abs(float64(value)) > 1e-6 { - t.Fatalf("output[%d] = %v, want zero reference within 1e-6", i, value) + for batch, nonzero := range wantNonzero { + for row, value := range got.Floats()[batch*rows : (batch+1)*rows] { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + t.Fatalf("output[%d,%d] = %v, want finite", batch, row, value) + } + if nonzero && math.Abs(float64(value)) <= 1e-4 { + t.Fatalf("output[%d,%d] = %v, want nonzero valid-selector output", batch, row, value) + } + if !nonzero && math.Abs(float64(value)) > 1e-6 { + t.Fatalf("output[%d,%d] = %v, want zero guarded output", batch, row, value) + } } } } From 026678587b75c42ab4939b03654aa9af9c97ed0a Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 06:59:27 +0000 Subject: [PATCH 27/58] create: canonicalize native GPT-OSS expert metadata Co-authored-by: Codex --- x/create/gptoss.go | 6 +++++- x/create/gptoss_test.go | 43 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/x/create/gptoss.go b/x/create/gptoss.go index 71f492f01eb..5c6dcec3df0 100644 --- a/x/create/gptoss.go +++ b/x/create/gptoss.go @@ -304,7 +304,11 @@ func (t *gptossImportTransform) planNativeExpertTensor(inv Inventory, blocksName blocks := inv.Tensors[blocksName] scales := inv.Tensors[scalesName] sources := []string{blocksName, scalesName} - metadata := t.prequantizedMetadata(strings.TrimSuffix(blocksName, "_blocks")+".weight", t.defaultMetadata()) + // The native _blocks/_scales layout is the authoritative direct-expert + // contract. Classification accepts it even when config.json is missing or + // stale, so never let global or per-tensor affine metadata relabel the + // preserved MXFP4 bytes into a runtime-incompatible format. + metadata := map[string]string{"quant_type": "mxfp4", "group_size": "32"} if strings.Contains(outName, ".gate_up_proj.weight") { gateWeight := strings.Replace(outName, "gate_up_proj", "gate_proj", 1) upWeight := strings.Replace(outName, "gate_up_proj", "up_proj", 1) diff --git a/x/create/gptoss_test.go b/x/create/gptoss_test.go index d4e4f17ed9d..63e38933000 100644 --- a/x/create/gptoss_test.go +++ b/x/create/gptoss_test.go @@ -235,6 +235,49 @@ func TestPlanGPTOSSNativeExperts(t *testing.T) { } } +func TestPlanGPTOSSNativeExpertsCanonicalizesStaleQuantizationMetadata(t *testing.T) { + for _, requested := range []string{"", "mxfp4"} { + t.Run("requested="+requested, func(t *testing.T) { + inv := gptossInventory(map[string]SourceTensor{ + "model.layers.0.mlp.experts.down_proj_blocks": {Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "U8", Shape: []int32{2, 16, 1, 16}}, + "model.layers.0.mlp.experts.down_proj_scales": {Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 16, 1}}, + "model.layers.0.mlp.experts.down_proj_bias": {Name: "model.layers.0.mlp.experts.down_proj_bias", Dtype: "BF16", Shape: []int32{2, 16}}, + }) + inv.Config.Quantization = sourceQuantization{QuantMethod: "affine", Mode: "affine", Bits: 4, GroupSize: 64} + inv.RawConfig = []byte(`{ + "architectures":["GptOssForCausalLM"], + "quantization":{ + "quant_method":"affine","mode":"affine","bits":4,"group_size":64, + "model.layers.0.mlp.experts.down_proj":{"mode":"affine","bits":4,"group_size":64} + } + }`) + + class, err := Classify(inv, requested) + if err != nil { + t.Fatalf("Classify() error = %v", err) + } + if class.Kind != SourcePrequantized || class.Quantize != "mxfp4" { + t.Fatalf("Classify() = %+v, want native mxfp4", class) + } + policy, err := newTensorImportTransform(inv) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + specs, err := Plan(inv, class, policy) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + experts, ok := specByName(specs, "blocks.0.experts") + if !ok { + t.Fatalf("missing expert blob; got %v", specNames(specs)) + } + if got := experts.Metadata; got["quant_type"] != "mxfp4" || got["group_size"] != "32" { + t.Fatalf("expert metadata = %v, want runtime-loadable mxfp4/group32", got) + } + }) + } +} + func TestPlanGPTOSSNativeDenseQuantizedBlob(t *testing.T) { inv := gptossInventory(map[string]SourceTensor{ "model.layers.0.self_attn.q_proj.weight": {Name: "model.layers.0.self_attn.q_proj.weight", Dtype: "U32", Shape: []int32{64, 8}}, From 761375c6d13c0d36f20a0681f8877369d5b8f6bd Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 07:01:04 +0000 Subject: [PATCH 28/58] create: require GPT-OSS native expert biases Co-authored-by: Codex --- x/create/classify.go | 19 +++++++++++++++++++ x/create/gptoss.go | 28 +++++++++++++--------------- x/create/gptoss_test.go | 41 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 15 deletions(-) diff --git a/x/create/classify.go b/x/create/classify.go index 93ebf0da714..976663e8450 100644 --- a/x/create/classify.go +++ b/x/create/classify.go @@ -55,6 +55,9 @@ func Classify(inv Inventory, requested string) (Classification, error) { return Classification{Kind: SourceFloat, Quantize: requested}, nil case SourcePrequantized: + if biasName, missing := missingGPTOSSNativeExpertBias(inv); missing { + return Classification{}, fmt.Errorf("incomplete GPT-OSS native MXFP4 checkpoint: required expert bias %q is missing", biasName) + } effective := detectPrequantizedQuantization(inv) if gptossNativeMXFP4Evidence(inv) { effective = "mxfp4" @@ -87,6 +90,22 @@ func Classify(inv Inventory, requested string) (Classification, error) { return Classification{}, fmt.Errorf("could not classify source model in %s", inv.Dir) } +func missingGPTOSSNativeExpertBias(inv Inventory) (string, bool) { + if inv.Config.Architecture() != "GptOssForCausalLM" { + return "", false + } + for _, name := range sortedTensorNames(inv) { + if !strings.Contains(name, ".experts.") || !strings.HasSuffix(name, "_blocks") { + continue + } + biasName := strings.TrimSuffix(name, "_blocks") + "_bias" + if !inv.Has(biasName) { + return biasName, true + } + } + return "", false +} + // detectPrequantizedQuantization reports a single quantization shared by the // source's recognized prequantized weights. A mixed or unrecognized source has // no single file type to record in the model manifest. diff --git a/x/create/gptoss.go b/x/create/gptoss.go index 5c6dcec3df0..d73257007e2 100644 --- a/x/create/gptoss.go +++ b/x/create/gptoss.go @@ -303,7 +303,12 @@ func (t *gptossImportTransform) planNativeExpertTensor(inv Inventory, blocksName blocks := inv.Tensors[blocksName] scales := inv.Tensors[scalesName] - sources := []string{blocksName, scalesName} + biasName := strings.TrimSuffix(blocksName, "_blocks") + "_bias" + if !inv.Has(biasName) { + return "", nil, nil, nil, false, fmt.Errorf("incomplete GPT-OSS native MXFP4 checkpoint: required expert bias %q is missing", biasName) + } + bias := inv.Tensors[biasName] + sources := []string{blocksName, scalesName, biasName} // The native _blocks/_scales layout is the authoritative direct-expert // contract. Classification accepts it even when config.json is missing or // stale, so never let global or per-tensor affine metadata relabel the @@ -318,16 +323,12 @@ func (t *gptossImportTransform) planNativeExpertTensor(inv Inventory, blocksName {Name: upWeight, Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSUpWeight}, {Name: upWeight + ".scale", Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSUpScale}, } - if biasName := strings.TrimSuffix(blocksName, "_blocks") + "_bias"; inv.Has(biasName) { - bias := inv.Tensors[biasName] - gateBias := strings.Replace(outName, "gate_up_proj.weight", "gate_proj.bias", 1) - upBias := strings.Replace(outName, "gate_up_proj.weight", "up_proj.bias", 1) - tensors = append(tensors, - TensorSpec{Name: gateBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSGateUpBias}, - TensorSpec{Name: upBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSUpBias}, - ) - sources = append(sources, biasName) - } + gateBias := strings.Replace(outName, "gate_up_proj.weight", "gate_proj.bias", 1) + upBias := strings.Replace(outName, "gate_up_proj.weight", "up_proj.bias", 1) + tensors = append(tensors, + TensorSpec{Name: gateBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSGateUpBias}, + TensorSpec{Name: upBias, Sources: []SourceTensor{bias}, Transform: TransformGPTOSSUpBias}, + ) return group, tensors, metadata, sources, true, nil } @@ -335,10 +336,7 @@ func (t *gptossImportTransform) planNativeExpertTensor(inv Inventory, blocksName {Name: outName, Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSPackedExpertWeight}, {Name: outName + ".scale", Sources: []SourceTensor{blocks, scales}, Transform: TransformGPTOSSPackedExpertScale}, } - if biasName := strings.TrimSuffix(blocksName, "_blocks") + "_bias"; inv.Has(biasName) { - tensors = append(tensors, TensorSpec{Name: strings.Replace(outName, ".weight", ".bias", 1), Sources: []SourceTensor{inv.Tensors[biasName]}}) - sources = append(sources, biasName) - } + tensors = append(tensors, TensorSpec{Name: strings.Replace(outName, ".weight", ".bias", 1), Sources: []SourceTensor{bias}}) return group, tensors, metadata, sources, true, nil } diff --git a/x/create/gptoss_test.go b/x/create/gptoss_test.go index 63e38933000..e362c8c7933 100644 --- a/x/create/gptoss_test.go +++ b/x/create/gptoss_test.go @@ -5,6 +5,7 @@ import ( "maps" "path/filepath" "slices" + "strings" "testing" st "github.com/ollama/ollama/x/safetensors" @@ -156,6 +157,46 @@ func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { } } +func TestGPTOSSNativeMXFP4MissingBiasIsRejected(t *testing.T) { + inv := gptossInventory(map[string]SourceTensor{ + "model.layers.0.mlp.experts.down_proj_blocks": {Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "U8", Shape: []int32{2, 16, 1, 16}}, + "model.layers.0.mlp.experts.down_proj_scales": {Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 16, 1}}, + }) + + for _, requested := range []string{"", "mxfp4"} { + if _, err := Classify(inv, requested); err == nil || !strings.Contains(err.Error(), "down_proj_bias") { + t.Fatalf("Classify(requested=%q) error = %v, want required bias error", requested, err) + } + } + + policy, err := newTensorImportTransform(inv) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + if _, err := Plan(inv, Classification{Kind: SourcePrequantized, Quantize: "mxfp4"}, policy); err == nil || !strings.Contains(err.Error(), "down_proj_bias") { + t.Fatalf("Plan() error = %v, want required bias error", err) + } +} + +func TestCreateGPTOSSNativeMXFP4MissingBiasIsRejectedWithoutRequest(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{ + "architectures":["GptOssForCausalLM"], + "quantization":{"quant_method":"mxfp4","mode":"mxfp4","bits":4,"group_size":32} + }`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.layers.0.mlp.experts.down_proj_blocks", "U8", []int32{1, 2, 1, 16}, make([]byte, 32)), + st.NewTensorDataFromBytes("model.layers.0.mlp.experts.down_proj_scales", "U8", []int32{1, 2, 1}, make([]byte, 2)), + }) + + err := Create("gptoss", dir, "", newCaptureStore(), func(string, LayerInfo, []LayerInfo, Classification) error { + return nil + }, func(string) {}) + if err == nil || !strings.Contains(err.Error(), "down_proj_bias") { + t.Fatalf("Create() error = %v, want required bias error", err) + } +} + func TestCreatePipelineReportsGPTOSSMXFP4FileType(t *testing.T) { dir := t.TempDir() writeConfigJSON(t, dir, `{ From 212f27024e114a343b6346a84412ee49d91ca879 Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 07:10:11 +0000 Subject: [PATCH 29/58] create: validate and share GPT-OSS expert transforms Co-authored-by: Codex --- x/create/classify.go | 21 ++++- x/create/gptoss.go | 117 ++++++++++++++++++------ x/create/gptoss_test.go | 157 +++++++++++++++++++++++++++++--- x/create/transform.go | 31 ++++++- x/create/writer.go | 3 +- x/safetensors/extractor.go | 115 ++++++++++++++++++++--- x/safetensors/extractor_test.go | 83 ++++++++++++++++- 7 files changed, 463 insertions(+), 64 deletions(-) diff --git a/x/create/classify.go b/x/create/classify.go index 976663e8450..c221b15b0a0 100644 --- a/x/create/classify.go +++ b/x/create/classify.go @@ -58,12 +58,16 @@ func Classify(inv Inventory, requested string) (Classification, error) { if biasName, missing := missingGPTOSSNativeExpertBias(inv); missing { return Classification{}, fmt.Errorf("incomplete GPT-OSS native MXFP4 checkpoint: required expert bias %q is missing", biasName) } + nativeMXFP4 := gptossNativeMXFP4Evidence(inv) + if gptossHasNativeExpertTensors(inv) && !nativeMXFP4 { + return Classification{}, fmt.Errorf("incomplete or malformed GPT-OSS native MXFP4 expert tensors") + } effective := detectPrequantizedQuantization(inv) - if gptossNativeMXFP4Evidence(inv) { + if nativeMXFP4 { effective = "mxfp4" } if requested != "" { - if effective == requested && gptossNativeMXFP4Evidence(inv) { + if effective == requested && nativeMXFP4 { return Classification{Kind: SourcePrequantized, Quantize: effective}, nil } return Classification{}, fmt.Errorf("cannot requantize an already-quantized source model (requested %q): only bf16/fp16/fp32 sources can be quantized", requested) @@ -90,6 +94,19 @@ func Classify(inv Inventory, requested string) (Classification, error) { return Classification{}, fmt.Errorf("could not classify source model in %s", inv.Dir) } +func gptossHasNativeExpertTensors(inv Inventory) bool { + if inv.Config.Architecture() != "GptOssForCausalLM" { + return false + } + for name := range inv.Tensors { + if strings.Contains(name, ".experts.") && + (strings.HasSuffix(name, "_blocks") || strings.HasSuffix(name, "_scales") || strings.HasSuffix(name, "_bias")) { + return true + } + } + return false +} + func missingGPTOSSNativeExpertBias(inv Inventory) (string, bool) { if inv.Config.Architecture() != "GptOssForCausalLM" { return "", false diff --git a/x/create/gptoss.go b/x/create/gptoss.go index d73257007e2..a000bb54237 100644 --- a/x/create/gptoss.go +++ b/x/create/gptoss.go @@ -376,9 +376,59 @@ func validateGPTOSSPackedMXFP4Inputs(name string, blocks, scales *safetensors.Te if blocks.Shape[3] != 16 { return fmt.Errorf("gpt-oss expert blocks %q trailing shape = %v, want [... 16]", blocks.Name, blocks.Shape) } + for i, dim := range blocks.Shape { + if dim <= 0 { + return fmt.Errorf("gpt-oss expert blocks %q dimension %d must be positive, got %d", blocks.Name, i, dim) + } + } + for i, dim := range scales.Shape { + if dim <= 0 { + return fmt.Errorf("gpt-oss expert scales %q dimension %d must be positive, got %d", scales.Name, i, dim) + } + } + blockBytes, err := checkedGPTOSSProduct("expert block byte size", blocks.Shape[0], blocks.Shape[1], blocks.Shape[2], blocks.Shape[3]) + if err != nil { + return fmt.Errorf("gpt-oss expert blocks %q: %w", blocks.Name, err) + } + scaleBytes, err := checkedGPTOSSProduct("expert scale byte size", scales.Shape[0], scales.Shape[1], scales.Shape[2]) + if err != nil { + return fmt.Errorf("gpt-oss expert scales %q: %w", scales.Name, err) + } + if blocks.Size != int64(blockBytes) { + return fmt.Errorf("gpt-oss expert blocks %q byte size = %d, want %d from shape %v", blocks.Name, blocks.Size, blockBytes, blocks.Shape) + } + if scales.Size != int64(scaleBytes) { + return fmt.Errorf("gpt-oss expert scales %q byte size = %d, want %d from shape %v", scales.Name, scales.Size, scaleBytes, scales.Shape) + } + if blocks.Shape[2] > (1<<31-1)/32 { + return fmt.Errorf("gpt-oss expert tensor %q output shape overflow for %d groups", name, blocks.Shape[2]) + } return nil } +func checkedGPTOSSProduct(label string, factors ...int32) (int, error) { + ints := make([]int, len(factors)) + for i, factor := range factors { + ints[i] = int(factor) + } + return checkedGPTOSSIntProduct(label, ints...) +} + +func checkedGPTOSSIntProduct(label string, factors ...int) (int, error) { + product := uint64(1) + maxInt := uint64(^uint(0) >> 1) + for _, factor := range factors { + if factor <= 0 { + return 0, fmt.Errorf("%s requires positive dimensions, got %v", label, factors) + } + if product > maxInt/uint64(factor) { + return 0, fmt.Errorf("%s overflow for dimensions %v", label, factors) + } + product *= uint64(factor) + } + return int(product), nil +} + func preservePackedExpertProjection(name string, blocks, scales *safetensors.TensorData) ([]*safetensors.TensorData, error) { if err := validateGPTOSSPackedMXFP4Inputs(name, blocks, scales); err != nil { return nil, err @@ -479,8 +529,12 @@ func preserveGPTOSSMXFP4Blocks(name string, source []byte, groups int) ([]byte, if groups <= 0 { return nil, fmt.Errorf("gpt-oss expert blocks %q group count must be positive, got %d", name, groups) } - if len(source)%(groups*16) != 0 { - return nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want multiple of row bytes %d", name, len(source), groups*16) + rowBytes, err := checkedGPTOSSIntProduct("expert block row byte size", groups, 16) + if err != nil { + return nil, fmt.Errorf("gpt-oss expert blocks %q: %w", name, err) + } + if len(source)%rowBytes != 0 { + return nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want multiple of row bytes %d", name, len(source), rowBytes) } out := make([]byte, len(source)) @@ -501,26 +555,8 @@ func decodeGPTOSSMXFP4Scale(scale byte) float32 { } func decodeGPTOSSMXFP4TensorValues(name string, blocks, scales *safetensors.TensorData) ([]float32, []int32, error) { - if blocks == nil || scales == nil { - return nil, nil, fmt.Errorf("gpt-oss expert tensor %q requires blocks and scales", name) - } - if blocks.Dtype != "U8" { - return nil, nil, fmt.Errorf("gpt-oss expert blocks %q dtype = %q, want U8", blocks.Name, blocks.Dtype) - } - if scales.Dtype != "U8" { - return nil, nil, fmt.Errorf("gpt-oss expert scales %q dtype = %q, want U8", scales.Name, scales.Dtype) - } - if len(blocks.Shape) != 4 { - return nil, nil, fmt.Errorf("gpt-oss expert blocks %q shape = %v, want [experts out groups 16]", blocks.Name, blocks.Shape) - } - if len(scales.Shape) != 3 { - return nil, nil, fmt.Errorf("gpt-oss expert scales %q shape = %v, want [experts out groups]", scales.Name, scales.Shape) - } - if blocks.Shape[0] != scales.Shape[0] || blocks.Shape[1] != scales.Shape[1] || blocks.Shape[2] != scales.Shape[2] { - return nil, nil, fmt.Errorf("gpt-oss expert tensor %q shape mismatch: blocks=%v scales=%v", name, blocks.Shape, scales.Shape) - } - if blocks.Shape[3] != 16 { - return nil, nil, fmt.Errorf("gpt-oss expert blocks %q trailing shape = %v, want [... 16]", blocks.Name, blocks.Shape) + if err := validateGPTOSSPackedMXFP4Inputs(name, blocks, scales); err != nil { + return nil, nil, err } blockBytes, err := io.ReadAll(blocks.Reader()) @@ -532,15 +568,21 @@ func decodeGPTOSSMXFP4TensorValues(name string, blocks, scales *safetensors.Tens return nil, nil, fmt.Errorf("read gpt-oss expert scales %q: %w", scales.Name, err) } - groupCount := int(blocks.Shape[0] * blocks.Shape[1] * blocks.Shape[2]) - if len(blockBytes) != groupCount*16 { - return nil, nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want %d", blocks.Name, len(blockBytes), groupCount*16) + groupCount, err := checkedGPTOSSProduct("expert group count", blocks.Shape[0], blocks.Shape[1], blocks.Shape[2]) + if err != nil { + return nil, nil, fmt.Errorf("gpt-oss expert tensor %q: %w", name, err) } - if len(scaleBytes) != groupCount { - return nil, nil, fmt.Errorf("gpt-oss expert scales %q byte length = %d, want %d", scales.Name, len(scaleBytes), groupCount) + if len(blockBytes) != int(blocks.Size) { + return nil, nil, fmt.Errorf("gpt-oss expert blocks %q byte length = %d, want %d", blocks.Name, len(blockBytes), blocks.Size) } - - values := make([]float32, groupCount*32) + if len(scaleBytes) != int(scales.Size) { + return nil, nil, fmt.Errorf("gpt-oss expert scales %q byte length = %d, want %d", scales.Name, len(scaleBytes), scales.Size) + } + valueCount, err := checkedGPTOSSProduct("dequantized expert value count", blocks.Shape[0], blocks.Shape[1], blocks.Shape[2], 32) + if err != nil { + return nil, nil, fmt.Errorf("gpt-oss expert tensor %q: %w", name, err) + } + values := make([]float32, valueCount) for i := range groupCount { src := blockBytes[i*16 : (i+1)*16] @@ -580,7 +622,7 @@ func dequantizeGPTOSSMXFP4Tensor(name string, blocks, scales *safetensors.Tensor func splitGateUpBiasTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { if td == nil { - return nil, nil + return nil, fmt.Errorf("gpt-oss gate/up expert bias is required") } if td.Dtype != "BF16" { return nil, fmt.Errorf("gpt-oss expert tensor %q dtype = %q, want BF16", td.Name, td.Dtype) @@ -588,6 +630,18 @@ func splitGateUpBiasTensor(td *safetensors.TensorData) ([]*safetensors.TensorDat if len(td.Shape) != 2 { return nil, fmt.Errorf("gpt-oss expert tensor %q shape = %v, want [experts out]", td.Name, td.Shape) } + for i, dim := range td.Shape { + if dim <= 0 { + return nil, fmt.Errorf("gpt-oss expert tensor %q dimension %d must be positive, got %d", td.Name, i, dim) + } + } + biasBytes, err := checkedGPTOSSProduct("expert bias byte size", td.Shape[0], td.Shape[1], 2) + if err != nil { + return nil, fmt.Errorf("gpt-oss expert tensor %q: %w", td.Name, err) + } + if td.Size != int64(biasBytes) { + return nil, fmt.Errorf("gpt-oss expert tensor %q byte size = %d, want %d from shape %v", td.Name, td.Size, biasBytes, td.Shape) + } experts, outDim := int(td.Shape[0]), int(td.Shape[1]) if outDim%2 != 0 { return nil, fmt.Errorf("gpt-oss expert tensor %q output dim = %d, want even gate/up rows", td.Name, outDim) @@ -602,6 +656,9 @@ func splitGateUpBiasTensor(td *safetensors.TensorData) ([]*safetensors.TensorDat if err != nil { return nil, fmt.Errorf("decode gpt-oss expert tensor %q: %w", td.Name, err) } + if len(values) != experts*outDim { + return nil, fmt.Errorf("gpt-oss expert tensor %q decoded values = %d, want %d", td.Name, len(values), experts*outDim) + } gateVals := make([]float32, experts*mid) upVals := make([]float32, experts*mid) diff --git a/x/create/gptoss_test.go b/x/create/gptoss_test.go index e362c8c7933..d53674f627e 100644 --- a/x/create/gptoss_test.go +++ b/x/create/gptoss_test.go @@ -1,6 +1,7 @@ package create import ( + "bytes" "io" "maps" "path/filepath" @@ -11,6 +12,23 @@ import ( st "github.com/ollama/ollama/x/safetensors" ) +type countingReaderAt struct { + data []byte + bytesRead int +} + +func (r *countingReaderAt) ReadAt(p []byte, off int64) (int, error) { + if off >= int64(len(r.data)) { + return 0, io.EOF + } + n := copy(p, r.data[off:]) + r.bytesRead += n + if n < len(p) { + return n, io.EOF + } + return n, nil +} + func TestNewTensorImportTransform_GptOSSRegistered(t *testing.T) { inv := gptossInventory(nil) transform, err := newTensorImportTransform(inv) @@ -119,18 +137,18 @@ func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { } for _, tt := range []struct { - name string - mutate func(*Inventory) - checkNoRequest bool + name string + mutate func(*Inventory) + rejectNoRequest bool }{ - {"missing scales", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_scales") }, false}, - {"missing bias", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_bias") }, false}, + {"missing scales", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_scales") }, true}, + {"missing bias", func(inv *Inventory) { delete(inv.Tensors, "model.layers.0.mlp.experts.down_proj_bias") }, true}, {"mismatched scales", func(inv *Inventory) { inv.Tensors["model.layers.0.mlp.experts.down_proj_scales"] = SourceTensor{Name: "model.layers.0.mlp.experts.down_proj_scales", Dtype: "U8", Shape: []int32{2, 15, 1}} - }, false}, + }, true}, {"wrong blocks dtype", func(inv *Inventory) { inv.Tensors["model.layers.0.mlp.experts.down_proj_blocks"] = SourceTensor{Name: "model.layers.0.mlp.experts.down_proj_blocks", Dtype: "BF16", Shape: []int32{2, 16, 1, 16}} - }, false}, + }, true}, {"architecture only", func(inv *Inventory) { clear(inv.Tensors) inv.Tensors["model.layers.0.weight"] = SourceTensor{Name: "model.layers.0.weight", Dtype: "U32", Shape: []int32{16, 4}} @@ -144,13 +162,9 @@ func TestClassifyGPTOSSNativeMXFP4TrustBoundary(t *testing.T) { if _, err := Classify(inv, "mxfp4"); err == nil { t.Fatal("unproven native preservation accepted") } - if tt.checkNoRequest { - class, err := Classify(inv, "") - if err != nil { - t.Fatalf("no-request classification error = %v", err) - } - if class.Kind != SourcePrequantized || class.Quantize != "" { - t.Fatalf("no-request classification = %+v, want unlabelled prequantized source", class) + if tt.rejectNoRequest { + if _, err := Classify(inv, ""); err == nil { + t.Fatal("malformed no-request native source accepted") } } }) @@ -408,6 +422,121 @@ func TestGPTOSSNativeGateUpSplitDequantizesLikeOriginal(t *testing.T) { } } +func TestGPTOSSPackedMXFP4ValidationRejectsMalformedResources(t *testing.T) { + validBlocks := func(shape []int32, size int64) *st.TensorData { + return st.NewTensorDataFromReaderAt("gate_up_blocks", "U8", shape, bytes.NewReader(nil), size) + } + validScales := func(shape []int32, size int64) *st.TensorData { + return st.NewTensorDataFromReaderAt("gate_up_scales", "U8", shape, bytes.NewReader(nil), size) + } + for _, tt := range []struct { + name string + blocks *st.TensorData + scales *st.TensorData + wantContain string + }{ + {name: "zero dimension", blocks: validBlocks([]int32{1, 0, 1, 16}, 0), scales: validScales([]int32{1, 0, 1}, 0), wantContain: "positive"}, + {name: "negative dimension", blocks: validBlocks([]int32{1, -2, 1, 16}, 0), scales: validScales([]int32{1, -2, 1}, 0), wantContain: "positive"}, + {name: "block byte mismatch", blocks: validBlocks([]int32{1, 2, 1, 16}, 31), scales: validScales([]int32{1, 2, 1}, 2), wantContain: "byte size"}, + {name: "scale byte mismatch", blocks: validBlocks([]int32{1, 2, 1, 16}, 32), scales: validScales([]int32{1, 2, 1}, 1), wantContain: "byte size"}, + { + name: "shape multiplication overflow", + blocks: validBlocks([]int32{2147483647, 2147483647, 2147483647, 16}, 0), + scales: validScales([]int32{2147483647, 2147483647, 2147483647}, 0), + wantContain: "overflow", + }, + } { + t.Run(tt.name, func(t *testing.T) { + err := validateGPTOSSPackedMXFP4Inputs("gate_up", tt.blocks, tt.scales) + if err == nil || !strings.Contains(err.Error(), tt.wantContain) { + t.Fatalf("validateGPTOSSPackedMXFP4Inputs() error = %v, want %q", err, tt.wantContain) + } + }) + } +} + +func TestGPTOSSGateUpTransformsShareOneSourceSplit(t *testing.T) { + blockData := make([]byte, 64) + scaleData := []byte{1, 2, 3, 4} + blocksReader := &countingReaderAt{data: blockData} + scalesReader := &countingReaderAt{data: scaleData} + sources := []*st.TensorData{ + st.NewTensorDataFromReaderAt("gate_up_blocks", "U8", []int32{1, 4, 1, 16}, blocksReader, int64(len(blockData))), + st.NewTensorDataFromReaderAt("gate_up_scales", "U8", []int32{1, 4, 1}, scalesReader, int64(len(scaleData))), + } + cache := newByteTransformCache() + for _, transform := range []Transform{ + TransformGPTOSSGateUpWeight, + TransformGPTOSSGateUpScale, + TransformGPTOSSUpWeight, + TransformGPTOSSUpScale, + } { + if _, err := applyByteTransform(TensorSpec{Name: string(transform), Transform: transform}, sources, cache); err != nil { + t.Fatalf("applyByteTransform(%s) error = %v", transform, err) + } + } + if blocksReader.bytesRead != len(blockData) || scalesReader.bytesRead != len(scaleData) { + t.Fatalf("source bytes read = blocks:%d scales:%d, want one split blocks:%d scales:%d", blocksReader.bytesRead, scalesReader.bytesRead, len(blockData), len(scaleData)) + } +} + +func TestGPTOSSGateUpTransformCacheIsBlobScopedAndCachesErrors(t *testing.T) { + shortBlocks := &countingReaderAt{data: []byte{1}} + scales := &countingReaderAt{data: []byte{1, 2, 3, 4}} + malformed := []*st.TensorData{ + st.NewTensorDataFromReaderAt("gate_up_blocks", "U8", []int32{1, 4, 1, 16}, shortBlocks, 64), + st.NewTensorDataFromReaderAt("gate_up_scales", "U8", []int32{1, 4, 1}, scales, 4), + } + cache := newByteTransformCache() + for _, transform := range []Transform{TransformGPTOSSGateUpWeight, TransformGPTOSSUpWeight} { + if _, err := applyByteTransform(TensorSpec{Name: string(transform), Transform: transform}, malformed, cache); err == nil { + t.Fatalf("applyByteTransform(%s) succeeded for truncated source", transform) + } + } + if shortBlocks.bytesRead != len(shortBlocks.data) || scales.bytesRead != len(scales.data) { + t.Fatalf("cached error reread source: blocks=%d scales=%d", shortBlocks.bytesRead, scales.bytesRead) + } + + validBlocks := &countingReaderAt{data: make([]byte, 64)} + validScales := &countingReaderAt{data: []byte{1, 2, 3, 4}} + valid := []*st.TensorData{ + st.NewTensorDataFromReaderAt("gate_up_blocks", "U8", []int32{1, 4, 1, 16}, validBlocks, 64), + st.NewTensorDataFromReaderAt("gate_up_scales", "U8", []int32{1, 4, 1}, validScales, 4), + } + if _, err := applyByteTransform( + TensorSpec{Name: "gate", Transform: TransformGPTOSSGateUpWeight}, + valid, + newByteTransformCache(), + ); err != nil { + t.Fatalf("fresh blob cache reused stale error: %v", err) + } + if validBlocks.bytesRead != len(validBlocks.data) || validScales.bytesRead != len(validScales.data) { + t.Fatalf("fresh blob cache did not read current sources: blocks=%d scales=%d", validBlocks.bytesRead, validScales.bytesRead) + } +} + +func TestGPTOSSGateUpBiasValidationRejectsMalformedResources(t *testing.T) { + for _, tt := range []struct { + name string + shape []int32 + size int64 + wantContain string + }{ + {name: "zero dimension", shape: []int32{1, 0}, size: 0, wantContain: "positive"}, + {name: "negative dimension", shape: []int32{1, -2}, size: 0, wantContain: "positive"}, + {name: "byte mismatch", shape: []int32{1, 4}, size: 7, wantContain: "byte size"}, + {name: "huge declared shape", shape: []int32{2147483647, 2147483646}, size: 0, wantContain: "byte size"}, + } { + t.Run(tt.name, func(t *testing.T) { + td := st.NewTensorDataFromReaderAt("gate_up_bias", "BF16", tt.shape, bytes.NewReader(nil), tt.size) + _, err := splitGateUpBiasTensor(td) + if err == nil || !strings.Contains(err.Error(), tt.wantContain) { + t.Fatalf("splitGateUpBiasTensor() error = %v, want %q", err, tt.wantContain) + } + }) + } +} + func gptossInventory(tensors map[string]SourceTensor) Inventory { if tensors == nil { tensors = map[string]SourceTensor{} diff --git a/x/create/transform.go b/x/create/transform.go index 329e441456b..e35eb05c2b3 100644 --- a/x/create/transform.go +++ b/x/create/transform.go @@ -8,10 +8,27 @@ import ( "github.com/ollama/ollama/x/safetensors" ) +type byteTransformResult struct { + tensors []*safetensors.TensorData + err error +} + +// byteTransformCache is scoped to one output blob. GPT-OSS gate/up companion +// TensorSpecs share the same source pair, so retaining their one split until +// that blob is stored avoids four full reads and split allocations without +// leaking transformed tensors across blobs, requests, or concurrent creates. +type byteTransformCache struct { + gptossGateUp map[string]byteTransformResult +} + +func newByteTransformCache() *byteTransformCache { + return &byteTransformCache{gptossGateUp: make(map[string]byteTransformResult)} +} + // applyByteTransform produces a TensorSpec's output tensor from its resolved // source tensors using only byte-level (non-MLX) operations. The MLX transform // (decode_fp8) and quantization are handled separately by the MLX writer path. -func applyByteTransform(ts TensorSpec, sources []*safetensors.TensorData) (*safetensors.TensorData, error) { +func applyByteTransform(ts TensorSpec, sources []*safetensors.TensorData, cache *byteTransformCache) (*safetensors.TensorData, error) { switch ts.Transform { case TransformNone: if len(sources) != 1 { @@ -66,10 +83,16 @@ func applyByteTransform(ts TensorSpec, sources []*safetensors.TensorData) (*safe if len(sources) != 2 { return nil, fmt.Errorf("transform %s expects block+scale sources, got %d", ts.Transform, len(sources)) } - out, err := preserveAndSplitGateUpTensor(ts.Name, sources[0], sources[1]) - if err != nil { - return nil, err + key := fmt.Sprintf("%s\x00%d\x00%s\x00%d", sources[0].Name, sources[0].Size, sources[1].Name, sources[1].Size) + result, ok := cache.gptossGateUp[key] + if !ok { + result.tensors, result.err = preserveAndSplitGateUpTensor(ts.Name, sources[0], sources[1]) + cache.gptossGateUp[key] = result + } + if result.err != nil { + return nil, result.err } + out := result.tensors switch ts.Transform { case TransformGPTOSSGateUpWeight: return out[0].WithName(ts.Name), nil diff --git a/x/create/writer.go b/x/create/writer.go index 13de6b1d649..c46f40c2497 100644 --- a/x/create/writer.go +++ b/x/create/writer.go @@ -38,6 +38,7 @@ func WriteBlobs(specs []BlobSpec, modelDir string, store BlobStore) ([]LayerInfo // writeBlob resolves each tensor's sources and produces the blob. func writeBlob(spec BlobSpec, src *sourceFiles, store BlobStore) (LayerInfo, error) { needsMLX := blobNeedsMLX(spec) + transformCache := newByteTransformCache() var ( tensors []*safetensors.TensorData items []quantizeItem @@ -54,7 +55,7 @@ func writeBlob(spec BlobSpec, src *sourceFiles, store BlobStore) (LayerInfo, err } items = append(items, quantizeItem{name: ts.Name, quantize: ts.Quantize, reader: reader, decodeFP8: needsFP8Decode(ts.Transform)}) } else { - td, err := applyByteTransform(ts, sources) + td, err := applyByteTransform(ts, sources, transformCache) if err != nil { return LayerInfo{}, fmt.Errorf("blob %s: tensor %s: %w", spec.Name, ts.Name, err) } diff --git a/x/safetensors/extractor.go b/x/safetensors/extractor.go index 28d79226c2f..8ff80e5556d 100644 --- a/x/safetensors/extractor.go +++ b/x/safetensors/extractor.go @@ -12,11 +12,13 @@ import ( // tensorInfo holds tensor metadata from safetensors headers. type tensorInfo struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` - DataOffsets [2]int `json:"data_offsets"` + Dtype string `json:"dtype"` + Shape []int32 `json:"shape"` + DataOffsets [2]int64 `json:"data_offsets"` } +const maxSafetensorsHeaderSize = 100 << 20 + // TensorExtractor extracts individual tensors from a safetensors file. // It provides io.Reader interfaces for each tensor's raw data, enabling // streaming writes to blobs without loading entire tensors into memory. @@ -64,7 +66,7 @@ func (td *TensorData) safetensorsHeader() []byte { td.Name: tensorInfo{ Dtype: td.Dtype, Shape: td.Shape, - DataOffsets: [2]int{0, int(td.Size)}, + DataOffsets: [2]int64{0, td.Size}, }, } headerJSON, _ := json.Marshal(header) @@ -155,14 +157,14 @@ func BuildPackedSafetensorsReader(tensors []*TensorData) io.Reader { func BuildPackedSafetensorsReaderWithMetadata(tensors []*TensorData, metadata map[string]string) io.Reader { // Build the header with sequential data offsets header := make(map[string]any, len(tensors)+1) - var offset int + var offset int64 for _, td := range tensors { header[td.Name] = tensorInfo{ Dtype: td.Dtype, Shape: td.Shape, - DataOffsets: [2]int{offset, offset + int(td.Size)}, + DataOffsets: [2]int64{offset, offset + td.Size}, } - offset += int(td.Size) + offset += td.Size } if len(metadata) > 0 { header["__metadata__"] = metadata @@ -203,9 +205,18 @@ func OpenForExtraction(path string) (*TensorExtractor, error) { f.Close() return nil, fmt.Errorf("failed to read header size: %w", err) } + stat, err := f.Stat() + if err != nil { + f.Close() + return nil, fmt.Errorf("failed to stat file: %w", err) + } + if headerSize > maxSafetensorsHeaderSize || headerSize > uint64(max(0, stat.Size()-8)) { + f.Close() + return nil, fmt.Errorf("invalid safetensors header size %d for file size %d", headerSize, stat.Size()) + } headerBytes := make([]byte, headerSize) - if _, err := f.Read(headerBytes); err != nil { + if _, err := io.ReadFull(f, headerBytes); err != nil { f.Close() return nil, fmt.Errorf("failed to read header: %w", err) } @@ -217,14 +228,96 @@ func OpenForExtraction(path string) (*TensorExtractor, error) { } delete(header, "__metadata__") + dataOffset := int64(8 + headerSize) + payloadSize := stat.Size() - dataOffset + names := make([]string, 0, len(header)) + for name := range header { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + info := header[name] + if err := validateTensorInfo(name, info, payloadSize); err != nil { + f.Close() + return nil, err + } + } + type tensorSpan struct { + name string + start, end int64 + } + spans := make([]tensorSpan, 0, len(names)) + for _, name := range names { + info := header[name] + spans = append(spans, tensorSpan{name: name, start: info.DataOffsets[0], end: info.DataOffsets[1]}) + } + sort.Slice(spans, func(i, j int) bool { + if spans[i].start != spans[j].start { + return spans[i].start < spans[j].start + } + return spans[i].name < spans[j].name + }) + for i := 1; i < len(spans); i++ { + if spans[i].start < spans[i-1].end { + f.Close() + return nil, fmt.Errorf("tensor %q data offsets overlap tensor %q", spans[i].name, spans[i-1].name) + } + } return &TensorExtractor{ file: f, - dataOffset: 8 + int64(headerSize), // 8 bytes for header size + header content + dataOffset: dataOffset, header: header, }, nil } +func validateTensorInfo(name string, info tensorInfo, payloadSize int64) error { + start, end := info.DataOffsets[0], info.DataOffsets[1] + if start < 0 || end < start { + return fmt.Errorf("tensor %q has invalid data offsets [%d %d]", name, start, end) + } + if end > payloadSize { + return fmt.Errorf("tensor %q data offsets [%d %d] exceed payload size %d", name, start, end, payloadSize) + } + width, ok := safetensorsDtypeWidth(info.Dtype) + if !ok { + return fmt.Errorf("tensor %q has unsupported dtype %q", name, info.Dtype) + } + elements := uint64(1) + for i, dim := range info.Shape { + if dim <= 0 { + return fmt.Errorf("tensor %q shape dimension %d must be positive, got %d", name, i, dim) + } + if elements > uint64(^uint64(0))/uint64(dim) { + return fmt.Errorf("tensor %q shape multiplication overflow for %v", name, info.Shape) + } + elements *= uint64(dim) + } + if elements > uint64(^uint64(0))/uint64(width) || elements*uint64(width) > uint64(^uint64(0)>>1) { + return fmt.Errorf("tensor %q byte size overflow for dtype %s shape %v", name, info.Dtype, info.Shape) + } + want := int64(elements * uint64(width)) + if got := end - start; got != want { + return fmt.Errorf("tensor %q byte size = %d, want %d for dtype %s shape %v", name, got, want, info.Dtype, info.Shape) + } + return nil +} + +func safetensorsDtypeWidth(dtype string) (int64, bool) { + switch dtype { + case "BOOL", "U8", "I8", "F8_E4M3", "F8_E4M3FN", "F8_E5M2", "F8_E5M2FNUZ": + return 1, true + case "U16", "I16", "F16", "BF16": + return 2, true + case "U32", "I32", "F32": + return 4, true + case "U64", "I64", "F64": + return 8, true + default: + return 0, false + } +} + // GetTensor returns tensor metadata and a reader for extracting a single tensor. func (te *TensorExtractor) GetTensor(name string) (*TensorData, error) { info, ok := te.header[name] @@ -232,8 +325,8 @@ func (te *TensorExtractor) GetTensor(name string) (*TensorData, error) { return nil, fmt.Errorf("tensor %q not found", name) } - start := te.dataOffset + int64(info.DataOffsets[0]) - size := int64(info.DataOffsets[1] - info.DataOffsets[0]) + start := te.dataOffset + info.DataOffsets[0] + size := info.DataOffsets[1] - info.DataOffsets[0] return &TensorData{ Name: name, diff --git a/x/safetensors/extractor_test.go b/x/safetensors/extractor_test.go index 0b9e1efe6e2..53b721aa934 100644 --- a/x/safetensors/extractor_test.go +++ b/x/safetensors/extractor_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "slices" + "strings" "testing" ) @@ -36,7 +37,7 @@ func createTestSafetensors(t *testing.T, path string, tensors map[string]struct header[name] = tensorInfo{ Dtype: info.dtype, Shape: info.shape, - DataOffsets: [2]int{offset, offset + len(info.data)}, + DataOffsets: [2]int64{int64(offset), int64(offset + len(info.data))}, } allData = append(allData, info.data...) offset += len(info.data) @@ -103,6 +104,84 @@ func TestOpenForExtraction(t *testing.T) { } } +func TestOpenForExtractionRejectsMalformedTensorHeaders(t *testing.T) { + for _, tt := range []struct { + name string + info tensorInfo + data []byte + wantContain string + }{ + {name: "zero dimension", info: tensorInfo{Dtype: "U8", Shape: []int32{1, 0}, DataOffsets: [2]int64{0, 0}}, wantContain: "positive"}, + {name: "negative offset", info: tensorInfo{Dtype: "U8", Shape: []int32{1}, DataOffsets: [2]int64{-1, 0}}, wantContain: "offsets"}, + {name: "reversed offsets", info: tensorInfo{Dtype: "U8", Shape: []int32{1}, DataOffsets: [2]int64{2, 1}}, data: []byte{0, 0}, wantContain: "offsets"}, + {name: "offset past payload", info: tensorInfo{Dtype: "U8", Shape: []int32{2}, DataOffsets: [2]int64{0, 2}}, data: []byte{0}, wantContain: "payload"}, + {name: "dtype size mismatch", info: tensorInfo{Dtype: "U32", Shape: []int32{2}, DataOffsets: [2]int64{0, 4}}, data: make([]byte, 4), wantContain: "byte size"}, + {name: "shape overflow", info: tensorInfo{Dtype: "U8", Shape: []int32{2147483647, 2147483647, 2147483647}, DataOffsets: [2]int64{0, 0}}, wantContain: "overflow"}, + } { + t.Run(tt.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "malformed.safetensors") + header, err := json.Marshal(map[string]tensorInfo{"weight": tt.info}) + if err != nil { + t.Fatal(err) + } + var file bytes.Buffer + if err := binary.Write(&file, binary.LittleEndian, uint64(len(header))); err != nil { + t.Fatal(err) + } + file.Write(header) + file.Write(tt.data) + if err := os.WriteFile(path, file.Bytes(), 0o600); err != nil { + t.Fatal(err) + } + + ext, err := OpenForExtraction(path) + if ext != nil { + ext.Close() + } + if err == nil || !strings.Contains(err.Error(), tt.wantContain) { + t.Fatalf("OpenForExtraction() error = %v, want %q", err, tt.wantContain) + } + }) + } +} + +func TestOpenForExtractionRejectsHeaderBeyondFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "truncated-header.safetensors") + var file bytes.Buffer + if err := binary.Write(&file, binary.LittleEndian, uint64(1024)); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, file.Bytes(), 0o600); err != nil { + t.Fatal(err) + } + if _, err := OpenForExtraction(path); err == nil || !strings.Contains(err.Error(), "header size") { + t.Fatalf("OpenForExtraction() error = %v, want bounded header-size error", err) + } +} + +func TestOpenForExtractionRejectsOverlappingTensorOffsets(t *testing.T) { + path := filepath.Join(t.TempDir(), "overlap.safetensors") + header, err := json.Marshal(map[string]tensorInfo{ + "a": {Dtype: "U8", Shape: []int32{2}, DataOffsets: [2]int64{0, 2}}, + "b": {Dtype: "U8", Shape: []int32{2}, DataOffsets: [2]int64{1, 3}}, + }) + if err != nil { + t.Fatal(err) + } + var file bytes.Buffer + if err := binary.Write(&file, binary.LittleEndian, uint64(len(header))); err != nil { + t.Fatal(err) + } + file.Write(header) + file.Write([]byte{0, 0, 0}) + if err := os.WriteFile(path, file.Bytes(), 0o600); err != nil { + t.Fatal(err) + } + if _, err := OpenForExtraction(path); err == nil || !strings.Contains(err.Error(), "overlap") { + t.Fatalf("OpenForExtraction() error = %v, want overlap error", err) + } +} + func TestGetTensor(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "test.safetensors") @@ -368,7 +447,7 @@ func TestOpenForExtraction_MetadataIgnored(t *testing.T) { "weight": tensorInfo{ Dtype: "F32", Shape: []int32{2}, - DataOffsets: [2]int{0, 8}, + DataOffsets: [2]int64{0, 8}, }, } headerJSON, _ := json.Marshal(header) From b04ed5d333bf158183cbc987dd84a8d66140e00a Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 07:14:00 +0000 Subject: [PATCH 30/58] safetensors: use named integer limits Co-authored-by: Codex --- x/safetensors/extractor.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/x/safetensors/extractor.go b/x/safetensors/extractor.go index 8ff80e5556d..45a4f8013ca 100644 --- a/x/safetensors/extractor.go +++ b/x/safetensors/extractor.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "io" + "math" "os" "sort" ) @@ -288,12 +289,12 @@ func validateTensorInfo(name string, info tensorInfo, payloadSize int64) error { if dim <= 0 { return fmt.Errorf("tensor %q shape dimension %d must be positive, got %d", name, i, dim) } - if elements > uint64(^uint64(0))/uint64(dim) { + if elements > math.MaxUint64/uint64(dim) { return fmt.Errorf("tensor %q shape multiplication overflow for %v", name, info.Shape) } elements *= uint64(dim) } - if elements > uint64(^uint64(0))/uint64(width) || elements*uint64(width) > uint64(^uint64(0)>>1) { + if elements > math.MaxUint64/uint64(width) || elements*uint64(width) > math.MaxInt64 { return fmt.Errorf("tensor %q byte size overflow for dtype %s shape %v", name, info.Dtype, info.Shape) } want := int64(elements * uint64(width)) From e49324c12d90beaf5cbf2e0d8a9cfb602c368832 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 19:18:20 +0000 Subject: [PATCH 31/58] mlx: add Apertus model runtime Register ApertusForCausalLM on the current MLX model contracts and fail closed on unsafe configuration, tensor shapes, dtypes, and quantization companions. Co-authored-by: Codex --- x/mlxrunner/imports.go | 1 + x/models/apertus/apertus.go | 813 +++++++++++++++++++++ x/models/apertus/apertus_test.go | 587 +++++++++++++++ x/models/apertus/forward_reference_test.go | 517 +++++++++++++ x/models/apertus/import_shape_test.go | 145 ++++ 5 files changed, 2063 insertions(+) create mode 100644 x/models/apertus/apertus.go create mode 100644 x/models/apertus/apertus_test.go create mode 100644 x/models/apertus/forward_reference_test.go create mode 100644 x/models/apertus/import_shape_test.go diff --git a/x/mlxrunner/imports.go b/x/mlxrunner/imports.go index 5084ee3b198..5e3fab8f21f 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -1,6 +1,7 @@ package mlxrunner import ( + _ "github.com/ollama/ollama/x/models/apertus" _ "github.com/ollama/ollama/x/models/cohere2_moe" _ "github.com/ollama/ollama/x/models/dflash" _ "github.com/ollama/ollama/x/models/gemma4" diff --git a/x/models/apertus/apertus.go b/x/models/apertus/apertus.go new file mode 100644 index 00000000000..d9402ce87a2 --- /dev/null +++ b/x/models/apertus/apertus.go @@ -0,0 +1,813 @@ +// Package apertus provides the Apertus text model implementation for MLX. +package apertus + +import ( + "encoding/json" + "fmt" + "math" + "strings" + + "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/quant" + "github.com/ollama/ollama/x/tokenizer" +) + +const ( + maxConfigDimension = int32(1 << 24) + maxLayers = int32(1024) + maxHeads = int32(1024) + // MLX dimensions are int32, but a tensor's element product is size_t. This + // ceiling covers the documented 70B/1.5 vocabulary matrix while rejecting + // implausible metadata before it reaches allocation. + maxArrayElements = uint64(1 << 34) +) + +func init() { + base.Register("ApertusForCausalLM", newModel) +} + +// RopeScaling carries the Llama 3 RoPE scaling block used by Apertus. +type RopeScaling struct { + Factor float32 `json:"factor"` + HighFreqFactor float32 `json:"high_freq_factor"` + LowFreqFactor float32 `json:"low_freq_factor"` + OriginalMaxPositionEmbeddings int32 `json:"original_max_position_embeddings"` + RopeType string `json:"rope_type,omitempty"` + Type string `json:"type,omitempty"` +} + +// Config holds Apertus model configuration. +type Config struct { + Architecture string `json:"-"` + ModelType string `json:"model_type"` + DType string `json:"dtype"` + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + VocabSize int32 `json:"vocab_size"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + RMSNormEps float32 `json:"rms_norm_eps"` + RopeTheta float32 `json:"rope_theta"` + RopeScaling RopeScaling `json:"rope_scaling"` + HiddenAct string `json:"hidden_act"` + QKNorm bool `json:"qk_norm"` + PostNorm bool `json:"post_norm"` + AttentionBias bool `json:"attention_bias"` + MLPBias bool `json:"mlp_bias"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + + QuantGroupSize int `json:"-"` + QuantBits int `json:"-"` + QuantMode string `json:"-"` + QuantType string `json:"-"` + TensorQuant map[string]*model.TensorQuantInfo `json:"-"` + + HeadDim int32 `json:"-"` + Scale float32 `json:"-"` + RopeFreqs *mlx.Array `json:"-"` +} + +// Model is an Apertus text model. +type Model struct { + EmbedTokens nn.EmbeddingLayer + Layers []*Layer + Norm *nn.RMSNorm + LMHead nn.LinearLayer + + tok *tokenizer.Tokenizer + *Config +} + +type Layer struct { + AttentionNorm *nn.RMSNorm + Attention *Attention + FFNNorm *nn.RMSNorm + MLP *MLP +} + +type Attention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + QNorm *nn.RMSNorm + KNorm *nn.RMSNorm +} + +type MLP struct { + UpProj nn.LinearLayer + DownProj nn.LinearLayer + Act *XIELU +} + +type XIELU struct { + AlphaP float32 + AlphaN float32 + Beta float32 + Eps float32 +} + +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 + } + ropeFreqs, err := Llama3Freqs( + cfg.HeadDim, + cfg.RopeTheta, + cfg.RopeScaling.Factor, + cfg.RopeScaling.LowFreqFactor, + cfg.RopeScaling.HighFreqFactor, + cfg.RopeScaling.OriginalMaxPositionEmbeddings, + ) + if err != nil { + return nil, fmt.Errorf("build llama3 rope frequencies: %w", err) + } + cfg.RopeFreqs = mlx.FromValues(ropeFreqs, len(ropeFreqs)) + + if qt := root.QuantType(); qt != "" { + cfg.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 data, err := root.Manifest.ReadConfig("generation_config.json"); err == nil { + tokConfig.GenerationConfigJSON = data + } + if data, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { + tokConfig.TokenizerConfigJSON = data + } + if data, err := root.Manifest.ReadConfig("special_tokens_map.json"); err == nil { + tokConfig.SpecialTokensMapJSON = data + } + tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) + if err != nil { + return nil, fmt.Errorf("parse tokenizer: %w", err) + } + + return &Model{Config: &cfg, Layers: make([]*Layer, int(cfg.NumHiddenLayers)), tok: tok}, nil +} + +func parseConfig(configData []byte) (Config, error) { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(configData, &envelope); err != nil { + return Config{}, fmt.Errorf("parse config envelope: %w", err) + } + active := configData + if textRaw, ok := envelope["text_config"]; ok { + active = textRaw + } + + var cfg Config + if err := json.Unmarshal(active, &cfg); err != nil { + return Config{}, fmt.Errorf("parse config: %w", err) + } + var archConfig struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + } + if err := json.Unmarshal(configData, &archConfig); err != nil { + return Config{}, fmt.Errorf("parse architecture: %w", err) + } + if len(archConfig.Architectures) > 0 && archConfig.Architectures[0] != "" { + cfg.Architecture = archConfig.Architectures[0] + } else { + cfg.Architecture = archConfig.ModelType + } + if cfg.Architecture == "" { + return Config{}, fmt.Errorf("missing architecture in config.json") + } + + for _, field := range []struct { + name string + value int32 + max int32 + }{ + {"hidden_size", cfg.HiddenSize, maxConfigDimension}, + {"intermediate_size", cfg.IntermediateSize, maxConfigDimension}, + {"num_hidden_layers", cfg.NumHiddenLayers, maxLayers}, + {"num_attention_heads", cfg.NumAttentionHeads, maxHeads}, + {"vocab_size", cfg.VocabSize, maxConfigDimension}, + {"max_position_embeddings", cfg.MaxPositionEmbeddings, maxConfigDimension}, + } { + if field.value <= 0 || field.value > field.max { + return Config{}, fmt.Errorf("invalid %s: %d (must be in [1,%d])", field.name, field.value, field.max) + } + } + if cfg.NumKeyValueHeads == 0 { + cfg.NumKeyValueHeads = cfg.NumAttentionHeads + } + if cfg.NumKeyValueHeads < 0 || cfg.NumKeyValueHeads > maxHeads { + return Config{}, fmt.Errorf("invalid num_key_value_heads: %d (must be in [1,%d])", cfg.NumKeyValueHeads, maxHeads) + } + 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 + if cfg.HeadDim%2 != 0 { + return Config{}, fmt.Errorf("head_dim must be even: %d", cfg.HeadDim) + } + if cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return Config{}, fmt.Errorf("num_attention_heads (%d) must be divisible by num_key_value_heads (%d)", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) + } + if _, err := checkedProduct("attention projection", uint64(cfg.NumAttentionHeads), uint64(cfg.HeadDim)); err != nil { + return Config{}, err + } + if _, err := checkedProduct("key/value projection", uint64(cfg.NumKeyValueHeads), uint64(cfg.HeadDim)); err != nil { + return Config{}, err + } + if _, err := checkedProduct("embedding", uint64(cfg.VocabSize), uint64(cfg.HiddenSize)); err != nil { + return Config{}, err + } + if _, err := checkedProduct("feed-forward projection", uint64(cfg.IntermediateSize), uint64(cfg.HiddenSize)); err != nil { + return Config{}, err + } + if _, err := checkedProduct("layer tensor descriptors", uint64(cfg.NumHiddenLayers), 15); err != nil { + return Config{}, err + } + if _, err := checkedProduct("RoPE frequency table", uint64(cfg.HeadDim), 1); err != nil { + return Config{}, err + } + + if !positiveFinite(cfg.RMSNormEps) { + return Config{}, fmt.Errorf("invalid rms_norm_eps: %v", cfg.RMSNormEps) + } + if !positiveFinite(cfg.RopeTheta) { + return Config{}, fmt.Errorf("invalid rope_theta: %v", cfg.RopeTheta) + } + if cfg.HiddenAct != "xielu" { + return Config{}, fmt.Errorf("unsupported hidden_act %q", cfg.HiddenAct) + } + if !cfg.QKNorm { + return Config{}, fmt.Errorf("unsupported qk_norm=false") + } + if cfg.PostNorm { + return Config{}, fmt.Errorf("unsupported post_norm=true") + } + if cfg.AttentionBias { + return Config{}, fmt.Errorf("unsupported attention_bias=true") + } + if cfg.MLPBias { + return Config{}, fmt.Errorf("unsupported mlp_bias=true") + } + if cfg.TieWordEmbeddings { + return Config{}, fmt.Errorf("unsupported tie_word_embeddings=true") + } + if ropeType := cfg.ropeType(); ropeType != "llama3" { + return Config{}, fmt.Errorf("unsupported rope scaling type %q", ropeType) + } + if !positiveFinite(cfg.RopeScaling.Factor) || + !positiveFinite(cfg.RopeScaling.LowFreqFactor) || + !positiveFinite(cfg.RopeScaling.HighFreqFactor) { + return Config{}, fmt.Errorf("invalid llama3 rope scaling factors") + } + if cfg.RopeScaling.HighFreqFactor <= cfg.RopeScaling.LowFreqFactor { + return Config{}, fmt.Errorf("high_freq_factor (%v) must exceed low_freq_factor (%v)", cfg.RopeScaling.HighFreqFactor, cfg.RopeScaling.LowFreqFactor) + } + if cfg.RopeScaling.OriginalMaxPositionEmbeddings <= 0 || cfg.RopeScaling.OriginalMaxPositionEmbeddings > maxConfigDimension { + return Config{}, fmt.Errorf("invalid original_max_position_embeddings: %d", cfg.RopeScaling.OriginalMaxPositionEmbeddings) + } + cfg.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim))) + return cfg, nil +} + +func positiveFinite(v float32) bool { + return v > 0 && !math.IsInf(float64(v), 0) && !math.IsNaN(float64(v)) +} + +func checkedProduct(name string, values ...uint64) (uint64, error) { + product := uint64(1) + for _, value := range values { + if value == 0 || product > maxArrayElements/value { + return 0, fmt.Errorf("%s element product exceeds %d", name, maxArrayElements) + } + product *= value + } + return product, nil +} + +func (c Config) ropeType() string { + if c.RopeScaling.RopeType != "" { + return strings.ToLower(c.RopeScaling.RopeType) + } + return strings.ToLower(c.RopeScaling.Type) +} + +// Llama3InvFreqs returns inverse RoPE frequencies matching Transformers. +func Llama3InvFreqs(headDim int32, base, factor, lowFreqFactor, highFreqFactor float32, originalContext int32) ([]float32, error) { + if headDim <= 0 || headDim > maxConfigDimension || headDim%2 != 0 { + return nil, fmt.Errorf("head_dim must be a bounded positive even number: %d", headDim) + } + if !positiveFinite(base) || !positiveFinite(factor) || !positiveFinite(lowFreqFactor) || !positiveFinite(highFreqFactor) || + highFreqFactor <= lowFreqFactor || originalContext <= 0 || originalContext > maxConfigDimension { + return nil, fmt.Errorf("invalid llama3 rope parameters") + } + inv := make([]float32, int(headDim/2)) + lowFreqWavelen := float64(originalContext) / float64(lowFreqFactor) + highFreqWavelen := float64(originalContext) / float64(highFreqFactor) + for i := range inv { + v := 1.0 / math.Pow(float64(base), float64(2*i)/float64(headDim)) + wavelen := 2 * math.Pi / v + switch { + case wavelen > lowFreqWavelen: + v /= float64(factor) + case wavelen >= highFreqWavelen: + smooth := (float64(originalContext)/wavelen - float64(lowFreqFactor)) / (float64(highFreqFactor) - float64(lowFreqFactor)) + v = (1-smooth)*v/float64(factor) + smooth*v + } + if v == 0 || math.IsNaN(v) || math.IsInf(v, 0) { + return nil, fmt.Errorf("invalid llama3 inverse frequency at index %d", i) + } + inv[i] = float32(v) + } + return inv, nil +} + +// Llama3Freqs returns the frequency values expected by MLX RoPEWithFreqs. +func Llama3Freqs(headDim int32, base, factor, lowFreqFactor, highFreqFactor float32, originalContext int32) ([]float32, error) { + inv, err := Llama3InvFreqs(headDim, base, factor, lowFreqFactor, highFreqFactor, originalContext) + if err != nil { + return nil, err + } + freqs := make([]float32, len(inv)) + for i, v := range inv { + freqs[i] = 1 / v + if !positiveFinite(freqs[i]) { + return nil, fmt.Errorf("invalid llama3 frequency at index %d", i) + } + } + return freqs, nil +} + +func qkNormShape(batch, seqLen, heads, headDim int32) []int32 { + return []int32{batch, heads, seqLen, headDim} +} + +func XIELUScalar(x, alphaPParam, alphaNParam, beta, eps float64) float64 { + alphaP := math.Log1p(math.Exp(alphaPParam)) + alphaN := beta + math.Log1p(math.Exp(alphaNParam)) + if x > 0 { + return alphaP*x*x + beta*x + } + return math.Expm1(math.Min(x, eps))*alphaN - x*alphaN + beta*x +} + +func validateTensors(tensors map[string]*mlx.Array, cfg *Config) error { + if cfg == nil { + return fmt.Errorf("missing Apertus config") + } + if err := validateMatrix(tensors, "model.embed_tokens", int(cfg.VocabSize), int(cfg.HiddenSize), cfg, true); err != nil { + return err + } + if err := validateDense(tensors, "model.norm.weight", []int{int(cfg.HiddenSize)}); err != nil { + return err + } + if err := validateMatrix(tensors, "lm_head", int(cfg.VocabSize), int(cfg.HiddenSize), cfg, false); err != nil { + return err + } + qOut := int(cfg.NumAttentionHeads * cfg.HeadDim) + kvOut := int(cfg.NumKeyValueHeads * cfg.HeadDim) + for i := range cfg.NumHiddenLayers { + prefix := fmt.Sprintf("model.layers.%d", i) + for _, spec := range []struct { + name string + shape []int + }{ + {prefix + ".attention_layernorm.weight", []int{int(cfg.HiddenSize)}}, + {prefix + ".feedforward_layernorm.weight", []int{int(cfg.HiddenSize)}}, + {prefix + ".self_attn.q_norm.weight", []int{int(cfg.HeadDim)}}, + {prefix + ".self_attn.k_norm.weight", []int{int(cfg.HeadDim)}}, + {prefix + ".mlp.act_fn.alpha_p", []int{1}}, + {prefix + ".mlp.act_fn.alpha_n", []int{1}}, + {prefix + ".mlp.act_fn.beta", nil}, + {prefix + ".mlp.act_fn.eps", nil}, + } { + if err := validateDense(tensors, spec.name, spec.shape); err != nil { + return fmt.Errorf("layer %d: %w", i, err) + } + } + for _, spec := range []struct { + path string + out, input int + }{ + {prefix + ".self_attn.q_proj", qOut, int(cfg.HiddenSize)}, + {prefix + ".self_attn.k_proj", kvOut, int(cfg.HiddenSize)}, + {prefix + ".self_attn.v_proj", kvOut, int(cfg.HiddenSize)}, + {prefix + ".self_attn.o_proj", int(cfg.HiddenSize), qOut}, + {prefix + ".mlp.up_proj", int(cfg.IntermediateSize), int(cfg.HiddenSize)}, + {prefix + ".mlp.down_proj", int(cfg.HiddenSize), int(cfg.IntermediateSize)}, + } { + if err := validateMatrix(tensors, spec.path, spec.out, spec.input, cfg, false); err != nil { + return fmt.Errorf("layer %d: %w", i, err) + } + } + } + return nil +} + +func requireArray(tensors map[string]*mlx.Array, name string) (*mlx.Array, error) { + t := tensors[name] + if t == nil || !t.Valid() { + return nil, fmt.Errorf("missing tensor %q", name) + } + return t, nil +} + +func validateShape(name string, t *mlx.Array, want []int) error { + if t == nil || !t.Valid() { + return fmt.Errorf("missing tensor %q", name) + } + got := t.Dims() + if len(got) != len(want) { + return fmt.Errorf("tensor %q shape %v, want %v", name, got, want) + } + for i := range want { + if got[i] != want[i] { + return fmt.Errorf("tensor %q shape %v, want %v", name, got, want) + } + } + return nil +} + +func isFloatDType(dtype mlx.DType) bool { + return dtype == mlx.DTypeBFloat16 || dtype == mlx.DTypeFloat16 || dtype == mlx.DTypeFloat32 +} + +func validateDense(tensors map[string]*mlx.Array, name string, want []int) error { + t, err := requireArray(tensors, name) + if err != nil { + return err + } + if err := validateShape(name, t, want); err != nil { + return err + } + if !isFloatDType(t.DType()) { + return fmt.Errorf("tensor %q dtype %s, want floating point", name, t.DType()) + } + for _, suffix := range []string{"_scale", "_qbias", ".global_scale", "_scale_2"} { + if companion := tensors[name+suffix]; companion != nil { + return fmt.Errorf("orphan quantization companion %q", name+suffix) + } + } + return nil +} + +func validateMatrix(tensors map[string]*mlx.Array, path string, out, input int, cfg *Config, embedding bool) error { + name := path + ".weight" + if err := validateExplicitQuantization(cfg, name); err != nil { + return err + } + weight, err := requireArray(tensors, name) + if err != nil { + for _, suffix := range []string{"_scale", "_qbias", ".global_scale", "_scale_2"} { + if tensors[name+suffix] != nil { + return fmt.Errorf("orphan quantization companion %q", name+suffix) + } + } + return err + } + if tensors[path+".bias"] != nil { + return fmt.Errorf("unexpected bias tensor %q", path+".bias") + } + scales := tensors[name+"_scale"] + qbiases := tensors[name+"_qbias"] + global := tensors[name+".global_scale"] + legacyGlobal := tensors[name+"_scale_2"] + if global != nil && legacyGlobal != nil { + return fmt.Errorf("duplicate global scale companions for %q", name) + } + if scales == nil { + if qbiases != nil || global != nil || legacyGlobal != nil { + return fmt.Errorf("incomplete quantization companions for %q", name) + } + if err := validateShape(name, weight, []int{out, input}); err != nil { + return err + } + if !isFloatDType(weight.DType()) { + return fmt.Errorf("tensor %q dtype %s, want floating point", name, weight.DType()) + } + return nil + } + + groupSize, bits, mode := model.ResolveLinearQuantParams( + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant, name, weight, scales, + ) + if groupSize <= 0 || input%groupSize != 0 { + return fmt.Errorf("tensor %q input size %d is incompatible with group size %d", name, input, groupSize) + } + scaleShape := []int{out, input / groupSize} + switch mode { + case "affine": + if bits != 4 && bits != 8 { + return fmt.Errorf("tensor %q has unsupported affine bit width %d", name, bits) + } + if input%(32/bits) != 0 { + return fmt.Errorf("tensor %q input size %d cannot be packed at %d bits", name, input, bits) + } + if err := validateShape(name, weight, []int{out, input / (32 / bits)}); err != nil { + return err + } + if weight.DType() != mlx.DTypeUint32 { + return fmt.Errorf("tensor %q dtype %s, want U32 packed affine", name, weight.DType()) + } + if err := validateShape(name+"_scale", scales, scaleShape); err != nil { + return err + } + if !isFloatDType(scales.DType()) { + return fmt.Errorf("tensor %q dtype %s, want floating point", name+"_scale", scales.DType()) + } + if err := validateShape(name+"_qbias", qbiases, scaleShape); err != nil { + return err + } + if !isFloatDType(qbiases.DType()) { + return fmt.Errorf("tensor %q dtype %s, want floating point", name+"_qbias", qbiases.DType()) + } + if global != nil || legacyGlobal != nil { + return fmt.Errorf("unexpected global scale for affine tensor %q", name) + } + case "nvfp4", "mxfp4": + if bits != 4 || input%8 != 0 { + return fmt.Errorf("tensor %q has invalid %s packing", name, mode) + } + if err := validateShape(name, weight, []int{out, input / 8}); err != nil { + return err + } + if weight.DType() != mlx.DTypeUint32 { + return fmt.Errorf("tensor %q dtype %s, want U32 packed %s", name, weight.DType(), mode) + } + if err := validateShape(name+"_scale", scales, scaleShape); err != nil { + return err + } + if scales.DType() != mlx.DTypeUint8 { + return fmt.Errorf("tensor %q dtype %s, want U8 %s scales", name+"_scale", scales.DType(), mode) + } + if qbiases != nil { + return fmt.Errorf("unexpected qbias companion for %s tensor %q", mode, name) + } + gs := global + if gs == nil { + gs = legacyGlobal + } + if mode == "mxfp4" && gs != nil { + return fmt.Errorf("unexpected global scale for mxfp4 tensor %q", name) + } + if gs != nil { + if err := validateGlobalScale(name, gs, out, embedding); err != nil { + return err + } + } + case "mxfp8": + if bits != 8 { + return fmt.Errorf("tensor %q has invalid mxfp8 bit width %d", name, bits) + } + if err := validateShape(name, weight, []int{out, input}); err != nil { + return err + } + if weight.DType() != mlx.DTypeUint8 { + return fmt.Errorf("tensor %q dtype %s, want U8 mxfp8", name, weight.DType()) + } + if err := validateShape(name+"_scale", scales, scaleShape); err != nil { + return err + } + if scales.DType() != mlx.DTypeUint8 || qbiases != nil || global != nil || legacyGlobal != nil { + return fmt.Errorf("invalid mxfp8 companions for tensor %q", name) + } + default: + return fmt.Errorf("tensor %q has unsupported quantization mode %q", name, mode) + } + return nil +} + +func validateExplicitQuantization(cfg *Config, name string) error { + if raw := strings.TrimSpace(cfg.QuantType); raw != "" { + canonical := quant.Canonical(raw) + if canonical == "" { + return fmt.Errorf("unsupported model quantization type %q", cfg.QuantType) + } + groupSize, bits, mode := model.QuantizationParams(canonical) + if cfg.QuantBits != bits || cfg.QuantMode != mode { + return fmt.Errorf("conflicting model quantization metadata for %q", cfg.QuantType) + } + if mode != "affine" && cfg.QuantGroupSize != groupSize { + return fmt.Errorf("conflicting model quantization group size %d for %q", cfg.QuantGroupSize, cfg.QuantType) + } + } + if cfg.TensorQuant == nil { + return nil + } + tq, ok := cfg.TensorQuant[name] + if !ok { + return nil + } + if tq == nil || strings.TrimSpace(tq.QuantType) == "" { + return fmt.Errorf("missing explicit quantization type for tensor %q", name) + } + canonical := quant.Canonical(tq.QuantType) + if canonical == "" { + return fmt.Errorf("unsupported quantization type %q for tensor %q", tq.QuantType, name) + } + groupSize, _, mode := model.QuantizationParams(canonical) + if tq.GroupSize < 0 || (mode != "affine" && tq.GroupSize > 0 && tq.GroupSize != groupSize) { + return fmt.Errorf("conflicting quantization group size %d for tensor %q type %q", tq.GroupSize, name, tq.QuantType) + } + return nil +} + +func validateGlobalScale(name string, scale *mlx.Array, out int, embedding bool) error { + if scale == nil || !scale.Valid() { + return fmt.Errorf("tensor %q has invalid global scale", name) + } + if scale.DType() != mlx.DTypeFloat32 { + return fmt.Errorf("tensor %q global scale dtype %s, want F32", name, scale.DType()) + } + dims := scale.Dims() + if len(dims) == 0 { + return nil + } + if !embedding && len(dims) == 1 && dims[0] == out { + return nil + } + if embedding { + return fmt.Errorf("tensor %q embedding global scale must be scalar, got shape %v", name, dims) + } + return fmt.Errorf("tensor %q global scale must be scalar or shape [%d], got %v", name, out, dims) +} + +// LoadWeights validates all model-owned tensor contracts before constructing layers. +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { + if err := validateTensors(tensors, m.Config); err != nil { + return err + } + linears := model.NewLinearFactory(tensors, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + m.EmbedTokens = model.MakeEmbeddingLayer(tensors, "model.embed_tokens", m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + m.Norm = nn.NewRMSNorm(tensors["model.norm.weight"], m.RMSNormEps) + m.LMHead = linears.Make("lm_head") + + for i := range m.NumHiddenLayers { + prefix := fmt.Sprintf("model.layers.%d", i) + act, err := newXIELU( + tensors[prefix+".mlp.act_fn.alpha_p"], tensors[prefix+".mlp.act_fn.alpha_n"], + tensors[prefix+".mlp.act_fn.beta"], tensors[prefix+".mlp.act_fn.eps"], + ) + if err != nil { + return fmt.Errorf("layer %d: load xielu activation parameters: %w", i, err) + } + m.Layers[i] = &Layer{ + AttentionNorm: nn.NewRMSNorm(tensors[prefix+".attention_layernorm.weight"], m.RMSNormEps), + FFNNorm: nn.NewRMSNorm(tensors[prefix+".feedforward_layernorm.weight"], m.RMSNormEps), + Attention: &Attention{ + QProj: linears.Make(prefix + ".self_attn.q_proj"), + KProj: linears.Make(prefix + ".self_attn.k_proj"), + VProj: linears.Make(prefix + ".self_attn.v_proj"), + OProj: linears.Make(prefix + ".self_attn.o_proj"), + QNorm: nn.NewRMSNorm(tensors[prefix+".self_attn.q_norm.weight"], m.RMSNormEps), + KNorm: nn.NewRMSNorm(tensors[prefix+".self_attn.k_norm.weight"], m.RMSNormEps), + }, + MLP: &MLP{UpProj: linears.Make(prefix + ".mlp.up_proj"), DownProj: linears.Make(prefix + ".mlp.down_proj"), Act: act}, + } + } + return nil +} + +func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) { + dims := b.InputIDs.Dims() + B, L := int32(dims[0]), int32(dims[1]) + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + h := mlx.Reshape(m.EmbedTokens.Forward(b.InputIDs), B, L, m.HiddenSize) + 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) + } + out := mlx.Reshape(m.Norm.Forward(h, m.RMSNormEps), B, L, m.HiddenSize) + return out, out +} + +func (m *Model) Unembed(x *mlx.Array) *mlx.Array { + dims := x.Dims() + return mlx.Reshape(m.LMHead.Forward(x), int32(dims[0]), int32(dims[1]), m.VocabSize) +} + +func (m *Model) Tokenizer() *tokenizer.Tokenizer { return m.tok } +func (m *Model) MaxContextLength() int { return int(m.MaxPositionEmbeddings) } +func (m *Model) NumLayers() int { return len(m.Layers) } +func (m *Model) NewCaches() []cache.Cache { + caches := make([]cache.Cache, len(m.Layers)) + for i := range caches { + caches[i] = cache.NewKVCache() + } + return caches +} + +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + h := mlx.Add(x, l.Attention.Forward(l.AttentionNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg)) + h = mlx.Reshape(h, B, L, cfg.HiddenSize) + return mlx.Reshape(mlx.Add(h, l.MLP.Forward(l.FFNNorm.Forward(h, cfg.RMSNormEps))), B, L, cfg.HiddenSize) +} + +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + q := mlx.Transpose(mlx.Reshape(a.QProj.Forward(x), B, L, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3) + q = headRMSNorm(a.QNorm, q, B, cfg.NumAttentionHeads, L, cfg.HeadDim, cfg.RMSNormEps) + k := mlx.Transpose(mlx.Reshape(a.KProj.Forward(x), B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3) + k = headRMSNorm(a.KNorm, k, B, cfg.NumKeyValueHeads, L, cfg.HeadDim, cfg.RMSNormEps) + v := mlx.Transpose(mlx.Reshape(a.VProj.Forward(x), B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3) + q = mlx.Reshape(mlx.RoPEWithFreqs(q, int(cfg.HeadDim), false, cfg.RopeTheta, 1, positions, cfg.RopeFreqs), B, cfg.NumAttentionHeads, L, cfg.HeadDim) + k = mlx.Reshape(mlx.RoPEWithFreqs(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1, positions, cfg.RopeFreqs), B, cfg.NumKeyValueHeads, L, cfg.HeadDim) + var kv nn.SDPAOption + if c != nil { + kv = nn.WithKVHistory(c.(cache.Attention).Update(b, k, v)) + } 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.HiddenSize) + return a.OProj.Forward(out) +} + +func headRMSNorm(norm *nn.RMSNorm, x *mlx.Array, batch, heads, seqLen, headDim int32, eps float32) *mlx.Array { + x = mlx.Reshape(x, -1, headDim) + return mlx.Reshape(norm.Forward(x, eps), batch, heads, seqLen, headDim) +} + +func (m *MLP) Forward(x *mlx.Array) *mlx.Array { + return m.DownProj.Forward(m.Act.Forward(m.UpProj.Forward(x))) +} + +func (a *XIELU) Forward(x *mlx.Array) *mlx.Array { + outDType := x.DType() + x = x.AsType(mlx.DTypeFloat32) + zero, one := mlx.FromValue[float32](0), mlx.FromValue[float32](1) + alphaP, alphaN := mlx.FromValue(a.AlphaP), mlx.FromValue(a.AlphaN) + beta, eps := mlx.FromValue(a.Beta), mlx.FromValue(a.Eps) + positive := mlx.Add(mlx.Mul(alphaP, mlx.Mul(x, x)), mlx.Mul(beta, x)) + expm1 := mlx.Sub(mlx.Exp(mlx.Minimum(x, eps)), one) + negative := mlx.Add(mlx.Mul(mlx.Sub(expm1, x), alphaN), mlx.Mul(beta, x)) + return mlx.Where(x.Greater(zero), positive, negative).AsType(outDType) +} + +func newXIELU(alphaPParam, alphaNParam, betaParam, epsParam *mlx.Array) (*XIELU, error) { + alphaP, err := scalarParam(alphaPParam) + if err != nil { + return nil, fmt.Errorf("alpha_p: %w", err) + } + alphaN, err := scalarParam(alphaNParam) + if err != nil { + return nil, fmt.Errorf("alpha_n: %w", err) + } + beta, err := scalarParam(betaParam) + if err != nil { + return nil, fmt.Errorf("beta: %w", err) + } + eps, err := scalarParam(epsParam) + if err != nil { + return nil, fmt.Errorf("eps: %w", err) + } + return &XIELU{AlphaP: float32(softplus64(float64(alphaP))), AlphaN: beta + float32(softplus64(float64(alphaN))), Beta: beta, Eps: eps}, nil +} + +func scalarParam(x *mlx.Array) (float32, error) { + if x == nil || !x.Valid() || x.Size() != 1 { + return 0, fmt.Errorf("expected scalar or single-element tensor") + } + x = x.AsType(mlx.DTypeFloat32) + mlx.Eval(x) + values := x.Floats() + if len(values) != 1 || math.IsNaN(float64(values[0])) || math.IsInf(float64(values[0]), 0) { + return 0, fmt.Errorf("expected one finite scalar value") + } + return values[0], nil +} + +func softplus64(x float64) float64 { + if x > 20 { + return x + } + if x < -20 { + return math.Exp(x) + } + return math.Log1p(math.Exp(x)) +} diff --git a/x/models/apertus/apertus_test.go b/x/models/apertus/apertus_test.go new file mode 100644 index 00000000000..7085edd45bc --- /dev/null +++ b/x/models/apertus/apertus_test.go @@ -0,0 +1,587 @@ +package apertus + +import ( + "fmt" + "math" + "os" + "path/filepath" + "strings" + "testing" + + imagemanifest "github.com/ollama/ollama/x/imagegen/manifest" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/mlxrunner/model/base" +) + +func TestRegistration(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + + root := minimalManifestRoot(t) + got, err := base.New(root) + if err != nil { + t.Fatal(err) + } + if _, ok := got.(*Model); !ok { + t.Fatalf("base.New() returned %T, want *apertus.Model", got) + } +} + +func TestParseConfigApertus8B(t *testing.T) { + cfg, err := parseConfig([]byte(`{ + "architectures": ["ApertusForCausalLM"], + "model_type": "apertus", + "dtype": "bfloat16", + "hidden_size": 4096, + "intermediate_size": 21504, + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "max_position_embeddings": 65536, + "rope_theta": 12000000, + "rope_scaling": { + "factor": 8, + "high_freq_factor": 4, + "low_freq_factor": 1, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + "type": "llama3" + }, + "hidden_act": "xielu", + "qk_norm": true, + "post_norm": false, + "attention_bias": false, + "mlp_bias": false, + "tie_word_embeddings": false, + "rms_norm_eps": 1e-5, + "vocab_size": 131072 + }`)) + if err != nil { + t.Fatal(err) + } + + if cfg.Architecture != "ApertusForCausalLM" { + t.Fatalf("Architecture = %q", cfg.Architecture) + } + if cfg.HeadDim != 128 { + t.Fatalf("HeadDim = %d, want 128", cfg.HeadDim) + } + if cfg.Scale != float32(1/math.Sqrt(128)) { + t.Fatalf("Scale = %v, want %v", cfg.Scale, float32(1/math.Sqrt(128))) + } + if cfg.MaxPositionEmbeddings != 65536 { + t.Fatalf("MaxPositionEmbeddings = %d, want 65536", cfg.MaxPositionEmbeddings) + } + if cfg.RopeScaling.OriginalMaxPositionEmbeddings != 8192 { + t.Fatalf("OriginalMaxPositionEmbeddings = %d, want 8192", cfg.RopeScaling.OriginalMaxPositionEmbeddings) + } +} + +func TestParseConfigRejectsUnsupportedVariants(t *testing.T) { + baseConfig := `{ + "architectures": ["ApertusForCausalLM"], + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "max_position_embeddings": 64, + "rope_theta": 12000000, + "rope_scaling": { + "factor": 8, + "high_freq_factor": 4, + "low_freq_factor": 1, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + }, + "hidden_act": "silu", + "qk_norm": true, + "rms_norm_eps": 1e-5, + "vocab_size": 128 + }` + + _, err := parseConfig([]byte(baseConfig)) + if err == nil || !strings.Contains(err.Error(), `unsupported hidden_act "silu"`) { + t.Fatalf("parseConfig error = %v, want unsupported hidden_act", err) + } +} + +func TestLlama3RoPEReferenceValues(t *testing.T) { + got, err := Llama3InvFreqs(8, 12000000, 8, 1, 4, 8192) + if err != nil { + t.Fatal(err) + } + want := []float32{ + 1, + 0.016990442, + 0.000036084392, + 0.00000061308978, + } + if len(got) != len(want) { + t.Fatalf("len(got) = %d, want %d", len(got), len(want)) + } + for i := range want { + if diff := math.Abs(float64(got[i] - want[i])); diff > 1e-9 { + t.Fatalf("inv_freq[%d] = %.12g, want %.12g", i, got[i], want[i]) + } + } + + freqs, err := Llama3Freqs(8, 12000000, 8, 1, 4, 8192) + if err != nil { + t.Fatal(err) + } + for i := range want { + wantFreq := 1 / want[i] + if diff := math.Abs(float64(freqs[i] - wantFreq)); diff/float64(wantFreq) > 1e-6 { + t.Fatalf("freq[%d] = %.12g, want reciprocal %.12g", i, freqs[i], wantFreq) + } + } +} + +func TestQKNormShape(t *testing.T) { + got := qkNormShape(2, 3, 4, 5) + want := []int32{2, 4, 3, 5} + for i := range want { + if got[i] != want[i] { + t.Fatalf("qkNormShape() = %v, want %v", got, want) + } + } +} + +func TestXIELUScalarParity(t *testing.T) { + tests := []struct { + x float64 + want float64 + }{ + {x: 2.0, want: 3.772588722239781}, + {x: -2.0, want: 0.35462209218399154}, + } + for _, tt := range tests { + got := XIELUScalar(tt.x, 0.0, 0.0, 0.5, -1e-6) + if diff := math.Abs(got - tt.want); diff > 1e-12 { + t.Fatalf("XIELUScalar(%v) = %.15g, want %.15g", tt.x, got, tt.want) + } + } +} + +func TestParseConfigBoundsAndNonFiniteValues(t *testing.T) { + valid := `{ + "architectures":["ApertusForCausalLM"],"hidden_size":16,"intermediate_size":32, + "num_hidden_layers":1,"num_attention_heads":4,"num_key_value_heads":2, + "vocab_size":128,"max_position_embeddings":64,"rms_norm_eps":1e-5,"rope_theta":12000000, + "rope_scaling":{"factor":8,"high_freq_factor":4,"low_freq_factor":1,"original_max_position_embeddings":8192,"rope_type":"llama3"}, + "hidden_act":"xielu","qk_norm":true}` + for _, tc := range []struct { + name, old, replacement, want string + }{ + {"zero hidden", `"hidden_size":16`, `"hidden_size":0`, "hidden_size"}, + {"layers first rejected", `"num_hidden_layers":1`, `"num_hidden_layers":1025`, "num_hidden_layers"}, + {"heads first rejected", `"num_attention_heads":4`, `"num_attention_heads":1025`, "num_attention_heads"}, + {"negative kv heads", `"num_key_value_heads":2`, `"num_key_value_heads":-1`, "num_key_value_heads"}, + {"context first rejected", `"max_position_embeddings":64`, `"max_position_embeddings":16777217`, "max_position_embeddings"}, + {"vocab first rejected", `"vocab_size":128`, `"vocab_size":16777217`, "vocab_size"}, + {"zero norm", `"rms_norm_eps":1e-5`, `"rms_norm_eps":0`, "rms_norm_eps"}, + {"overflowing norm", `"rms_norm_eps":1e-5`, `"rms_norm_eps":1e100`, "rms_norm_eps"}, + {"zero theta", `"rope_theta":12000000`, `"rope_theta":0`, "rope_theta"}, + {"equal rope factors", `"high_freq_factor":4`, `"high_freq_factor":1`, "high_freq_factor"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := parseConfig([]byte(strings.Replace(valid, tc.old, tc.replacement, 1))) + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("parseConfig error = %v, want %q", err, tc.want) + } + }) + } + + maxima := strings.NewReplacer( + `"hidden_size":16`, `"hidden_size":2048`, + `"intermediate_size":32`, `"intermediate_size":2`, + `"num_hidden_layers":1`, `"num_hidden_layers":1024`, + `"num_attention_heads":4`, `"num_attention_heads":1024`, + `"num_key_value_heads":2`, `"num_key_value_heads":1024`, + `"vocab_size":128`, `"vocab_size":2`, + `"max_position_embeddings":64`, `"max_position_embeddings":16777216`, + ).Replace(valid) + if _, err := parseConfig([]byte(maxima)); err != nil { + t.Fatalf("documented maxima rejected: %v", err) + } + if got, err := checkedProduct("boundary", maxArrayElements, 1); err != nil || got != maxArrayElements { + t.Fatalf("greatest array product = %d, %v", got, err) + } + if _, err := checkedProduct("boundary", maxArrayElements, 2); err == nil { + t.Fatal("first overflowing array product accepted") + } + future70B := strings.NewReplacer( + `"hidden_size":16`, `"hidden_size":8192`, + `"intermediate_size":32`, `"intermediate_size":43008`, + `"num_hidden_layers":1`, `"num_hidden_layers":80`, + `"num_attention_heads":4`, `"num_attention_heads":64`, + `"num_key_value_heads":2`, `"num_key_value_heads":8`, + `"vocab_size":128`, `"vocab_size":266752`, + `"max_position_embeddings":64`, `"max_position_embeddings":262144`, + ).Replace(valid) + if _, err := parseConfig([]byte(future70B)); err != nil { + t.Fatalf("documented 70B plus recorded 1.5 dimensions rejected: %v", err) + } +} + +func TestTensorValidationDenseAndPacked(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + cfg := tinyConfig() + tensors := tinyDenseTensors(cfg) + if err := validateTensors(tensors, cfg); err != nil { + t.Fatalf("valid dense tensors rejected: %v", err) + } + + for _, tc := range []struct { + name string + edit func(map[string]*mlx.Array) + }{ + {"missing", func(ts map[string]*mlx.Array) { delete(ts, "model.layers.0.self_attn.k_proj.weight") }}, + {"wrong shape", func(ts map[string]*mlx.Array) { + ts["model.layers.0.mlp.up_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, 31, 16) + }}, + {"wrong dtype", func(ts map[string]*mlx.Array) { ts["model.norm.weight"] = mlx.Zeros(mlx.DTypeUint32, 16) }}, + {"orphan scale", func(ts map[string]*mlx.Array) { + ts["model.layers.0.self_attn.q_proj.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, 16, 1) + }}, + } { + t.Run(tc.name, func(t *testing.T) { + ts := tinyDenseTensors(cfg) + tc.edit(ts) + if err := validateTensors(ts, cfg); err == nil { + t.Fatal("malformed dense tensors accepted") + } + }) + } + + name := "model.layers.0.self_attn.q_proj.weight" + tensors[name] = mlx.Zeros(mlx.DTypeUint32, 16, 2) + tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, 16, 1) + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {QuantType: "nvfp4", GroupSize: 16}} + if err := validateTensors(tensors, cfg); err != nil { + t.Fatalf("valid NVFP4 tensors rejected: %v", err) + } + for _, tc := range []struct { + name string + edit func(map[string]*mlx.Array) + }{ + {"packed shape", func(ts map[string]*mlx.Array) { ts[name] = mlx.Zeros(mlx.DTypeUint32, 16, 3) }}, + {"packed dtype", func(ts map[string]*mlx.Array) { ts[name] = mlx.Zeros(mlx.DTypeUint8, 16, 2) }}, + {"scale shape", func(ts map[string]*mlx.Array) { ts[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, 16, 2) }}, + {"scale dtype", func(ts map[string]*mlx.Array) { ts[name+"_scale"] = mlx.Zeros(mlx.DTypeBFloat16, 16, 1) }}, + {"qbias", func(ts map[string]*mlx.Array) { ts[name+"_qbias"] = mlx.Zeros(mlx.DTypeUint8, 16, 1) }}, + {"duplicate global", func(ts map[string]*mlx.Array) { + ts[name+".global_scale"] = mlx.Zeros(mlx.DTypeFloat32, 1) + ts[name+"_scale_2"] = mlx.Zeros(mlx.DTypeFloat32, 1) + }}, + } { + t.Run(tc.name, func(t *testing.T) { + ts := tinyDenseTensors(cfg) + ts[name] = mlx.Zeros(mlx.DTypeUint32, 16, 2) + ts[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, 16, 1) + tc.edit(ts) + if err := validateTensors(ts, cfg); err == nil { + t.Fatal("malformed packed tensors accepted") + } + }) + } + if err := validateTensors(tensors, cfg); err != nil { + t.Fatalf("valid packed tensors rejected after failures: %v", err) + } +} + +func TestMatrixQuantizationIdentitiesFailClosed(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + const name = "projection.weight" + for _, quantType := range []string{"int4", "int8", "nvfp4", "mxfp4", "mxfp8", "FP4", "Q8"} { + t.Run("supported model "+quantType, func(t *testing.T) { + tensors, cfg := quantizedMatrix(name, quantType, 4, 64) + if err := validateMatrix(tensors, "projection", 4, 64, cfg, false); err != nil { + t.Fatalf("supported quantization %q rejected: %v", quantType, err) + } + }) + t.Run("supported tensor "+quantType, func(t *testing.T) { + tensors, cfg := quantizedMatrix(name, quantType, 4, 64) + cfg.QuantType, cfg.QuantMode, cfg.QuantBits, cfg.QuantGroupSize = "", "", 0, 0 + groupSize, _, _ := model.QuantizationParams(quantType) + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {QuantType: quantType, GroupSize: groupSize}} + if err := validateMatrix(tensors, "projection", 4, 64, cfg, false); err != nil { + t.Fatalf("supported per-tensor quantization %q rejected: %v", quantType, err) + } + }) + } + + t.Run("supported dense", func(t *testing.T) { + if err := validateMatrix(map[string]*mlx.Array{name: mlx.Zeros(mlx.DTypeBFloat16, 4, 64)}, "projection", 4, 64, &Config{}, false); err != nil { + t.Fatalf("dense matrix rejected: %v", err) + } + }) + + for _, tc := range []struct { + name string + edit func(*Config) + want string + }{ + { + name: "unknown model identity", + edit: func(cfg *Config) { cfg.QuantType = "future8" }, + want: "unsupported model quantization type", + }, + { + name: "conflicting model identity", + edit: func(cfg *Config) { + cfg.QuantType = "nvfp4" + cfg.QuantBits = 8 + cfg.QuantMode = "affine" + }, + want: "conflicting model quantization metadata", + }, + { + name: "unknown tensor identity", + edit: func(cfg *Config) { + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {QuantType: "future8", GroupSize: 32}} + }, + want: "unsupported quantization type", + }, + { + name: "empty tensor identity", + edit: func(cfg *Config) { + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {GroupSize: 32}} + }, + want: "missing explicit quantization type", + }, + { + name: "nil tensor identity", + edit: func(cfg *Config) { cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: nil} }, + want: "missing explicit quantization type", + }, + { + name: "conflicting tensor group size", + edit: func(cfg *Config) { + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {QuantType: "nvfp4", GroupSize: 32}} + }, + want: "conflicting quantization group size", + }, + { + name: "negative tensor group size", + edit: func(cfg *Config) { + cfg.TensorQuant = map[string]*model.TensorQuantInfo{name: {QuantType: "int8", GroupSize: -1}} + }, + want: "conflicting quantization group size", + }, + } { + t.Run(tc.name, func(t *testing.T) { + tensors, cfg := quantizedMatrix(name, "int8", 4, 64) + tc.edit(cfg) + err := validateMatrix(tensors, "projection", 4, 64, cfg, false) + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("validateMatrix error = %v, want %q", err, tc.want) + } + }) + } + + t.Run("unknown identity rejected before dense fallback", func(t *testing.T) { + err := validateMatrix( + map[string]*mlx.Array{name: mlx.Zeros(mlx.DTypeBFloat16, 4, 64)}, + "projection", 4, 64, &Config{QuantType: "future8", QuantBits: 8, QuantMode: "affine"}, false, + ) + if err == nil || !strings.Contains(err.Error(), "unsupported model quantization type") { + t.Fatalf("validateMatrix error = %v, want unknown identity rejection", err) + } + }) +} + +func TestNVFP4GlobalScaleShapes(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + const ( + linearName = "projection.weight" + embeddingName = "embedding.weight" + ) + for _, tc := range []struct { + name string + embedding bool + shape []int + invalid bool + wantOK bool + }{ + {name: "linear scalar", shape: nil, wantOK: true}, + {name: "linear per row", shape: []int{4}, wantOK: true}, + {name: "linear invalid handle", invalid: true}, + {name: "linear wrong rank same size", shape: []int{4, 1}}, + {name: "linear transposed wrong rank same size", shape: []int{1, 4}}, + {name: "linear wrong length", shape: []int{3}}, + {name: "embedding scalar", embedding: true, shape: nil, wantOK: true}, + {name: "embedding per row", embedding: true, shape: []int{4}}, + {name: "embedding wrong rank same size", embedding: true, shape: []int{4, 1}}, + {name: "embedding single element vector is not scalar", embedding: true, shape: []int{1}}, + } { + t.Run(tc.name, func(t *testing.T) { + name := linearName + path := "projection" + if tc.embedding { + name = embeddingName + path = "embedding" + } + tensors, cfg := quantizedMatrix(name, "nvfp4", 4, 64) + if tc.invalid { + tensors[name+".global_scale"] = &mlx.Array{} + } else { + tensors[name+".global_scale"] = mlx.Zeros(mlx.DTypeFloat32, tc.shape...) + } + err := validateMatrix(tensors, path, 4, 64, cfg, tc.embedding) + if tc.wantOK && err != nil { + t.Fatalf("valid global scale rejected: %v", err) + } + if !tc.wantOK && err == nil { + t.Fatal("invalid global scale accepted") + } + }) + } +} + +func quantizedMatrix(name, quantType string, out, input int) (map[string]*mlx.Array, *Config) { + groupSize, bits, mode := model.QuantizationParams(quantType) + tensors := map[string]*mlx.Array{} + switch mode { + case "affine": + tensors[name] = mlx.Zeros(mlx.DTypeUint32, out, input/(32/bits)) + tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeBFloat16, out, input/groupSize) + tensors[name+"_qbias"] = mlx.Zeros(mlx.DTypeBFloat16, out, input/groupSize) + case "nvfp4", "mxfp4": + tensors[name] = mlx.Zeros(mlx.DTypeUint32, out, input/8) + tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, out, input/groupSize) + case "mxfp8": + tensors[name] = mlx.Zeros(mlx.DTypeUint8, out, input) + tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, out, input/groupSize) + } + return tensors, &Config{QuantType: quantType, QuantGroupSize: groupSize, QuantBits: bits, QuantMode: mode} +} + +func TestCacheCount(t *testing.T) { + m := &Model{Config: &Config{}, Layers: make([]*Layer, 7)} + if got := len(m.NewCaches()); got != 7 { + t.Fatalf("cache count = %d, want 7", got) + } +} + +func tinyConfig() *Config { + return &Config{ + HiddenSize: 16, IntermediateSize: 32, NumHiddenLayers: 1, NumAttentionHeads: 4, + NumKeyValueHeads: 2, VocabSize: 128, HeadDim: 4, QuantMode: "", TensorQuant: map[string]*model.TensorQuantInfo{}, + } +} + +func tinyDenseTensors(cfg *Config) map[string]*mlx.Array { + hidden, intermediate, vocab := int(cfg.HiddenSize), int(cfg.IntermediateSize), int(cfg.VocabSize) + qOut, kvOut, headDim := int(cfg.NumAttentionHeads*cfg.HeadDim), int(cfg.NumKeyValueHeads*cfg.HeadDim), int(cfg.HeadDim) + tensors := map[string]*mlx.Array{ + "model.embed_tokens.weight": mlx.Zeros(mlx.DTypeBFloat16, vocab, hidden), + "model.norm.weight": mlx.Zeros(mlx.DTypeBFloat16, hidden), + "lm_head.weight": mlx.Zeros(mlx.DTypeBFloat16, vocab, hidden), + } + for i := range cfg.NumHiddenLayers { + prefix := fmt.Sprintf("model.layers.%d", i) + tensors[prefix+".attention_layernorm.weight"] = mlx.Zeros(mlx.DTypeBFloat16, hidden) + tensors[prefix+".feedforward_layernorm.weight"] = mlx.Zeros(mlx.DTypeBFloat16, hidden) + tensors[prefix+".self_attn.q_norm.weight"] = mlx.Zeros(mlx.DTypeBFloat16, headDim) + tensors[prefix+".self_attn.k_norm.weight"] = mlx.Zeros(mlx.DTypeBFloat16, headDim) + tensors[prefix+".self_attn.q_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, qOut, hidden) + tensors[prefix+".self_attn.k_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, kvOut, hidden) + tensors[prefix+".self_attn.v_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, kvOut, hidden) + tensors[prefix+".self_attn.o_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, hidden, hidden) + tensors[prefix+".mlp.up_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, intermediate, hidden) + tensors[prefix+".mlp.down_proj.weight"] = mlx.Zeros(mlx.DTypeBFloat16, hidden, intermediate) + tensors[prefix+".mlp.act_fn.alpha_p"] = mlx.Zeros(mlx.DTypeBFloat16, 1) + tensors[prefix+".mlp.act_fn.alpha_n"] = mlx.Zeros(mlx.DTypeBFloat16, 1) + tensors[prefix+".mlp.act_fn.beta"] = mlx.Zeros(mlx.DTypeBFloat16) + tensors[prefix+".mlp.act_fn.eps"] = mlx.Zeros(mlx.DTypeBFloat16) + } + return tensors +} + +func minimalManifestRoot(t *testing.T) *model.Root { + t.Helper() + + dir := t.TempDir() + configDigest := writeManifestBlob(t, dir, "config", []byte(`{ + "architectures": ["ApertusForCausalLM"], + "model_type": "apertus", + "dtype": "bfloat16", + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "max_position_embeddings": 64, + "rope_theta": 12000000, + "rope_scaling": { + "factor": 8, + "high_freq_factor": 4, + "low_freq_factor": 1, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + }, + "hidden_act": "xielu", + "qk_norm": true, + "post_norm": false, + "attention_bias": false, + "mlp_bias": false, + "tie_word_embeddings": false, + "rms_norm_eps": 1e-5, + "vocab_size": 3 + }`)) + tokenizerDigest := writeManifestBlob(t, dir, "tokenizer", []byte(`{ + "model": { + "type": "BPE", + "vocab": {"": 0, "hello": 1, "world": 2}, + "merges": [] + }, + "added_tokens": [ + {"id": 0, "content": "", "special": true} + ] + }`)) + + return &model.Root{ + Manifest: &imagemanifest.ModelManifest{ + BlobDir: dir, + Manifest: &imagemanifest.Manifest{ + SchemaVersion: 2, + MediaType: "application/vnd.ollama.image.model", + Layers: []imagemanifest.ManifestLayer{ + { + MediaType: "application/vnd.ollama.image.json", + Digest: configDigest, + Size: 1, + Name: "config.json", + }, + { + MediaType: "application/vnd.ollama.image.json", + Digest: tokenizerDigest, + Size: 1, + Name: "tokenizer.json", + }, + }, + }, + }, + } +} + +func writeManifestBlob(t *testing.T, dir, name string, data []byte) string { + t.Helper() + + digest := "sha256:" + name + path := filepath.Join(dir, "sha256-"+name) + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + return digest +} diff --git a/x/models/apertus/forward_reference_test.go b/x/models/apertus/forward_reference_test.go new file mode 100644 index 00000000000..0a735742d72 --- /dev/null +++ b/x/models/apertus/forward_reference_test.go @@ -0,0 +1,517 @@ +package apertus + +import ( + "context" + "fmt" + "math" + "os" + "runtime/debug" + "strings" + "testing" + + "github.com/ollama/ollama/x/internal/mlxthread" + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/mlxrunner/model/base" +) + +func TestForwardReference(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + + refPath := firstNonEmpty(os.Getenv("APERTUS_REF_PATH"), os.Getenv("PORTING_APERTUS_REF_PATH")) + if refPath == "" { + t.Skip("set APERTUS_REF_PATH or PORTING_APERTUS_REF_PATH to an HF activation safetensors reference") + } + modelName := firstNonEmpty(os.Getenv("APERTUS_MODEL_NAME"), os.Getenv("PORTING_APERTUS_MODEL_NAME"), "apertus-mlx") + + thread, err := mlxthread.Start("apertus-forward-reference", func() error { + if err := mlx.CheckInit(); err != nil { + return err + } + if mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + configureCompileForTest(t) + return nil + }) + if err != nil { + t.Skipf("MLX not available: %v", err) + } + defer func() { + if err := thread.Stop(context.Background(), func() { + mlx.Sweep() + mlx.ClearCache() + }); err != nil { + t.Fatal(err) + } + }() + + if err := thread.Do(context.Background(), func() (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("panic in forward reference test: %v\n%s", recovered, debug.Stack()) + } + }() + runForwardReference(t, refPath, modelName) + return nil + }); err != nil { + t.Fatal(err) + } +} + +func TestForwardCacheReference(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + + refPath := firstNonEmpty(os.Getenv("APERTUS_CACHE_REF_PATH"), os.Getenv("PORTING_APERTUS_CACHE_REF_PATH")) + if refPath == "" { + t.Skip("set APERTUS_CACHE_REF_PATH or PORTING_APERTUS_CACHE_REF_PATH to an HF cached decode safetensors reference") + } + modelName := firstNonEmpty(os.Getenv("APERTUS_MODEL_NAME"), os.Getenv("PORTING_APERTUS_MODEL_NAME"), "apertus-mlx") + + thread, err := mlxthread.Start("apertus-forward-cache-reference", func() error { + if err := mlx.CheckInit(); err != nil { + return err + } + if mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + configureCompileForTest(t) + return nil + }) + if err != nil { + t.Skipf("MLX not available: %v", err) + } + defer func() { + if err := thread.Stop(context.Background(), func() { + mlx.Sweep() + mlx.ClearCache() + }); err != nil { + t.Fatal(err) + } + }() + + if err := thread.Do(context.Background(), func() (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("panic in cache reference test: %v\n%s", recovered, debug.Stack()) + } + }() + runForwardCacheReference(t, refPath, modelName) + return nil + }); err != nil { + t.Fatal(err) + } +} + +func runForwardReference(t *testing.T, refPath, modelName string) { + t.Helper() + ref := loadReferenceFiltered(t, refPath, map[string]bool{ + "input_ids": true, + "logits": true, + "model.embed_tokens": true, + "model.layers.0": true, + "model.layers.0.attention_layernorm": true, + "model.layers.0.self_attn": true, + "model.layers.0.self_attn.q_proj": true, + "model.layers.0.self_attn.q_norm": true, + "model.layers.0.self_attn.k_proj": true, + "model.layers.0.self_attn.k_norm": true, + "model.layers.0.self_attn.v_proj": true, + "model.layers.0.mlp.up_proj": true, + "model.layers.0.mlp.act_fn": true, + "model.layers.7": true, + "model.layers.15": true, + "model.layers.31": true, + "model.norm": true, + }) + + inputIDs := ref["input_ids"] + if inputIDs == nil { + panic("reference is missing input_ids") + } + tokens := inputIDs.AsType(mlx.DTypeInt32) + mlx.Pin(tokens) + defer mlx.Unpin(tokens) + mlx.Eval(tokens) + dims := tokens.Dims() + if len(dims) != 2 { + panic(fmt.Sprintf("input_ids shape = %v, want rank 2", dims)) + } + B, L := int32(dims[0]), int32(dims[1]) + + m := loadImportedModel(t, modelName) + b := &batch.Batch{ + InputIDs: tokens, + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{L}, + } + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + + h := m.EmbedTokens.Forward(tokens) + compareReference(t, "model.embed_tokens", h, ref["model.embed_tokens"], 0.9999) + + layer0 := m.Layers[0] + hNorm := layer0.AttentionNorm.Forward(h, m.RMSNormEps) + compareReference(t, "model.layers.0.attention_layernorm", hNorm, ref["model.layers.0.attention_layernorm"], 0.9999) + + q := layer0.Attention.QProj.Forward(hNorm) + compareReference(t, "model.layers.0.self_attn.q_proj", q, ref["model.layers.0.self_attn.q_proj"], 0.9999) + q = mlx.Reshape(q, B, L, m.NumAttentionHeads, m.HeadDim) + q = mlx.Transpose(q, 0, 2, 1, 3) + q = headRMSNorm(layer0.Attention.QNorm, q, B, m.NumAttentionHeads, L, m.HeadDim, m.RMSNormEps) + compareReference(t, "model.layers.0.self_attn.q_norm", q, ref["model.layers.0.self_attn.q_norm"], 0.9999) + + k := layer0.Attention.KProj.Forward(hNorm) + compareReference(t, "model.layers.0.self_attn.k_proj", k, ref["model.layers.0.self_attn.k_proj"], 0.9999) + k = mlx.Reshape(k, B, L, m.NumKeyValueHeads, m.HeadDim) + k = mlx.Transpose(k, 0, 2, 1, 3) + k = headRMSNorm(layer0.Attention.KNorm, k, B, m.NumKeyValueHeads, L, m.HeadDim, m.RMSNormEps) + compareReference(t, "model.layers.0.self_attn.k_norm", k, ref["model.layers.0.self_attn.k_norm"], 0.9999) + + v := layer0.Attention.VProj.Forward(hNorm) + compareReference(t, "model.layers.0.self_attn.v_proj", v, ref["model.layers.0.self_attn.v_proj"], 0.9999) + + attnOut := layer0.Attention.Forward(hNorm, b, nil, positions, B, L, m.Config) + compareReference(t, "model.layers.0.self_attn", attnOut, ref["model.layers.0.self_attn"], 0.999) + + h = mlx.Add(h, attnOut) + ffnInput := layer0.FFNNorm.Forward(h, m.RMSNormEps) + mlpUp := layer0.MLP.UpProj.Forward(ffnInput) + compareReference(t, "model.layers.0.mlp.up_proj", mlpUp, ref["model.layers.0.mlp.up_proj"], 0.999) + mlpAct := layer0.MLP.Act.Forward(mlpUp) + compareReference(t, "model.layers.0.mlp.act_fn", mlpAct, ref["model.layers.0.mlp.act_fn"], 0.997) + + h = mlx.Add(h, layer0.MLP.DownProj.Forward(mlpAct)) + compareReference(t, "model.layers.0", h, ref["model.layers.0"], 0.999) + + for i := 1; i < len(m.Layers); i++ { + h = m.Layers[i].Forward(h, b, nil, positions, B, L, m.Config) + switch i { + case 7, 15, 31: + key := fmt.Sprintf("model.layers.%d", i) + compareReference(t, key, h, ref[key], 0.999) + } + } + + h = m.Norm.Forward(h, m.RMSNormEps) + compareReference(t, "model.norm", h, ref["model.norm"], 0.997) + if ref["logits"] != nil { + logits := mlx.Reshape(m.LMHead.Forward(h), B, L, m.VocabSize) + compareReference(t, "logits", logits, ref["logits"], 0.997) + compareLogitArgmax(t, "logits", logits, ref["logits"]) + } +} + +func configureCompileForTest(t *testing.T) { + t.Helper() + switch strings.ToLower(os.Getenv("PORTING_APERTUS_COMPILE")) { + case "1", "true", "enable", "enabled", "on": + t.Log("enabling MLX compile") + mlx.EnableCompile() + case "0", "false", "disable", "disabled", "off": + t.Log("disabling MLX compile") + mlx.DisableCompile() + } +} + +func runForwardCacheReference(t *testing.T, refPath, modelName string) { + t.Helper() + ref := loadReferenceFiltered(t, refPath, map[string]bool{ + "input_ids": true, + "prefill_input_ids": true, + "logits": true, + "model.embed_tokens": true, + "model.layers.0": true, + "model.layers.0.self_attn": true, + "model.layers.0.self_attn.q_norm": true, + "model.layers.0.self_attn.k_norm": true, + "model.layers.0.mlp.up_proj": true, + "model.layers.0.mlp.act_fn": true, + "model.layers.7": true, + "model.layers.15": true, + "model.layers.31": true, + "model.norm": true, + }) + + prefillIDs := ref["prefill_input_ids"] + if prefillIDs == nil { + panic("reference is missing prefill_input_ids") + } + decodeIDs := ref["input_ids"] + if decodeIDs == nil { + panic("reference is missing input_ids") + } + prefill := prefillIDs.AsType(mlx.DTypeInt32) + decode := decodeIDs.AsType(mlx.DTypeInt32) + mlx.Pin(prefill, decode) + defer mlx.Unpin(prefill, decode) + mlx.Eval(prefill, decode) + + prefillDims := prefill.Dims() + decodeDims := decode.Dims() + if len(prefillDims) != 2 { + panic(fmt.Sprintf("prefill_input_ids shape = %v, want rank 2", prefillDims)) + } + if len(decodeDims) != 2 { + panic(fmt.Sprintf("input_ids shape = %v, want rank 2", decodeDims)) + } + if prefillDims[0] != 1 || decodeDims[0] != 1 { + panic(fmt.Sprintf("only batch size 1 is supported, got prefill=%v decode=%v", prefillDims, decodeDims)) + } + + m := loadImportedModel(t, modelName) + caches := m.NewCaches() + + prefillLen := int32(prefillDims[1]) + m.Forward(&batch.Batch{ + InputIDs: prefill, + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{prefillLen}, + }, caches) + mlx.Sweep() + + decodeLen := int32(decodeDims[1]) + decodeBatch := &batch.Batch{ + InputIDs: decode, + SeqOffsets: []int32{prefillLen}, + SeqQueryLens: []int32{decodeLen}, + } + positions := mlx.FromValues(decodeBatch.SeqOffsets, len(decodeBatch.SeqOffsets)) + + h := mlx.Reshape(m.EmbedTokens.Forward(decode), 1, decodeLen, m.HiddenSize) + compareReference(t, "model.embed_tokens", h, ref["model.embed_tokens"], 0.9999) + + layer0 := m.Layers[0] + hNorm := layer0.AttentionNorm.Forward(h, m.RMSNormEps) + q := layer0.Attention.QProj.Forward(hNorm) + q = mlx.Reshape(q, 1, decodeLen, m.NumAttentionHeads, m.HeadDim) + q = mlx.Transpose(q, 0, 2, 1, 3) + q = headRMSNorm(layer0.Attention.QNorm, q, 1, m.NumAttentionHeads, decodeLen, m.HeadDim, m.RMSNormEps) + compareReference(t, "model.layers.0.self_attn.q_norm", q, ref["model.layers.0.self_attn.q_norm"], 0.9999) + + k := layer0.Attention.KProj.Forward(hNorm) + k = mlx.Reshape(k, 1, decodeLen, m.NumKeyValueHeads, m.HeadDim) + k = mlx.Transpose(k, 0, 2, 1, 3) + k = headRMSNorm(layer0.Attention.KNorm, k, 1, m.NumKeyValueHeads, decodeLen, m.HeadDim, m.RMSNormEps) + compareReference(t, "model.layers.0.self_attn.k_norm", k, ref["model.layers.0.self_attn.k_norm"], 0.9999) + + attnOut := layer0.Attention.Forward(hNorm, decodeBatch, caches[0], positions, 1, decodeLen, m.Config) + compareReference(t, "model.layers.0.self_attn", attnOut, ref["model.layers.0.self_attn"], 0.999) + + h = mlx.Add(h, attnOut) + ffnInput := layer0.FFNNorm.Forward(h, m.RMSNormEps) + mlpUp := layer0.MLP.UpProj.Forward(ffnInput) + compareReference(t, "model.layers.0.mlp.up_proj", mlpUp, ref["model.layers.0.mlp.up_proj"], 0.999) + mlpAct := layer0.MLP.Act.Forward(mlpUp) + compareReference(t, "model.layers.0.mlp.act_fn", mlpAct, ref["model.layers.0.mlp.act_fn"], 0.997) + + h = mlx.Add(h, layer0.MLP.DownProj.Forward(mlpAct)) + compareReference(t, "model.layers.0", h, ref["model.layers.0"], 0.999) + + for i := 1; i < len(m.Layers); i++ { + h = m.Layers[i].Forward(h, decodeBatch, caches[i], positions, 1, decodeLen, m.Config) + switch i { + case 7, 15, 31: + key := fmt.Sprintf("model.layers.%d", i) + compareReference(t, key, h, ref[key], 0.999) + } + } + + h = m.Norm.Forward(h, m.RMSNormEps) + compareReference(t, "model.norm", h, ref["model.norm"], 0.997) + if ref["logits"] != nil { + logits := mlx.Reshape(m.LMHead.Forward(h), 1, decodeLen, m.VocabSize) + compareReference(t, "logits", logits, ref["logits"], 0.997) + compareLogitArgmax(t, "logits", logits, ref["logits"]) + } +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func loadImportedModel(t *testing.T, modelName string) *Model { + t.Helper() + + root, err := model.Open(modelName) + if err != nil { + panic(fmt.Sprintf("open imported model %q: %v", modelName, err)) + } + defer root.Close() + + bm, err := base.New(root) + if err != nil { + panic(fmt.Sprintf("construct model: %v", err)) + } + m, ok := bm.(*Model) + if !ok { + panic(fmt.Sprintf("expected *apertus.Model, got %T", bm)) + } + + tensors := loadTensorsForTest(t, root) + if err := m.LoadWeights(tensors); err != nil { + panic(fmt.Sprintf("load weights: %v", err)) + } + + collected := mlx.Collect(m) + for _, arr := range collected { + mlx.Pin(arr) + } + mlx.Sweep() + mlx.Eval(collected...) + + return m +} + +func loadTensorsForTest(t *testing.T, root *model.Root) map[string]*mlx.Array { + t.Helper() + + raw := make(map[string]*mlx.Array) + seen := make(map[string]bool) + for _, layer := range root.Manifest.GetTensorLayers("") { + if seen[layer.Digest] { + continue + } + seen[layer.Digest] = true + for name, arr := range mlx.Load(root.Manifest.BlobPath(layer.Digest)) { + raw[name] = arr + } + } + + scaleBaseNames := make(map[string]bool) + out := make(map[string]*mlx.Array, len(raw)) + for name, arr := range raw { + if strings.HasSuffix(name, ".scale") { + baseName := strings.TrimSuffix(name, ".scale") + out[baseName+"_scale"] = arr + scaleBaseNames[baseName] = true + } + } + for name, arr := range raw { + if strings.HasSuffix(name, ".scale") { + continue + } + if strings.HasSuffix(name, ".bias") && !strings.HasSuffix(name, ".weight_qbias") { + baseName := strings.TrimSuffix(name, ".bias") + if scaleBaseNames[baseName] { + out[baseName+"_qbias"] = arr + continue + } + } + out[name] = arr + } + return out +} + +func loadReferenceFiltered(t *testing.T, path string, keep map[string]bool) map[string]*mlx.Array { + t.Helper() + if _, err := os.Stat(path); err != nil { + panic(fmt.Sprintf("reference data not available at %s: %v", path, err)) + } + + out := make(map[string]*mlx.Array) + for name, arr := range mlx.Load(path) { + if keep[name] { + out[name] = arr + } + } + if len(out) == 0 { + panic(fmt.Sprintf("no reference tensors loaded from %s", path)) + } + for _, arr := range out { + mlx.Pin(arr) + } + mlx.Sweep() + return out +} + +func compareReference(t *testing.T, name string, got, want *mlx.Array, minCos float64) { + t.Helper() + if want == nil { + panic(fmt.Sprintf("reference is missing %q", name)) + } + + wantShape := want.Dims() + shape := make([]int32, len(wantShape)) + for i, dim := range wantShape { + shape[i] = int32(dim) + } + got = mlx.Reshape(got, shape...) + got = got.AsType(mlx.DTypeFloat32) + want = want.AsType(mlx.DTypeFloat32) + + size := 1 + for _, dim := range wantShape { + size *= dim + } + gotFlat := mlx.Reshape(got, int32(size)) + wantFlat := mlx.Reshape(want, int32(size)) + diff := mlx.Sub(gotFlat, wantFlat).Abs() + dot := mlx.Sum(mlx.Mul(gotFlat, wantFlat), 0, false) + gotNorm := mlx.Sum(mlx.Mul(gotFlat, gotFlat), 0, false).Sqrt() + wantNorm := mlx.Sum(mlx.Mul(wantFlat, wantFlat), 0, false).Sqrt() + cosine := mlx.Div(dot, mlx.Mul(gotNorm, wantNorm)) + maxDiff := diff.MaxAxis(0, false) + meanDiff := mlx.Mean(diff, 0, false) + mlx.Eval(cosine, maxDiff, meanDiff) + + cos := scalarFloat(cosine) + maxDiffValue := scalarFloat(maxDiff) + meanDiffValue := scalarFloat(meanDiff) + if maxDiffValue <= 1e-6 { + t.Logf("%s: shape=%v exact max_diff=%.6g mean_diff=%.6g", name, wantShape, maxDiffValue, meanDiffValue) + return + } + t.Logf("%s: shape=%v cos=%.8f max_diff=%.6g mean_diff=%.6g", name, wantShape, cos, maxDiffValue, meanDiffValue) + if math.IsNaN(cos) || cos < minCos { + t.Errorf("%s cosine similarity %.8f below %.8f", name, cos, minCos) + } +} + +func compareLogitArgmax(t *testing.T, name string, got, want *mlx.Array) { + t.Helper() + if want == nil { + panic(fmt.Sprintf("reference is missing %q", name)) + } + wantShape := want.Dims() + if len(wantShape) != 3 { + panic(fmt.Sprintf("%s shape = %v, want [B,L,V]", name, wantShape)) + } + shape := make([]int32, len(wantShape)) + for i, dim := range wantShape { + shape[i] = int32(dim) + } + got = mlx.Reshape(got, shape...) + last := wantShape[1] - 1 + gotLast := got.Slice(mlx.Slice(), mlx.Slice(last, last+1), mlx.Slice()).Squeeze(1) + wantLast := want.Slice(mlx.Slice(), mlx.Slice(last, last+1), mlx.Slice()).Squeeze(1) + gotID := gotLast.Argmax(-1, false).AsType(mlx.DTypeInt32) + wantID := wantLast.Argmax(-1, false).AsType(mlx.DTypeInt32) + mlx.Eval(gotID, wantID) + gotValue := gotID.Int() + wantValue := wantID.Int() + t.Logf("%s last-token argmax got=%d want=%d", name, gotValue, wantValue) + if gotValue != wantValue { + t.Errorf("%s last-token argmax = %d, want %d", name, gotValue, wantValue) + } +} + +func scalarFloat(x *mlx.Array) float64 { + x = x.AsType(mlx.DTypeFloat32) + mlx.Eval(x) + values := x.Floats() + if len(values) == 0 { + return math.NaN() + } + return float64(values[0]) +} diff --git a/x/models/apertus/import_shape_test.go b/x/models/apertus/import_shape_test.go new file mode 100644 index 00000000000..52f3af22c29 --- /dev/null +++ b/x/models/apertus/import_shape_test.go @@ -0,0 +1,145 @@ +package apertus + +import ( + "encoding/binary" + "encoding/json" + "io" + "os" + "reflect" + "testing" + + imagemanifest "github.com/ollama/ollama/x/imagegen/manifest" + "github.com/ollama/ollama/x/tokenizer" +) + +func TestImportedApertusTensorShapes(t *testing.T) { + if os.Getenv("OLLAMA_MODELS") == "" { + t.Skip("set OLLAMA_MODELS to the imported model cache to validate imported tensor shapes") + } + + m, err := imagemanifest.LoadManifest("apertus-mlx") + if err != nil { + t.Fatalf("load imported manifest: %v", err) + } + + got := map[string][]int{} + for _, layer := range m.GetTensorLayers("") { + header, err := readSafetensorsHeader(m.BlobPath(layer.Digest)) + if err != nil { + t.Fatalf("read tensor layer %s: %v", layer.Name, err) + } + for name, meta := range header { + if name == "__metadata__" { + continue + } + var info struct { + DType string `json:"dtype"` + Shape []int `json:"shape"` + } + if err := json.Unmarshal(meta, &info); err != nil { + t.Fatalf("parse tensor %s metadata: %v", name, err) + } + got[name] = info.Shape + } + } + + want := map[string][]int{ + "model.embed_tokens.weight": {131072, 4096}, + "lm_head.weight": {131072, 4096}, + "model.norm.weight": {4096}, + "model.layers.0.attention_layernorm.weight": {4096}, + "model.layers.0.feedforward_layernorm.weight": {4096}, + "model.layers.0.self_attn.q_proj.weight": {4096, 4096}, + "model.layers.0.self_attn.k_proj.weight": {1024, 4096}, + "model.layers.0.self_attn.v_proj.weight": {1024, 4096}, + "model.layers.0.self_attn.o_proj.weight": {4096, 4096}, + "model.layers.0.self_attn.q_norm.weight": {128}, + "model.layers.0.self_attn.k_norm.weight": {128}, + "model.layers.0.mlp.up_proj.weight": {21504, 4096}, + "model.layers.0.mlp.down_proj.weight": {4096, 21504}, + "model.layers.0.mlp.act_fn.alpha_p": {1}, + "model.layers.0.mlp.act_fn.alpha_n": {1}, + "model.layers.0.mlp.act_fn.beta": {}, + "model.layers.0.mlp.act_fn.eps": {}, + "model.layers.31.self_attn.q_proj.weight": {4096, 4096}, + "model.layers.31.self_attn.k_proj.weight": {1024, 4096}, + "model.layers.31.mlp.down_proj.weight": {4096, 21504}, + "model.layers.31.mlp.act_fn.eps": {}, + } + + for name, wantShape := range want { + gotShape, ok := got[name] + if !ok { + t.Fatalf("imported tensor %s missing", name) + } + if !reflect.DeepEqual(gotShape, wantShape) { + t.Fatalf("imported tensor %s shape = %v, want %v", name, gotShape, wantShape) + } + } + if len(got) != 451 { + t.Fatalf("imported tensor count = %d, want 451", len(got)) + } +} + +func TestImportedApertusEOSTokens(t *testing.T) { + if os.Getenv("OLLAMA_MODELS") == "" { + t.Skip("set OLLAMA_MODELS to the imported model cache to validate imported EOS tokens") + } + + m, err := imagemanifest.LoadManifest("apertus-mlx") + if err != nil { + t.Fatalf("load imported manifest: %v", err) + } + + tokData, err := m.ReadConfig("tokenizer.json") + if err != nil { + t.Fatalf("read tokenizer.json: %v", err) + } + configData, err := m.ReadConfig("config.json") + if err != nil { + t.Fatalf("read config.json: %v", err) + } + tokConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData} + if genConfigData, err := m.ReadConfig("generation_config.json"); err == nil { + tokConfig.GenerationConfigJSON = genConfigData + } + if tokConfigData, err := m.ReadConfig("tokenizer_config.json"); err == nil { + tokConfig.TokenizerConfigJSON = tokConfigData + } + if specialTokensMapData, err := m.ReadConfig("special_tokens_map.json"); err == nil { + tokConfig.SpecialTokensMapJSON = specialTokensMapData + } + + tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) + if err != nil { + t.Fatalf("load tokenizer: %v", err) + } + + want := []int32{2, 68, 72} + if got := tok.EOSTokens(); !reflect.DeepEqual(got, want) { + t.Fatalf("EOSTokens() = %v, want %v", got, want) + } +} + +func readSafetensorsHeader(path string) (map[string]json.RawMessage, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + var headerSize uint64 + if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil { + return nil, err + } + header := make([]byte, headerSize) + if _, err := io.ReadFull(f, header); err != nil { + return nil, err + } + + var out map[string]json.RawMessage + if err := json.Unmarshal(header, &out); err != nil { + return nil, err + } + return out, nil +} From 05db22150249606843168d06fca9be60f01d76aa Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 19:46:35 +0000 Subject: [PATCH 32/58] model: add Apertus chat rendering and parsing Co-authored-by: Codex --- model/parsers/apertus.go | 323 +++++++++++++++++++++++ model/parsers/apertus_test.go | 205 +++++++++++++++ model/parsers/parsers.go | 2 + model/renderers/apertus.go | 440 ++++++++++++++++++++++++++++++++ model/renderers/apertus_test.go | 138 ++++++++++ model/renderers/renderer.go | 2 + 6 files changed, 1110 insertions(+) create mode 100644 model/parsers/apertus.go create mode 100644 model/parsers/apertus_test.go create mode 100644 model/renderers/apertus.go create mode 100644 model/renderers/apertus_test.go diff --git a/model/parsers/apertus.go b/model/parsers/apertus.go new file mode 100644 index 00000000000..c7396f577ba --- /dev/null +++ b/model/parsers/apertus.go @@ -0,0 +1,323 @@ +package parsers + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "strings" + "unicode" + + "github.com/ollama/ollama/api" +) + +const ( + apertusToolOpenTag = "<|tools_prefix|>" + apertusToolCloseTag = "<|tools_suffix|>" + apertusAssistantOpenTag = "<|assistant_start|>" + apertusAssistantCloseTag = "<|assistant_end|>" + // A tool call is one model response fragment. Keep malformed streams from + // retaining unbounded output while allowing substantially larger calls than + // the model's normal tool grammar produces. + maxApertusToolCallBytes = 1 << 20 +) + +type apertusParserState uint8 + +const ( + apertusContent apertusParserState = iota + apertusToolCalls +) + +type ApertusParser struct { + state apertusParserState + acc strings.Builder + allowedTool map[string]struct{} + initErr error + callIndex int + pendingBare bool +} + +func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool { + p.state = apertusContent + p.acc.Reset() + p.allowedTool = make(map[string]struct{}, len(tools)) + p.initErr = nil + p.callIndex = 0 + p.pendingBare = false + for _, tool := range tools { + name := tool.Function.Name + if !apertusIdentifier(name) { + p.initErr = fmt.Errorf("invalid apertus tool name %q", name) + continue + } + if _, exists := p.allowedTool[name]; exists { + p.initErr = fmt.Errorf("duplicate apertus tool name %q", name) + continue + } + p.allowedTool[name] = struct{}{} + } + return tools +} + +func (p *ApertusParser) Add(s string, done bool) (content, thinking string, calls []api.ToolCall, err error) { + if p.initErr != nil { + return "", "", nil, p.initErr + } + if p.pendingFragmentWouldExceed(s) { + return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) + } + p.acc.WriteString(s) + if p.state == apertusToolCalls && p.acc.Len() > maxApertusToolCallBytes { + return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) + } + + var out strings.Builder + for { + current := p.acc.String() + switch p.state { + case apertusContent: + if idx := strings.Index(current, apertusToolOpenTag); idx >= 0 { + p.pendingBare = false + out.WriteString(cleanApertusContent(current[:idx])) + p.acc.Reset() + p.acc.WriteString(current[idx+len(apertusToolOpenTag):]) + p.state = apertusToolCalls + if p.acc.Len() > maxApertusToolCallBytes { + return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) + } + continue + } + if done { + cleaned := cleanApertusContent(current) + if p.looksLikeToolCall(cleaned) { + if parsed, parseErr := p.parseToolCalls(cleaned); parseErr == nil { + p.acc.Reset() + p.pendingBare = false + return out.String(), "", parsed, nil + } else if !isSoftApertusToolParseError(parseErr) { + return "", "", nil, parseErr + } + } + p.acc.Reset() + p.pendingBare = false + out.WriteString(cleaned) + return out.String(), "", calls, nil + } + if p.looksLikeToolCallStart(current) { + p.pendingBare = true + if p.acc.Len() > maxApertusToolCallBytes { + return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) + } + return out.String(), "", nil, nil + } + p.pendingBare = false + overlapLen := max( + overlap(current, apertusToolOpenTag), + overlap(current, apertusAssistantOpenTag), + overlap(current, apertusAssistantCloseTag), + ) + n := len(current) - overlapLen + if n == len(current) { + n -= trailingWhitespaceLen(current) + } + if n > 0 { + emit := current[:n] + if n < len(current) { + emit = strings.TrimRightFunc(emit, unicode.IsSpace) + } + out.WriteString(cleanApertusContent(emit)) + p.acc.Reset() + p.acc.WriteString(current[n:]) + } + return out.String(), "", calls, nil + + case apertusToolCalls: + if idx := strings.Index(current, apertusToolCloseTag); idx >= 0 { + parsed, parseErr := p.parseToolCalls(current[:idx]) + if parseErr != nil { + if !isSoftApertusToolParseError(parseErr) { + return "", "", nil, parseErr + } + out.WriteString(cleanApertusContent(current[:idx])) + } else { + calls = append(calls, parsed...) + } + p.acc.Reset() + p.acc.WriteString(strings.TrimLeftFunc(current[idx+len(apertusToolCloseTag):], unicode.IsSpace)) + p.state = apertusContent + p.pendingBare = false + continue + } + if done { + parsed, parseErr := p.parseToolCalls(current) + if parseErr != nil { + if !isSoftApertusToolParseError(parseErr) { + return "", "", nil, fmt.Errorf("unterminated apertus tool call: %w", parseErr) + } + out.WriteString(cleanApertusContent(current)) + } else { + calls = append(calls, parsed...) + } + p.acc.Reset() + p.state = apertusContent + p.pendingBare = false + return out.String(), "", calls, nil + } + return out.String(), "", calls, nil + } + } +} + +func (p *ApertusParser) pendingFragmentWouldExceed(s string) bool { + if p.acc.Len() > maxApertusToolCallBytes { + return true + } + if p.state == apertusToolCalls || p.pendingBare { + return len(s) > maxApertusToolCallBytes-p.acc.Len() + } + if p.acc.Len() > 0 && strings.HasPrefix(apertusToolOpenTag, p.acc.String()) { + remainingOpener := len(apertusToolOpenTag) - p.acc.Len() + return len(s) > remainingOpener+maxApertusToolCallBytes + } + + probeLen := min(len(s), maxApertusToolCallBytes+1-p.acc.Len()) + probe := p.acc.String() + s[:probeLen] + if idx := strings.Index(probe, apertusToolOpenTag); idx >= 0 { + return p.acc.Len()+len(s)-idx-len(apertusToolOpenTag) > maxApertusToolCallBytes + } + return apertusPendingBarePrefix(probe) && p.acc.Len()+len(s) > maxApertusToolCallBytes +} + +func apertusPendingBarePrefix(s string) bool { + s = strings.TrimLeftFunc(s, unicode.IsSpace) + for { + switch { + case strings.HasPrefix(s, apertusAssistantOpenTag): + s = strings.TrimLeftFunc(strings.TrimPrefix(s, apertusAssistantOpenTag), unicode.IsSpace) + case strings.HasPrefix(s, apertusAssistantCloseTag): + s = strings.TrimLeftFunc(strings.TrimPrefix(s, apertusAssistantCloseTag), unicode.IsSpace) + default: + return strings.TrimSpace(s) == "" || strings.HasPrefix("[{", s) || strings.HasPrefix("{", s) || strings.HasPrefix(s, "[{") || strings.HasPrefix(s, "{") + } + } +} + +func cleanApertusContent(s string) string { + s = strings.ReplaceAll(s, apertusAssistantOpenTag, "") + s = strings.ReplaceAll(s, apertusAssistantCloseTag, "") + return strings.TrimRightFunc(s, unicode.IsSpace) +} +func (p *ApertusParser) HasToolSupport() bool { return true } +func (p *ApertusParser) HasThinkingSupport() bool { return false } +func (p *ApertusParser) PreservedTokens() []string { + return []string{apertusToolOpenTag, apertusToolCloseTag, apertusAssistantOpenTag, apertusAssistantCloseTag} +} + +func (p *ApertusParser) parseToolCalls(raw string) ([]api.ToolCall, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, errors.New("empty apertus tool call") + } + var entries []json.RawMessage + if raw[0] == '[' { + if err := json.Unmarshal([]byte(raw), &entries); err != nil { + return nil, err + } + } else { + entries = []json.RawMessage{json.RawMessage(raw)} + } + var calls []api.ToolCall + for _, entry := range entries { + name, rawArgs, err := apertusCallEntry(entry) + if err != nil { + return nil, err + } + if _, ok := p.allowedTool[name]; !ok { + return nil, fmt.Errorf("unknown apertus tool %q", name) + } + args := api.NewToolCallFunctionArguments() + if len(rawArgs) > 0 && !bytes.Equal(rawArgs, []byte("null")) { + if err := json.Unmarshal(rawArgs, &args); err != nil { + var encoded string + if stringErr := json.Unmarshal(rawArgs, &encoded); stringErr != nil { + return nil, err + } + if err := json.Unmarshal([]byte(encoded), &args); err != nil { + return nil, err + } + } + } + calls = append(calls, api.ToolCall{Function: api.ToolCallFunction{Index: p.callIndex, Name: name, Arguments: args}}) + p.callIndex++ + } + return calls, nil +} + +func apertusCallEntry(raw json.RawMessage) (string, json.RawMessage, error) { + dec := json.NewDecoder(bytes.NewReader(raw)) + tok, err := dec.Token() + if err != nil { + return "", nil, err + } + if d, ok := tok.(json.Delim); !ok || d != '{' { + return "", nil, errors.New("apertus tool call must be an object") + } + if !dec.More() { + return "", nil, errors.New("apertus tool call object must contain exactly one function name") + } + nameTok, err := dec.Token() + if err != nil { + return "", nil, err + } + name, ok := nameTok.(string) + if !ok { + return "", nil, errors.New("invalid apertus tool name") + } + var args json.RawMessage + if err := dec.Decode(&args); err != nil { + return "", nil, err + } + if dec.More() { + return "", nil, errors.New("apertus tool call object must contain exactly one function name") + } + if _, err := dec.Token(); err != nil { + return "", nil, err + } + if dec.More() { + return "", nil, errors.New("invalid trailing apertus tool call data") + } + return name, args, nil +} + +func isSoftApertusToolParseError(err error) bool { + var syntaxErr *json.SyntaxError + if errors.As(err, &syntaxErr) { + return true + } + var typeErr *json.UnmarshalTypeError + return errors.As(err, &typeErr) +} + +func (p *ApertusParser) looksLikeToolCall(s string) bool { + s = strings.TrimSpace(s) + return len(p.allowedTool) > 0 && (strings.HasPrefix(s, "[{") || strings.HasPrefix(s, "{")) +} + +func (p *ApertusParser) looksLikeToolCallStart(s string) bool { + s = strings.TrimSpace(s) + return len(p.allowedTool) > 0 && s != "" && (strings.HasPrefix("[{", s) || strings.HasPrefix("{", s) || strings.HasPrefix(s, "[{") || strings.HasPrefix(s, "{")) +} + +func apertusIdentifier(s string) bool { + if s == "" { + return false + } + for i := range len(s) { + c := s[i] + if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_' || c == '$' || (i > 0 && c >= '0' && c <= '9')) { + return false + } + } + return true +} diff --git a/model/parsers/apertus_test.go b/model/parsers/apertus_test.go new file mode 100644 index 00000000000..a2fa77ef4d3 --- /dev/null +++ b/model/parsers/apertus_test.go @@ -0,0 +1,205 @@ +package parsers + +import ( + "strings" + "testing" + + "github.com/ollama/ollama/api" +) + +func TestApertusParserGrammarAndStreamingBoundaries(t *testing.T) { + input := `prefix <|tools_prefix|>[{"get_weather":{"city":"Bern"}}]<|tools_suffix|><|assistant_end|>` + for i := range len(input) + 1 { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("get_weather")}, nil, nil) + content, _, calls, err := p.Add(input[:min(i, len(input))], false) + if err != nil { + t.Fatalf("split %d: %v", i, err) + } + content2, _, calls2, err := p.Add(input[min(i, len(input)):], true) + if err != nil { + t.Fatalf("split %d: %v", i, err) + } + if content+content2 != "prefix" || len(calls)+len(calls2) != 1 { + t.Fatalf("split %d: content=%q calls=%d", i, content+content2, len(calls)+len(calls2)) + } + } +} + +func TestApertusParserMultipleUnknownMalformedAndDuplicateDeclarations(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one"), apertusParserTool("two")}, nil, nil) + _, _, calls, err := p.Add(`<|tools_prefix|>[{"one":{}},{"two":{"x":1}}]<|tools_suffix|>`, true) + if err != nil || len(calls) != 2 || calls[0].Function.Index != 0 || calls[1].Function.Index != 1 { + t.Fatalf("calls=%#v err=%v", calls, err) + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(`<|tools_prefix|>[{"other":{}}]<|tools_suffix|>`, true); err == nil { + t.Fatal("unknown call accepted") + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + content, _, calls, err := p.Add(`<|tools_prefix|>[{"one":}]<|tools_suffix|>`, true) + if err != nil || content != `[{"one":}]` || len(calls) != 0 { + t.Fatalf("malformed content=%q calls=%d err=%v", content, len(calls), err) + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("same"), apertusParserTool("same")}, nil, nil) + if _, _, _, err := p.Add("anything", true); err == nil { + t.Fatal("duplicate declaration accepted") + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("bad.name")}, nil, nil) + if _, _, _, err := p.Add("anything", true); err == nil { + t.Fatal("separator-bearing declaration accepted") + } +} + +func TestApertusParserBoundsToolPayload(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(apertusToolOpenTag+strings.Repeat("x", maxApertusToolCallBytes+1), false); err == nil { + t.Fatal("oversized tool call accepted") + } +} + +func TestApertusParserBoundsBareToolPayload(t *testing.T) { + for _, prefix := range []string{"[", "{"} { + t.Run(prefix, func(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(prefix, false); err != nil { + t.Fatal(err) + } + if _, _, _, err := p.Add(strings.Repeat("x", maxApertusToolCallBytes-len(prefix)), true); err != nil { + t.Fatalf("exact limit rejected: %v", err) + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(prefix, false); err != nil { + t.Fatal(err) + } + if _, _, _, err := p.Add(strings.Repeat("x", maxApertusToolCallBytes-len(prefix)+1), false); err == nil { + t.Fatal("first over-limit byte accepted") + } + }) + } +} + +func TestApertusParserBoundsRetainedPrefixesBeforeAppend(t *testing.T) { + tests := []struct { + name string + prefix string + body string + }{ + {"whitespace", " ", "[{"}, + {"assistant-token", "<|assistant_", "end|>[{"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(tt.prefix, false); err != nil { + t.Fatal(err) + } + exact := tt.body + strings.Repeat("x", maxApertusToolCallBytes-p.acc.Len()-len(tt.body)) + if _, _, _, err := p.Add(exact, true); err != nil { + t.Fatalf("exact limit rejected: %v", err) + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if _, _, _, err := p.Add(tt.prefix, false); err != nil { + t.Fatal(err) + } + retained := p.acc.Len() + over := tt.body + strings.Repeat("x", maxApertusToolCallBytes-retained-len(tt.body)+1) + if _, _, _, err := p.Add(over, false); err == nil { + t.Fatal("first over-limit byte accepted") + } + if p.acc.Len() != retained { + t.Fatalf("retained %d bytes after rejection, want %d", p.acc.Len(), retained) + } + if _, _, _, err := p.Add("", true); err != nil { + t.Fatalf("retained prefix did not recover: %v", err) + } + }) + } +} + +func TestApertusParserTaggedPayloadBoundaryBeforeAppend(t *testing.T) { + tests := []struct { + name string + prefix string + body string + }{ + {"complete", "", apertusToolOpenTag}, + } + for i := 1; i < len(apertusToolOpenTag); i++ { + tests = append(tests, struct { + name string + prefix string + body string + }{name: "split", prefix: apertusToolOpenTag[:i], body: apertusToolOpenTag[i:]}) + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if tt.prefix != "" { + if _, _, _, err := p.Add(tt.prefix, false); err != nil { + t.Fatal(err) + } + } + exact := tt.body + strings.Repeat("x", maxApertusToolCallBytes) + if _, _, _, err := p.Add(exact, true); err != nil { + t.Fatalf("exact payload limit rejected: %v", err) + } + p = &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + if tt.prefix != "" { + if _, _, _, err := p.Add(tt.prefix, false); err != nil { + t.Fatal(err) + } + } + retained := p.acc.Len() + over := tt.body + strings.Repeat("x", maxApertusToolCallBytes+1) + if _, _, _, err := p.Add(over, false); err == nil { + t.Fatal("first payload byte over limit accepted") + } + if p.acc.Len() != retained { + t.Fatalf("retained %d bytes after rejection, want %d", p.acc.Len(), retained) + } + }) + } +} + +func TestApertusParserBareAndTaggedHardErrorsMatch(t *testing.T) { + cases := []string{ + `[{"other":{}}]`, + `[{"one":{},"two":{}}]`, + } + for _, raw := range cases { + t.Run(raw, func(t *testing.T) { + for _, framed := range []string{raw, apertusToolOpenTag + raw + apertusToolCloseTag} { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one"), apertusParserTool("two")}, nil, nil) + if _, _, _, err := p.Add(framed, true); err == nil { + t.Fatalf("hard-invalid %q accepted", framed) + } + } + }) + } + for _, framed := range []string{`[{"one":}]`, apertusToolOpenTag + `[{"one":}]` + apertusToolCloseTag} { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, nil) + content, _, calls, err := p.Add(framed, true) + if err != nil || content != `[{"one":}]` || len(calls) != 0 { + t.Fatalf("malformed %q content=%q calls=%d err=%v", framed, content, len(calls), err) + } + } +} + +func apertusParserTool(name string) api.Tool { + return api.Tool{Type: "function", Function: api.ToolFunction{Name: name}} +} diff --git a/model/parsers/parsers.go b/model/parsers/parsers.go index 4d2ace546e1..cd9062d414a 100644 --- a/model/parsers/parsers.go +++ b/model/parsers/parsers.go @@ -48,6 +48,8 @@ func ParserForName(name string) Parser { var p Parser switch name { + case "apertus": + p = &ApertusParser{} case "qwen3": p = &Qwen3Parser{hasThinkingSupport: false, defaultThinking: false} case "qwen3-thinking": diff --git a/model/renderers/apertus.go b/model/renderers/apertus.go new file mode 100644 index 00000000000..86e377bd277 --- /dev/null +++ b/model/renderers/apertus.go @@ -0,0 +1,440 @@ +package renderers + +import ( + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/ollama/ollama/api" +) + +const ( + apertusSystemStart = "<|system_start|>" + apertusSystemEnd = "<|system_end|>" + apertusDeveloperStart = "<|developer_start|>" + apertusDeveloperEnd = "<|developer_end|>" + apertusUserStart = "<|user_start|>" + apertusUserEnd = "<|user_end|>" + apertusAssistantStart = "<|assistant_start|>" + apertusAssistantEnd = "<|assistant_end|>" + apertusToolsPrefix = "<|tools_prefix|>" + apertusToolsSuffix = "<|tools_suffix|>" + apertusImageToken = "<|image|>" + maxApertusSchemaDepth = 32 + maxApertusSchemaNodes = 4096 +) + +type ApertusRenderer struct{} + +func (r *ApertusRenderer) LeadingBOS() string { return "" } + +func (r *ApertusRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) { + if think != nil && think.Bool() { + return "", fmt.Errorf("apertus renderer does not support thinking") + } + if err := validateApertusTools(tools); err != nil { + return "", err + } + declaredTools := make(map[string]struct{}, len(tools)) + for _, tool := range tools { + declaredTools[tool.Function.Name] = struct{}{} + } + for _, message := range messages { + if err := validateApertusText(message.Content); err != nil { + return "", err + } + } + var sb strings.Builder + start := 0 + if len(messages) > 0 && messages[0].Role == "system" { + sb.WriteString(apertusSystemStart) + sb.WriteString(r.renderContent(messages[0])) + sb.WriteString(apertusSystemEnd) + start = 1 + } else { + sb.WriteString(apertusSystemStart) + sb.WriteString("You are Apertus, a helpful assistant created by the SwissAI initiative.\nKnowledge cutoff: 2024-04\nCurrent date: ") + sb.WriteString(time.Now().Format("2006-01-02")) + sb.WriteString(apertusSystemEnd) + } + sb.WriteString(apertusDeveloperStart) + sb.WriteString("Deliberation: disabled\n") + if len(tools) == 0 { + sb.WriteString("Tool Capabilities: disabled") + } else { + sb.WriteString("Tool Capabilities:\n") + renderApertusTools(&sb, tools) + } + sb.WriteString(apertusDeveloperEnd) + inAssistant, inTool := false, false + pendingToolResults := 0 + toolResultsStarted := false + closeAssistant := func() { + if inTool { + sb.WriteString("]") + inTool = false + } + if inAssistant { + sb.WriteString(apertusAssistantEnd) + inAssistant = false + } + } + for _, message := range messages[start:] { + switch message.Role { + case "user", "system": + if pendingToolResults > 0 && toolResultsStarted { + return "", fmt.Errorf("apertus tool results are incomplete") + } + closeAssistant() + if message.Role == "user" { + sb.WriteString(apertusUserStart) + sb.WriteString(r.renderContent(message)) + sb.WriteString(apertusUserEnd) + } else { + sb.WriteString(apertusSystemStart) + sb.WriteString(r.renderContent(message)) + sb.WriteString(apertusSystemEnd) + } + pendingToolResults = 0 + toolResultsStarted = false + case "assistant": + if pendingToolResults > 0 && toolResultsStarted { + return "", fmt.Errorf("apertus tool results are incomplete") + } + pendingToolResults = 0 + toolResultsStarted = false + if !inAssistant { + sb.WriteString(apertusAssistantStart) + inAssistant = true + } + if inTool { + sb.WriteString("]") + inTool = false + } + sb.WriteString(message.Content) + if len(message.ToolCalls) > 0 { + if err := validateApertusHistoricalCalls(message.ToolCalls, declaredTools); err != nil { + return "", err + } + if err := renderApertusToolCalls(&sb, message.ToolCalls); err != nil { + return "", err + } + pendingToolResults = len(message.ToolCalls) + } + case "tool": + if !inAssistant || pendingToolResults == 0 { + return "", fmt.Errorf("apertus tool message does not follow an assistant tool call") + } + if !inTool { + sb.WriteString("[") + inTool = true + } else { + sb.WriteString(", ") + } + sb.WriteString(message.Content) + toolResultsStarted = true + pendingToolResults-- + default: + return "", fmt.Errorf("unsupported apertus message role %q", message.Role) + } + } + if inTool { + sb.WriteString("]") + } + if inAssistant && pendingToolResults == 0 { + sb.WriteString(apertusAssistantEnd) + } + last := "" + if len(messages) > 0 { + last = messages[len(messages)-1].Role + } + if last != "assistant" && !(len(tools) > 0 && last == "user") { + sb.WriteString(apertusAssistantStart) + } + return sb.String(), nil +} + +func validateApertusHistoricalCalls(calls []api.ToolCall, declared map[string]struct{}) error { + if len(declared) == 0 { + return fmt.Errorf("apertus assistant tool calls require declarations") + } + for _, call := range calls { + name := call.Function.Name + if _, ok := declared[name]; !ok { + return fmt.Errorf("undeclared apertus assistant tool %q", name) + } + } + return nil +} + +func (r *ApertusRenderer) renderContent(message api.Message) string { + return strings.Repeat(apertusImageToken, len(message.Images)) + message.Content +} + +func renderApertusTools(sb *strings.Builder, tools []api.Tool) { + for i, tool := range tools { + if tool.Function.Description != "" { + sb.WriteString("// ") + sb.WriteString(tool.Function.Description) + sb.WriteString("\n") + } + sb.WriteString("type ") + sb.WriteString(tool.Function.Name) + if tool.Function.Parameters.Properties == nil || tool.Function.Parameters.Properties.Len() == 0 { + sb.WriteString(" = () => any;") + } else { + sb.WriteString(" = (_: {\n") + required := make(map[string]bool, len(tool.Function.Parameters.Required)) + for _, n := range tool.Function.Parameters.Required { + required[n] = true + } + n := 0 + total := tool.Function.Parameters.Properties.Len() + for name, prop := range tool.Function.Parameters.Properties.All() { + if prop.Description != "" { + sb.WriteString("// ") + sb.WriteString(prop.Description) + sb.WriteString("\n") + } + sb.WriteString(name) + if !required[name] { + sb.WriteString("?") + } + sb.WriteString(": ") + sb.WriteString(apertusTypeScriptType(prop)) + n++ + if n < total { + sb.WriteString(",\n") + } else { + sb.WriteString("\n") + } + } + sb.WriteString("}) => any;") + } + if i < len(tools)-1 { + sb.WriteString("\n") + } + } +} + +func apertusTypeScriptType(prop api.ToolProperty) string { + if len(prop.AnyOf) > 0 { + parts := make([]string, 0, len(prop.AnyOf)) + for _, p := range prop.AnyOf { + parts = append(parts, apertusTypeScriptType(p)) + } + return strings.Join(parts, " | ") + } + if len(prop.Enum) > 0 { + parts := make([]string, 0, len(prop.Enum)) + for _, v := range prop.Enum { + if s, ok := v.(string); ok { + b, _ := json.Marshal(s) + parts = append(parts, string(b)) + } else { + parts = append(parts, fmt.Sprint(v)) + } + } + return strings.Join(parts, " | ") + } + typ := "any" + if len(prop.Type) > 0 { + typ = prop.Type[0] + } + switch typ { + case "array": + return apertusArrayType(prop.Items) + case "integer", "number": + return "number" + case "boolean": + return "boolean" + case "string": + return "string" + case "object": + if prop.Properties != nil && prop.Properties.Len() > 0 { + return apertusObjectType(prop.Properties, prop.Required) + } + return "object" + default: + if len(prop.Type) > 1 { + parts := make([]string, 0, len(prop.Type)) + for _, t := range prop.Type { + parts = append(parts, apertusTypeScriptType(api.ToolProperty{Type: api.PropertyType{t}})) + } + return strings.Join(parts, " | ") + } + return "any" + } +} + +func apertusArrayType(items any) string { + if items == nil { + return "any[]" + } + var p api.ToolProperty + b, err := json.Marshal(items) + if err != nil || json.Unmarshal(b, &p) != nil { + return "any[]" + } + inner := apertusTypeScriptType(p) + if inner == "object | object" || len(inner) > 50 { + inner = "any" + } + return inner + "[]" +} + +func apertusObjectType(properties *api.ToolPropertiesMap, required []string) string { + set := make(map[string]bool, len(required)) + for _, n := range required { + set[n] = true + } + var sb strings.Builder + sb.WriteString("{\n") + i := 0 + for name, prop := range properties.All() { + sb.WriteString(name) + if !set[name] { + sb.WriteString("?") + } + sb.WriteString(": ") + sb.WriteString(apertusTypeScriptType(prop)) + i++ + if i < properties.Len() { + sb.WriteString(", ") + } + } + sb.WriteString("}") + return sb.String() +} + +func renderApertusToolCalls(sb *strings.Builder, calls []api.ToolCall) error { + sb.WriteString(apertusToolsPrefix) + sb.WriteString("[") + for i, call := range calls { + if !apertusIdentifier(call.Function.Name) { + return fmt.Errorf("invalid apertus tool name %q", call.Function.Name) + } + if i > 0 { + sb.WriteString(", ") + } + args, err := json.Marshal(call.Function.Arguments) + if err != nil { + return err + } + name, err := json.Marshal(call.Function.Name) + if err != nil { + return err + } + sb.WriteString("{") + sb.Write(name) + sb.WriteString(": ") + sb.Write(args) + sb.WriteString("}") + } + sb.WriteString("]") + sb.WriteString(apertusToolsSuffix) + return nil +} + +func validateApertusTools(tools []api.Tool) error { + names := make(map[string]struct{}, len(tools)) + for _, tool := range tools { + name := tool.Function.Name + if !apertusIdentifier(name) { + return fmt.Errorf("invalid apertus tool name %q", name) + } + if _, ok := names[name]; ok { + return fmt.Errorf("duplicate apertus tool name %q", name) + } + names[name] = struct{}{} + if err := validateApertusText(tool.Function.Description); err != nil { + return err + } + budget := maxApertusSchemaNodes + if err := validateApertusSchema(tool.Function.Parameters.Properties, tool.Function.Parameters.Required, 0, &budget); err != nil { + return fmt.Errorf("invalid apertus tool %q schema: %w", name, err) + } + } + return nil +} + +func validateApertusSchema(props *api.ToolPropertiesMap, required []string, depth int, budget *int) error { + if depth > maxApertusSchemaDepth { + return fmt.Errorf("schema exceeds depth %d", maxApertusSchemaDepth) + } + declared := make(map[string]struct{}) + if props != nil { + for name, prop := range props.All() { + if !apertusIdentifier(name) { + return fmt.Errorf("invalid property name %q", name) + } + declared[name] = struct{}{} + if err := validateApertusProperty(prop, depth+1, budget); err != nil { + return err + } + } + } + for _, name := range required { + if _, ok := declared[name]; !ok { + return fmt.Errorf("required property %q is not declared", name) + } + } + return nil +} + +func validateApertusProperty(prop api.ToolProperty, depth int, budget *int) error { + if depth > maxApertusSchemaDepth { + return fmt.Errorf("schema exceeds depth %d", maxApertusSchemaDepth) + } + if *budget == 0 { + return fmt.Errorf("schema exceeds %d nodes", maxApertusSchemaNodes) + } + *budget-- + if err := validateApertusText(prop.Description); err != nil { + return err + } + if err := validateApertusSchema(prop.Properties, prop.Required, depth+1, budget); err != nil { + return err + } + for _, nested := range prop.AnyOf { + if err := validateApertusProperty(nested, depth+1, budget); err != nil { + return err + } + } + if prop.Items != nil { + b, err := json.Marshal(prop.Items) + if err != nil { + return err + } + var nested api.ToolProperty + if err := json.Unmarshal(b, &nested); err == nil { + if err := validateApertusProperty(nested, depth+1, budget); err != nil { + return err + } + } + } + return nil +} + +func validateApertusText(s string) error { + for _, token := range []string{apertusSystemStart, apertusSystemEnd, apertusDeveloperStart, apertusDeveloperEnd, apertusUserStart, apertusUserEnd, apertusAssistantStart, apertusAssistantEnd, apertusToolsPrefix, apertusToolsSuffix, apertusImageToken} { + if strings.Contains(s, token) { + return fmt.Errorf("apertus content contains reserved token %q", token) + } + } + return nil +} + +func apertusIdentifier(s string) bool { + if s == "" { + return false + } + for i := range len(s) { + c := s[i] + if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_' || c == '$' || (i > 0 && c >= '0' && c <= '9')) { + return false + } + } + return true +} diff --git a/model/renderers/apertus_test.go b/model/renderers/apertus_test.go new file mode 100644 index 00000000000..b04bfac9481 --- /dev/null +++ b/model/renderers/apertus_test.go @@ -0,0 +1,138 @@ +package renderers + +import ( + "strings" + "testing" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/model/parsers" +) + +func TestApertusRendererGrammar(t *testing.T) { + tool := apertusRendererTool("get_weather") + got, err := (&ApertusRenderer{}).Render([]api.Message{{Role: "system", Content: "Be concise."}, {Role: "user", Content: "Hello"}}, []api.Tool{tool}, nil) + if err != nil { + t.Fatal(err) + } + for _, want := range []string{"<|system_start|>Be concise.<|system_end|>", "<|developer_start|>Deliberation: disabled\nTool Capabilities:\n", "type get_weather = (_: {\ncity: string\n}) => any;", "<|user_start|>Hello<|user_end|>"} { + if !strings.Contains(got, want) { + t.Fatalf("missing %q in %q", want, got) + } + } + if strings.HasSuffix(got, apertusAssistantStart) { + t.Fatal("tool decision must not append assistant generation prompt") + } +} + +func TestApertusRendererHistoryAndToolCalls(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("city", "Bern") + got, err := (&ApertusRenderer{}).Render([]api.Message{{Role: "user", Content: "weather"}, {Role: "assistant", ToolCalls: []api.ToolCall{{Function: api.ToolCallFunction{Name: "get_weather", Arguments: args}}}}, {Role: "tool", Content: `{"temperature":22}`}}, []api.Tool{apertusRendererTool("get_weather")}, nil) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(got, `<|tools_prefix|>[{"get_weather": {"city":"Bern"}}]<|tools_suffix|>[{"temperature":22}]<|assistant_end|><|assistant_start|>`) { + t.Fatalf("unexpected history %q", got) + } +} + +func TestApertusRendererToolHistoryStateAndIdentity(t *testing.T) { + r := &ApertusRenderer{} + args := api.NewToolCallFunctionArguments() + args.Set("city", "Bern") + call := api.ToolCall{Function: api.ToolCallFunction{Name: "get_weather", Arguments: args}} + if _, err := r.Render([]api.Message{{Role: "assistant", Content: "text"}, {Role: "tool", Content: `{}`}}, []api.Tool{apertusRendererTool("get_weather")}, nil); err == nil { + t.Fatal("tool after assistant text accepted") + } + if _, err := r.Render([]api.Message{{Role: "assistant", ToolCalls: []api.ToolCall{call}}}, nil, nil); err == nil { + t.Fatal("assistant call without declarations accepted") + } + undeclared := call + undeclared.Function.Name = "other" + if _, err := r.Render([]api.Message{{Role: "assistant", ToolCalls: []api.ToolCall{undeclared}}}, []api.Tool{apertusRendererTool("get_weather")}, nil); err == nil { + t.Fatal("undeclared assistant call accepted") + } + second := call + second.Function.Arguments = api.NewToolCallFunctionArguments() + second.Function.Arguments.Set("city", "Zurich") + got, err := r.Render([]api.Message{{Role: "assistant", ToolCalls: []api.ToolCall{call, second}}, {Role: "tool", Content: `{"weather":1}`}, {Role: "tool", Content: `{"weather":2}`}}, []api.Tool{apertusRendererTool("get_weather")}, nil) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(got, `[{"weather":1}, {"weather":2}]`) { + t.Fatalf("ordered multi-result grouping missing: %q", got) + } +} + +func TestApertusRendererRepeatedCallRoundTripAndPendingTransitions(t *testing.T) { + tool := apertusRendererTool("get_weather") + p := &parsers.ApertusParser{} + p.Init([]api.Tool{tool}, nil, nil) + _, _, calls, err := p.Add(`<|tools_prefix|>[{"get_weather":{"city":"Bern"}},{"get_weather":{"city":"Zurich"}}]<|tools_suffix|>`, true) + if err != nil || len(calls) != 2 || calls[0].Function.Index != 0 || calls[1].Function.Index != 1 { + t.Fatalf("calls=%#v err=%v", calls, err) + } + r := &ApertusRenderer{} + terminal := []api.Message{{Role: "assistant", ToolCalls: calls}} + if _, err := r.Render(terminal, []api.Tool{tool}, nil); err != nil { + t.Fatalf("terminal pending call rejected: %v", err) + } + partial := append(append([]api.Message{}, terminal...), api.Message{Role: "tool", Content: `{}`}) + for _, next := range []api.Message{{Role: "user", Content: "next"}, {Role: "system", Content: "next"}, {Role: "assistant", Content: "next"}} { + messages := append(append([]api.Message{}, partial...), next) + if _, err := r.Render(messages, []api.Tool{tool}, nil); err == nil { + t.Fatalf("partial tool results followed by %s accepted", next.Role) + } + } + if _, err := r.Render(append(partial, api.Message{Role: "tool", Content: `{}`}, api.Message{Role: "tool", Content: `{}`}), []api.Tool{tool}, nil); err == nil { + t.Fatal("extra tool result accepted") + } + if _, err := r.Render(append(partial, api.Message{Role: "tool", Content: `{}`}, api.Message{Role: "user", Content: "next"}), []api.Tool{tool}, nil); err != nil { + t.Fatalf("complete ordered results rejected: %v", err) + } +} + +func TestApertusRendererRejectsAmbiguousOrUnsafeSchema(t *testing.T) { + r := &ApertusRenderer{} + if _, err := r.Render(nil, []api.Tool{apertusRendererTool("same"), apertusRendererTool("same")}, nil); err == nil { + t.Fatal("duplicate tools accepted") + } + if _, err := r.Render(nil, []api.Tool{apertusRendererTool("bad.name")}, nil); err == nil { + t.Fatal("separator-bearing tool accepted") + } + missing := apertusRendererTool("missing") + missing.Function.Parameters.Required = []string{"not_declared"} + if _, err := r.Render(nil, []api.Tool{missing}, nil); err == nil { + t.Fatal("required/property mismatch accepted") + } + deep := api.ToolProperty{Type: api.PropertyType{"object"}} + for range maxApertusSchemaDepth + 2 { + props := api.NewToolPropertiesMap() + props.Set("next", deep) + deep = api.ToolProperty{Type: api.PropertyType{"object"}, Properties: props} + } + tool := apertusRendererTool("deep") + tool.Function.Parameters.Properties.Set("root", deep) + if _, err := r.Render(nil, []api.Tool{tool}, nil); err == nil { + t.Fatal("deep schema accepted") + } + wide := apertusRendererTool("wide") + for i := range maxApertusSchemaNodes { + wide.Function.Parameters.Properties.Set("field"+strings.Repeat("x", i+1), api.ToolProperty{Type: api.PropertyType{"string"}}) + } + if _, err := r.Render(nil, []api.Tool{wide}, nil); err == nil { + t.Fatal("oversized schema accepted") + } + if _, err := r.Render([]api.Message{{Role: "user", Content: "x" + apertusAssistantStart}}, nil, nil); err == nil { + t.Fatal("control token accepted") + } + if _, err := r.Render([]api.Message{{Role: "tool", Content: "orphan"}}, nil, nil); err == nil { + t.Fatal("orphan tool accepted") + } +} + +func apertusRendererTool(name string) api.Tool { + props := api.NewToolPropertiesMap() + props.Set("city", api.ToolProperty{Type: api.PropertyType{"string"}}) + return api.Tool{Type: "function", Function: api.ToolFunction{Name: name, Parameters: api.ToolFunctionParameters{Type: "object", Required: []string{"city"}, Properties: props}}} +} diff --git a/model/renderers/renderer.go b/model/renderers/renderer.go index 93d3b0eabfb..c685b238857 100644 --- a/model/renderers/renderer.go +++ b/model/renderers/renderer.go @@ -57,6 +57,8 @@ func rendererForName(name string) Renderer { return constructor() } switch name { + case "apertus": + return &ApertusRenderer{} case "qwen3-coder": renderer := &Qwen3CoderRenderer{} return renderer From d260be754115ebec61a9a87aa8491b334c17c0d9 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 20:14:23 +0000 Subject: [PATCH 33/58] create: infer Apertus metadata Co-authored-by: Codex --- x/create/client/create.go | 13 ++++++ x/create/client/create_test.go | 72 ++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/x/create/client/create.go b/x/create/client/create.go index 64400b9c561..b0ab78a34dd 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -627,6 +627,15 @@ func isGPTOSSFamily(s string) bool { return strings.Contains(s, "gptoss") || strings.Contains(s, "gpt_oss") || strings.Contains(s, "gpt-oss") } +func isApertusFamily(s string) bool { + switch strings.ToLower(s) { + case "apertus", "apertusforcausallm": + return true + default: + return false + } +} + func qwen35RendererName(modelDir string) string { template := readChatTemplate(modelDir) if strings.Contains(template, "resolved_reasoning_effort") && @@ -707,6 +716,8 @@ func getParserName(modelDir string) string { func parserNameForIdentifier(modelDir, s string) string { s = strings.ToLower(s) switch { + case isApertusFamily(s): + return "apertus" case strings.HasPrefix(s, "museglimmer") || s == "muse_glimmer": return "glimmer" case strings.Contains(s, "laguna"): @@ -774,6 +785,8 @@ func getRendererName(modelDir string) string { func rendererNameForIdentifier(modelDir, s string) string { s = strings.ToLower(s) switch { + case isApertusFamily(s): + return "apertus" case strings.HasPrefix(s, "museglimmer") || s == "muse_glimmer": return "glimmer" case strings.Contains(s, "laguna"): diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index c3f07297439..83a3a61548b 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -473,6 +473,78 @@ func TestInferSafetensorsCapabilities(t *testing.T) { } } +func TestApertusMetadataInference(t *testing.T) { + tests := []struct { + name string + configJSON string + parser string + renderer string + caps []string + }{ + { + name: "architecture wins over top-level and nested model type", + configJSON: `{"architectures":["ApertusForCausalLM"],"model_type":"qwen3","llm_config":{"model_type":"gpt_oss"}}`, + parser: "apertus", + renderer: "apertus", + caps: []string{"completion", "tools"}, + }, + { + name: "top-level wins over nested model type", + configJSON: `{"model_type":"apertus","llm_config":{"model_type":"qwen3"}}`, + parser: "apertus", + renderer: "apertus", + caps: []string{"completion", "tools"}, + }, + { + name: "nested model type", + configJSON: `{"model_type":"wrapper","llm_config":{"model_type":"apertus"}}`, + parser: "apertus", + renderer: "apertus", + caps: []string{"completion", "tools"}, + }, + { + name: "near matches are not Apertus", + configJSON: `{"architectures":["ApertusForCausalLMExtra"],"model_type":"apertus-next","llm_config":{"model_type":"not_apertus"}}`, + caps: []string{"completion"}, + }, + { + name: "missing identifiers", + configJSON: `{}`, + caps: []string{"completion"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644); err != nil { + t.Fatal(err) + } + if got := getParserName(dir); got != tt.parser { + t.Fatalf("parser = %q, want %q", got, tt.parser) + } + if got := getRendererName(dir); got != tt.renderer { + t.Fatalf("renderer = %q, want %q", got, tt.renderer) + } + if got := inferSafetensorsCapabilities(dir, getParserName(dir)); !slices.Equal(got, tt.caps) { + t.Fatalf("capabilities = %v, want %v", got, tt.caps) + } + }) + } + + for _, invalid := range []string{"not json", `{"architectures":"ApertusForCausalLM"}`} { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(invalid), 0o644); err != nil { + t.Fatal(err) + } + if got := getParserName(dir); got != "" { + t.Fatalf("invalid config parser = %q, want empty", got) + } + if got := getRendererName(dir); got != "" { + t.Fatalf("invalid config renderer = %q, want empty", got) + } + } +} + func TestCreateModelfileLayersIncludesParameters(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) From 4c6296a855bc0f631951a75b338ceb89b4966a50 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 20:20:12 +0000 Subject: [PATCH 34/58] model: enable Apertus thinking mode Co-authored-by: Codex --- model/parsers/apertus.go | 73 ++++++++++++++++++++++++++++----- model/parsers/apertus_test.go | 43 +++++++++++++++++++ model/renderers/apertus.go | 22 +++++++--- model/renderers/apertus_test.go | 20 +++++++++ x/create/client/create_test.go | 6 +-- 5 files changed, 146 insertions(+), 18 deletions(-) diff --git a/model/parsers/apertus.go b/model/parsers/apertus.go index c7396f577ba..2415dd6b82a 100644 --- a/model/parsers/apertus.go +++ b/model/parsers/apertus.go @@ -16,6 +16,8 @@ const ( apertusToolCloseTag = "<|tools_suffix|>" apertusAssistantOpenTag = "<|assistant_start|>" apertusAssistantCloseTag = "<|assistant_end|>" + apertusInnerOpenTag = "<|inner_prefix|>" + apertusInnerCloseTag = "<|inner_suffix|>" // A tool call is one model response fragment. Keep malformed streams from // retaining unbounded output while allowing substantially larger calls than // the model's normal tool grammar produces. @@ -26,6 +28,7 @@ type apertusParserState uint8 const ( apertusContent apertusParserState = iota + apertusThinking apertusToolCalls ) @@ -36,6 +39,7 @@ type ApertusParser struct { initErr error callIndex int pendingBare bool + thinking bool } func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool { @@ -45,6 +49,7 @@ func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkVa p.initErr = nil p.callIndex = 0 p.pendingBare = false + p.thinking = thinkValue != nil && thinkValue.Bool() for _, tool := range tools { name := tool.Function.Name if !apertusIdentifier(name) { @@ -72,12 +77,21 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) } - var out strings.Builder + var out, thought strings.Builder for { current := p.acc.String() switch p.state { case apertusContent: - if idx := strings.Index(current, apertusToolOpenTag); idx >= 0 { + innerIdx := strings.Index(current, apertusInnerOpenTag) + toolIdx := strings.Index(current, apertusToolOpenTag) + if innerIdx >= 0 && (toolIdx < 0 || innerIdx < toolIdx) { + out.WriteString(cleanApertusContent(current[:innerIdx])) + p.acc.Reset() + p.acc.WriteString(current[innerIdx+len(apertusInnerOpenTag):]) + p.state = apertusThinking + continue + } + if idx := toolIdx; idx >= 0 { p.pendingBare = false out.WriteString(cleanApertusContent(current[:idx])) p.acc.Reset() @@ -94,7 +108,7 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call if parsed, parseErr := p.parseToolCalls(cleaned); parseErr == nil { p.acc.Reset() p.pendingBare = false - return out.String(), "", parsed, nil + return out.String(), thought.String(), parsed, nil } else if !isSoftApertusToolParseError(parseErr) { return "", "", nil, parseErr } @@ -102,20 +116,21 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call p.acc.Reset() p.pendingBare = false out.WriteString(cleaned) - return out.String(), "", calls, nil + return out.String(), thought.String(), calls, nil } if p.looksLikeToolCallStart(current) { p.pendingBare = true if p.acc.Len() > maxApertusToolCallBytes { return "", "", nil, fmt.Errorf("apertus tool call exceeds %d bytes", maxApertusToolCallBytes) } - return out.String(), "", nil, nil + return out.String(), thought.String(), nil, nil } p.pendingBare = false overlapLen := max( overlap(current, apertusToolOpenTag), overlap(current, apertusAssistantOpenTag), overlap(current, apertusAssistantCloseTag), + overlap(current, apertusInnerOpenTag), ) n := len(current) - overlapLen if n == len(current) { @@ -130,7 +145,43 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call p.acc.Reset() p.acc.WriteString(current[n:]) } - return out.String(), "", calls, nil + return out.String(), thought.String(), calls, nil + + case apertusThinking: + if idx := strings.Index(current, apertusInnerCloseTag); idx >= 0 { + inner := current[:idx] + if p.thinking { + thought.WriteString(inner) + } else { + out.WriteString(cleanApertusContent(inner)) + } + p.acc.Reset() + p.acc.WriteString(strings.TrimLeftFunc(current[idx+len(apertusInnerCloseTag):], unicode.IsSpace)) + p.state = apertusContent + continue + } + if done { + if p.thinking { + thought.WriteString(current) + } else { + out.WriteString(cleanApertusContent(current)) + } + p.acc.Reset() + p.state = apertusContent + return out.String(), thought.String(), calls, nil + } + n := len(current) - overlap(current, apertusInnerCloseTag) + if n > 0 { + emit := current[:n] + if p.thinking { + thought.WriteString(emit) + } else { + out.WriteString(cleanApertusContent(emit)) + } + p.acc.Reset() + p.acc.WriteString(current[n:]) + } + return out.String(), thought.String(), calls, nil case apertusToolCalls: if idx := strings.Index(current, apertusToolCloseTag); idx >= 0 { @@ -162,9 +213,9 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call p.acc.Reset() p.state = apertusContent p.pendingBare = false - return out.String(), "", calls, nil + return out.String(), thought.String(), calls, nil } - return out.String(), "", calls, nil + return out.String(), thought.String(), calls, nil } } } @@ -206,12 +257,14 @@ func apertusPendingBarePrefix(s string) bool { func cleanApertusContent(s string) string { s = strings.ReplaceAll(s, apertusAssistantOpenTag, "") s = strings.ReplaceAll(s, apertusAssistantCloseTag, "") + s = strings.ReplaceAll(s, apertusInnerOpenTag, "") + s = strings.ReplaceAll(s, apertusInnerCloseTag, "") return strings.TrimRightFunc(s, unicode.IsSpace) } func (p *ApertusParser) HasToolSupport() bool { return true } -func (p *ApertusParser) HasThinkingSupport() bool { return false } +func (p *ApertusParser) HasThinkingSupport() bool { return true } func (p *ApertusParser) PreservedTokens() []string { - return []string{apertusToolOpenTag, apertusToolCloseTag, apertusAssistantOpenTag, apertusAssistantCloseTag} + return []string{apertusToolOpenTag, apertusToolCloseTag, apertusAssistantOpenTag, apertusAssistantCloseTag, apertusInnerOpenTag, apertusInnerCloseTag} } func (p *ApertusParser) parseToolCalls(raw string) ([]api.ToolCall, error) { diff --git a/model/parsers/apertus_test.go b/model/parsers/apertus_test.go index a2fa77ef4d3..8ec0e668b84 100644 --- a/model/parsers/apertus_test.go +++ b/model/parsers/apertus_test.go @@ -200,6 +200,49 @@ func TestApertusParserBareAndTaggedHardErrorsMatch(t *testing.T) { } } +func TestApertusParserThinkingAndTools(t *testing.T) { + for _, think := range []*api.ThinkValue{nil, {Value: false}, {Value: "low"}, {Value: true}} { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("one")}, nil, think) + content, thinking, calls, err := p.Add(`<|inner_prefix|>reason<|inner_suffix|><|tools_prefix|>[{"one":{}}]<|tools_suffix|>`, true) + if err != nil || len(calls) != 1 { + t.Fatalf("think=%v calls=%d err=%v", think, len(calls), err) + } + if think != nil && think.Bool() { + if thinking != "reason" || content != "" { + t.Fatalf("thinking=%q content=%q", thinking, content) + } + } else if content != "reason" || thinking != "" { + t.Fatalf("content=%q thinking=%q", content, thinking) + } + } +} + +func TestApertusParserSplitThinkingTags(t *testing.T) { + p := &ApertusParser{} + p.Init(nil, nil, &api.ThinkValue{Value: true}) + + var content, thinking string + for _, chunk := range []struct { + text string + done bool + }{ + {"<|inner_pre", false}, + {"fix|>reason<|inner_suf", false}, + {"fix|>answer", true}, + } { + gotContent, gotThinking, _, err := p.Add(chunk.text, chunk.done) + if err != nil { + t.Fatal(err) + } + content += gotContent + thinking += gotThinking + } + if thinking != "reason" || content != "answer" { + t.Fatalf("thinking=%q content=%q", thinking, content) + } +} + func apertusParserTool(name string) api.Tool { return api.Tool{Type: "function", Function: api.ToolFunction{Name: name}} } diff --git a/model/renderers/apertus.go b/model/renderers/apertus.go index 86e377bd277..b88f2ac702c 100644 --- a/model/renderers/apertus.go +++ b/model/renderers/apertus.go @@ -21,6 +21,8 @@ const ( apertusToolsPrefix = "<|tools_prefix|>" apertusToolsSuffix = "<|tools_suffix|>" apertusImageToken = "<|image|>" + apertusInnerOpenTag = "<|inner_prefix|>" + apertusInnerCloseTag = "<|inner_suffix|>" maxApertusSchemaDepth = 32 maxApertusSchemaNodes = 4096 ) @@ -30,9 +32,7 @@ type ApertusRenderer struct{} func (r *ApertusRenderer) LeadingBOS() string { return "" } func (r *ApertusRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) { - if think != nil && think.Bool() { - return "", fmt.Errorf("apertus renderer does not support thinking") - } + thinkingEnabled := think != nil && think.Bool() if err := validateApertusTools(tools); err != nil { return "", err } @@ -44,6 +44,9 @@ func (r *ApertusRenderer) Render(messages []api.Message, tools []api.Tool, think if err := validateApertusText(message.Content); err != nil { return "", err } + if err := validateApertusText(message.Thinking); err != nil { + return "", err + } } var sb strings.Builder start := 0 @@ -59,7 +62,11 @@ func (r *ApertusRenderer) Render(messages []api.Message, tools []api.Tool, think sb.WriteString(apertusSystemEnd) } sb.WriteString(apertusDeveloperStart) - sb.WriteString("Deliberation: disabled\n") + if thinkingEnabled { + sb.WriteString("Deliberation: enabled\n") + } else { + sb.WriteString("Deliberation: disabled\n") + } if len(tools) == 0 { sb.WriteString("Tool Capabilities: disabled") } else { @@ -112,6 +119,11 @@ func (r *ApertusRenderer) Render(messages []api.Message, tools []api.Tool, think sb.WriteString("]") inTool = false } + if thinkingEnabled && message.Thinking != "" { + sb.WriteString(apertusInnerOpenTag) + sb.WriteString(message.Thinking) + sb.WriteString(apertusInnerCloseTag) + } sb.WriteString(message.Content) if len(message.ToolCalls) > 0 { if err := validateApertusHistoricalCalls(message.ToolCalls, declaredTools); err != nil { @@ -418,7 +430,7 @@ func validateApertusProperty(prop api.ToolProperty, depth int, budget *int) erro } func validateApertusText(s string) error { - for _, token := range []string{apertusSystemStart, apertusSystemEnd, apertusDeveloperStart, apertusDeveloperEnd, apertusUserStart, apertusUserEnd, apertusAssistantStart, apertusAssistantEnd, apertusToolsPrefix, apertusToolsSuffix, apertusImageToken} { + for _, token := range []string{apertusSystemStart, apertusSystemEnd, apertusDeveloperStart, apertusDeveloperEnd, apertusUserStart, apertusUserEnd, apertusAssistantStart, apertusAssistantEnd, apertusToolsPrefix, apertusToolsSuffix, apertusImageToken, apertusInnerOpenTag, apertusInnerCloseTag} { if strings.Contains(s, token) { return fmt.Errorf("apertus content contains reserved token %q", token) } diff --git a/model/renderers/apertus_test.go b/model/renderers/apertus_test.go index b04bfac9481..9858dcf36ac 100644 --- a/model/renderers/apertus_test.go +++ b/model/renderers/apertus_test.go @@ -24,6 +24,26 @@ func TestApertusRendererGrammar(t *testing.T) { } } +func TestApertusRendererThinkingModes(t *testing.T) { + message := []api.Message{{Role: "assistant", Thinking: "reason", Content: "answer"}} + for _, think := range []*api.ThinkValue{nil, {Value: false}, {Value: "low"}, {Value: true}} { + got, err := (&ApertusRenderer{}).Render(message, nil, think) + if err != nil { + t.Fatal(err) + } + enabled := think != nil && think.Bool() + if strings.Contains(got, apertusInnerOpenTag) != enabled || strings.Contains(got, "Deliberation: enabled") != enabled { + t.Fatalf("think=%v got=%q", think, got) + } + if !enabled && strings.Contains(got, "reason") { + t.Fatalf("disabled thinking leaked: %q", got) + } + } + if _, err := (&ApertusRenderer{}).Render([]api.Message{{Role: "assistant", Thinking: apertusAssistantStart}}, nil, &api.ThinkValue{Value: true}); err == nil { + t.Fatal("thinking control token accepted") + } +} + func TestApertusRendererHistoryAndToolCalls(t *testing.T) { args := api.NewToolCallFunctionArguments() args.Set("city", "Bern") diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 83a3a61548b..638943d689c 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -486,21 +486,21 @@ func TestApertusMetadataInference(t *testing.T) { configJSON: `{"architectures":["ApertusForCausalLM"],"model_type":"qwen3","llm_config":{"model_type":"gpt_oss"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools"}, + caps: []string{"completion", "tools", "thinking"}, }, { name: "top-level wins over nested model type", configJSON: `{"model_type":"apertus","llm_config":{"model_type":"qwen3"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools"}, + caps: []string{"completion", "tools", "thinking"}, }, { name: "nested model type", configJSON: `{"model_type":"wrapper","llm_config":{"model_type":"apertus"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools"}, + caps: []string{"completion", "tools", "thinking"}, }, { name: "near matches are not Apertus", From 55c4341251a64411452e6bf5c1e6b45d069a4b90 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 20:26:33 +0000 Subject: [PATCH 35/58] create: keep Apertus NVFP4 imports uniform Co-authored-by: Codex --- x/create/apertus.go | 55 ++++++++++++++++++++++++++++++ x/create/apertus_test.go | 72 ++++++++++++++++++++++++++++++++++++++++ x/create/create.go | 1 + 3 files changed, 128 insertions(+) create mode 100644 x/create/apertus.go create mode 100644 x/create/apertus_test.go diff --git a/x/create/apertus.go b/x/create/apertus.go new file mode 100644 index 00000000000..13b7f1935a9 --- /dev/null +++ b/x/create/apertus.go @@ -0,0 +1,55 @@ +package create + +import ( + "encoding/json" + "strings" +) + +type apertusImportTransform struct{} + +func newApertusImportTransform(json.RawMessage) (quantizePolicy, error) { + return apertusImportTransform{}, nil +} + +func (apertusImportTransform) quantizationType(name string, shape []int32, quantize string) string { + // This policy is deliberately limited to Apertus NVFP4 imports. Unknown or + // mismatched requests remain at source precision rather than falling through + // to a different shared quantization recipe. + if normalizeQuantType(quantize) != "nvfp4" { + return "" + } + + stackedExpert := isStackedExpertWeight(name) + if !stackedExpert && !ShouldQuantize(name, "") { + return "" + } + if len(shape) != 2 && !(len(shape) == 3 && stackedExpert) { + return "" + } + + elems, ok := apertusElementCount(shape) + if !ok || elems < 1024 || isRoutingGate(name) || isApertusNonlinearTensor(name) || !isAligned(shape, "nvfp4") { + return "" + } + + return "nvfp4" +} + +func apertusElementCount(shape []int32) (uint64, bool) { + elements := uint64(1) + for _, dim := range shape { + if dim <= 0 { + return 0, false + } + d := uint64(dim) + if elements > ^uint64(0)/d { + return 0, false + } + elements *= d + } + return elements, true +} + +func isApertusNonlinearTensor(name string) bool { + return strings.Contains(name, ".act_fn.") +} diff --git a/x/create/apertus_test.go b/x/create/apertus_test.go new file mode 100644 index 00000000000..c3c4db58acd --- /dev/null +++ b/x/create/apertus_test.go @@ -0,0 +1,72 @@ +package create + +import ( + "math" + "testing" +) + +func TestApertusImportTransform(t *testing.T) { + transform, err := newTensorImportTransform(Inventory{ + Config: sourceModelConfig{Architectures: []string{"ApertusForCausalLM"}}, + }) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + if _, ok := transform.(apertusImportTransform); !ok { + t.Fatalf("newTensorImportTransform() = %T, want apertusImportTransform", transform) + } + for _, architecture := range []string{"GptOssForCausalLM", "Qwen3_5ForCausalLM"} { + if _, ok := tensorImportTransformRegistry[architecture]; !ok { + t.Fatalf("registry lost %s", architecture) + } + } +} + +func TestApertusQuantizationPolicy(t *testing.T) { + transform := apertusImportTransform{} + for _, tt := range []struct { + name, tensor, quantize, want string + shape []int32 + }{ + {"lm head retains requested nvfp4", "lm_head.weight", "NVFP4", "nvfp4", []int32{131072, 4096}}, + {"k projection retains requested nvfp4", "model.layers.0.self_attn.k_proj.weight", "nvfp4", "nvfp4", []int32{1024, 4096}}, + {"v projection retains requested nvfp4", "model.layers.0.self_attn.v_proj.weight", "nvfp4", "nvfp4", []int32{1024, 4096}}, + {"down projection retains requested nvfp4", "model.layers.0.mlp.down_proj.weight", "nvfp4", "nvfp4", []int32{4096, 21504}}, + {"stacked rank three projection retains requested nvfp4", "model.layers.0.mlp.experts.down_proj.weight", "nvfp4", "nvfp4", []int32{8, 4096, 21504}}, + {"embeddings remain source precision", "model.embed_tokens.weight", "nvfp4", "", []int32{131072, 4096}}, + {"routing remains source precision", "model.layers.0.mlp.gate.weight", "nvfp4", "", []int32{4096, 4096}}, + {"small remains source precision", "model.layers.0.mlp.down_proj.weight", "nvfp4", "", []int32{16, 16}}, + {"nonlinear remains source precision", "model.layers.0.mlp.act_fn.alpha.weight", "nvfp4", "", []int32{1024, 4096}}, + {"misaligned remains source precision", "model.layers.0.mlp.down_proj.weight", "nvfp4", "", []int32{1024, 4095}}, + {"unknown quantization remains source precision", "model.layers.0.mlp.down_proj.weight", "unknown", "", []int32{1024, 4096}}, + {"mismatched rank remains source precision", "model.layers.0.mlp.down_proj.weight", "nvfp4", "", []int32{2, 1024, 4096}}, + } { + t.Run(tt.name, func(t *testing.T) { + if got := transform.quantizationType(tt.tensor, tt.shape, tt.quantize); got != tt.want { + t.Fatalf("quantizationType(%q, %v, %q) = %q, want %q", tt.tensor, tt.shape, tt.quantize, got, tt.want) + } + }) + } +} + +func TestApertusElementCount(t *testing.T) { + for _, tt := range []struct { + name string + shape []int32 + want uint64 + ok bool + }{ + {"positive rank two", []int32{1024, 4096}, 4194304, true}, + {"greatest accepted max-dimension product", []int32{math.MaxInt32, math.MaxInt32, 4}, 18446744056529682436, true}, + {"first rejected max-dimension product", []int32{math.MaxInt32, math.MaxInt32, 5}, 0, false}, + {"zero dimension", []int32{1024, 0}, 0, false}, + {"negative dimension", []int32{1024, -1}, 0, false}, + } { + t.Run(tt.name, func(t *testing.T) { + got, ok := apertusElementCount(tt.shape) + if got != tt.want || ok != tt.ok { + t.Fatalf("apertusElementCount(%v) = (%d, %t), want (%d, %t)", tt.shape, got, ok, tt.want, tt.ok) + } + }) + } +} diff --git a/x/create/create.go b/x/create/create.go index df23786e05b..4ba0fdff3bb 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -493,6 +493,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "Gemma4ForConditionalGeneration": newGemma4ImportTransform, "Gemma4UnifiedForCausalLM": newGemma4ImportTransform, "Gemma4UnifiedForConditionalGeneration": newGemma4ImportTransform, + "ApertusForCausalLM": newApertusImportTransform, "gemma4_unified": newGemma4ImportTransform, "gemma4_unified_text": newGemma4ImportTransform, "LagunaForCausalLM": newLagunaImportTransform, From 7dcd6a31aaf612bc31134229f408097f3c4b8fa6 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 18:56:23 +0000 Subject: [PATCH 36/58] create: restore Apertus MXFP8 imports Co-authored-by: Codex --- x/create/apertus.go | 12 ++++++++---- x/create/apertus_test.go | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/x/create/apertus.go b/x/create/apertus.go index 13b7f1935a9..205944b22b7 100644 --- a/x/create/apertus.go +++ b/x/create/apertus.go @@ -12,10 +12,14 @@ func newApertusImportTransform(json.RawMessage) (quantizePolicy, error) { } func (apertusImportTransform) quantizationType(name string, shape []int32, quantize string) string { - // This policy is deliberately limited to Apertus NVFP4 imports. Unknown or - // mismatched requests remain at source precision rather than falling through - // to a different shared quantization recipe. - if normalizeQuantType(quantize) != "nvfp4" { + base := normalizeQuantType(quantize) + if base == "mxfp8" { + // Apertus has a bespoke NVFP4 policy, but its MXFP8 imports use the + // established general policy. Returning source precision here turns a + // 70B MXFP8 request into a BF16-sized artifact that cannot be admitted. + return GetTensorQuantization(name, shape, quantize) + } + if base != "nvfp4" { return "" } diff --git a/x/create/apertus_test.go b/x/create/apertus_test.go index c3c4db58acd..7afa874f8e2 100644 --- a/x/create/apertus_test.go +++ b/x/create/apertus_test.go @@ -49,6 +49,26 @@ func TestApertusQuantizationPolicy(t *testing.T) { } } +func TestApertusMXFP8UsesGeneralQuantizationPolicy(t *testing.T) { + transform := apertusImportTransform{} + for _, tt := range []struct { + name, tensor, want string + shape []int32 + }{ + {"70B lm head", "lm_head.weight", "mxfp8", []int32{131072, 8192}}, + {"70B attention projection", "model.layers.0.self_attn.q_proj.weight", "mxfp8", []int32{8192, 8192}}, + {"70B MLP projection", "model.layers.0.mlp.down_proj.weight", "mxfp8", []int32{8192, 28672}}, + {"embedding remains source precision", "model.embed_tokens.weight", "", []int32{131072, 8192}}, + {"routing remains source precision", "model.layers.0.mlp.gate.weight", "", []int32{8192, 8192}}, + } { + t.Run(tt.name, func(t *testing.T) { + if got := transform.quantizationType(tt.tensor, tt.shape, "mxfp8"); got != tt.want { + t.Fatalf("quantizationType(%q, %v, mxfp8) = %q, want %q", tt.tensor, tt.shape, got, tt.want) + } + }) + } +} + func TestApertusElementCount(t *testing.T) { for _, tt := range []struct { name string From 2bfda88b47681b47ae66de93d94ab7134e4705e8 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 19:14:00 +0000 Subject: [PATCH 37/58] apertus: accept packed MXFP8 import tensors Co-authored-by: Codex --- x/models/apertus/apertus.go | 10 +++++++--- x/models/apertus/apertus_test.go | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/x/models/apertus/apertus.go b/x/models/apertus/apertus.go index d9402ce87a2..8b9ad7dbe8f 100644 --- a/x/models/apertus/apertus.go +++ b/x/models/apertus/apertus.go @@ -582,11 +582,15 @@ func validateMatrix(tensors map[string]*mlx.Array, path string, out, input int, if bits != 8 { return fmt.Errorf("tensor %q has invalid mxfp8 bit width %d", name, bits) } - if err := validateShape(name, weight, []int{out, input}); err != nil { + // MLX stores four 8-bit MXFP8 values in each U32 word. + if input%4 != 0 { + return fmt.Errorf("tensor %q has mxfp8 input width %d, which is not divisible by 4", name, input) + } + if err := validateShape(name, weight, []int{out, input / 4}); err != nil { return err } - if weight.DType() != mlx.DTypeUint8 { - return fmt.Errorf("tensor %q dtype %s, want U8 mxfp8", name, weight.DType()) + if weight.DType() != mlx.DTypeUint32 { + return fmt.Errorf("tensor %q dtype %s, want U32 packed mxfp8", name, weight.DType()) } if err := validateShape(name+"_scale", scales, scaleShape); err != nil { return err diff --git a/x/models/apertus/apertus_test.go b/x/models/apertus/apertus_test.go index 7085edd45bc..9c92abccfc6 100644 --- a/x/models/apertus/apertus_test.go +++ b/x/models/apertus/apertus_test.go @@ -460,7 +460,7 @@ func quantizedMatrix(name, quantType string, out, input int) (map[string]*mlx.Ar tensors[name] = mlx.Zeros(mlx.DTypeUint32, out, input/8) tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, out, input/groupSize) case "mxfp8": - tensors[name] = mlx.Zeros(mlx.DTypeUint8, out, input) + tensors[name] = mlx.Zeros(mlx.DTypeUint32, out, input/4) tensors[name+"_scale"] = mlx.Zeros(mlx.DTypeUint8, out, input/groupSize) } return tensors, &Config{QuantType: quantType, QuantGroupSize: groupSize, QuantBits: bits, QuantMode: mode} From d4ad2d96531795500d4a8fe52a57570e4e95bd3a Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 15:36:54 +0000 Subject: [PATCH 38/58] create: preserve Apertus family metadata Co-authored-by: Codex --- x/create/client/create.go | 3 +++ x/create/client/create_test.go | 43 ++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/x/create/client/create.go b/x/create/client/create.go index b0ab78a34dd..1056f17342e 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -469,6 +469,9 @@ func inferModelFamily(modelDir string) string { if isGPTOSSFamily(identifier) { return "gptoss" } + if isApertusFamily(identifier) { + return "apertus" + } } return "" diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 638943d689c..84cef9c7f35 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -669,6 +669,49 @@ func TestNewManifestWriter_PopulatesGPTOSSFamily(t *testing.T) { } } +func TestNewManifestWriter_PopulatesApertusFamily(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + modelDir := t.TempDir() + if err := os.WriteFile(filepath.Join(modelDir, "config.json"), []byte(`{ + "architectures": ["ApertusForCausalLM"], + "model_type": "apertus" + }`), 0o644); err != nil { + t.Fatal(err) + } + + opts := CreateOptions{ModelName: "apertus-family-test", ModelDir: modelDir} + writer := newManifestWriter(opts, []string{"completion", "tools", "thinking"}, "apertus", "apertus") + if err := writer(opts.ModelName, create.LayerInfo{}, nil, create.Classification{Kind: create.SourceFloat, Quantize: "nvfp4"}); err != nil { + t.Fatalf("newManifestWriter() error = %v", err) + } + + name := model.ParseName(opts.ModelName) + mf, err := manifest.ParseNamedManifest(name) + if err != nil { + t.Fatal(err) + } + configPath, err := manifest.BlobsPath(mf.Config.Digest) + if err != nil { + t.Fatal(err) + } + data, err := os.ReadFile(configPath) + if err != nil { + t.Fatal(err) + } + + var cfg model.ConfigV2 + if err := json.Unmarshal(data, &cfg); err != nil { + t.Fatal(err) + } + if cfg.ModelFamily != "apertus" { + t.Fatalf("ModelFamily = %q, want apertus", cfg.ModelFamily) + } + if !slices.Equal(cfg.ModelFamilies, []string{"apertus"}) { + t.Fatalf("ModelFamilies = %v, want [apertus]", cfg.ModelFamilies) + } +} + func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) From 0d61ba2ed8994b598ebe48bc0cdcba165df66c37 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 23 Aug 2026 16:26:47 +0000 Subject: [PATCH 39/58] create: correct Apertus 1.0 capabilities Co-authored-by: Codex --- server/images.go | 9 ++++++++- server/images_test.go | 13 +++++++++++++ server/model_list_cache.go | 2 +- x/create/client/create.go | 34 +++++++++++++++++++++++++++++++++- x/create/client/create_test.go | 6 +++--- 5 files changed, 58 insertions(+), 6 deletions(-) diff --git a/server/images.go b/server/images.go index 9515e35580b..e4b81d564c6 100644 --- a/server/images.go +++ b/server/images.go @@ -441,13 +441,20 @@ func (m *Model) parserCapabilities(capabilities []model.Capability) []model.Capa if builtinParser.HasToolSupport() { capabilities = appendCapability(capabilities, model.CapabilityTools) } - if builtinParser.HasThinkingSupport() { + if builtinParser.HasThinkingSupport() && !isApertus1p0SafetensorsConfig(m.Config) { capabilities = appendCapability(capabilities, model.CapabilityThinking) } return capabilities } +func isApertus1p0SafetensorsConfig(cfg model.ConfigV2) bool { + return cfg.ModelFormat == "safetensors" && + cfg.ModelFamily == "apertus" && + cfg.Parser == "apertus" && + cfg.Renderer == "apertus" +} + func (m *Model) modelFamilyCapabilities(capabilities []model.Capability) []model.Capability { isGptoss := slices.Contains([]string{"gptoss", "gpt-oss"}, m.Config.ModelFamily) if isGptoss { diff --git a/server/images_test.go b/server/images_test.go index a3de7fc0fc1..677a59b9784 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -361,6 +361,19 @@ func TestModelCapabilities(t *testing.T) { model Model expectedCaps []model.Capability }{ + { + name: "Apertus 1.0 parser does not imply thinking", + model: Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", + ModelFamily: "apertus", + Parser: "apertus", + Renderer: "apertus", + }, + Template: chatTemplate, + }, + expectedCaps: []model.Capability{model.CapabilityTools}, + }, { name: "model with image generation capability via config", model: Model{ diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 9443a157668..8b3a3346112 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -382,7 +382,7 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum hasTags := openingTag != "" && closingTag != "" isGptoss := slices.Contains([]string{"gptoss", "gpt-oss"}, cfg.ModelFamily) if !slices.Contains(summary.Capabilities, model.CapabilityThinking) && - (hasTags || isGptoss || (builtinParser != nil && builtinParser.HasThinkingSupport())) { + (hasTags || isGptoss || (builtinParser != nil && builtinParser.HasThinkingSupport() && !isApertus1p0SafetensorsConfig(cfg))) { summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityThinking) } } diff --git a/x/create/client/create.go b/x/create/client/create.go index 1056f17342e..e070e50b5d5 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -352,7 +352,7 @@ func inferSafetensorsCapabilities(modelDir, parserName string) []string { capabilities = append(capabilities, "tools") } - if caps.thinking || (builtinParser != nil && builtinParser.HasThinkingSupport()) { + if caps.thinking || (builtinParser != nil && builtinParser.HasThinkingSupport() && !isApertus1p0ModelDir(modelDir, parserName)) { capabilities = append(capabilities, "thinking") } @@ -568,6 +568,38 @@ func detectCapabilities(modelDir string) modelCapabilities { } } +// isApertus1p0ModelDir recognizes the original text-only Apertus release. +// Parser capability alone is insufficient because the legacy grammar can +// parse inner spans for history compatibility even though the 1.0 checkpoint +// itself was not trained to produce thinking output. +func isApertus1p0ModelDir(modelDir, parserName string) bool { + if parserName != "apertus" { + return false + } + + data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) + if err != nil { + return false + } + var cfg struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + LLMConfig struct { + ModelType string `json:"model_type"` + } `json:"llm_config"` + } + if json.Unmarshal(data, &cfg) != nil { + return false + } + + for _, identifier := range append(slices.Clone(cfg.Architectures), cfg.ModelType, cfg.LLMConfig.ModelType) { + if isApertusFamily(identifier) { + return true + } + } + return false +} + // readChatTemplate returns the model's chat template, preferring the // chat_template field of tokenizer_config.json and falling back to a standalone // chat_template.jinja. It returns "" when neither is present. diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 84cef9c7f35..5be574acc65 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -486,21 +486,21 @@ func TestApertusMetadataInference(t *testing.T) { configJSON: `{"architectures":["ApertusForCausalLM"],"model_type":"qwen3","llm_config":{"model_type":"gpt_oss"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools", "thinking"}, + caps: []string{"completion", "tools"}, }, { name: "top-level wins over nested model type", configJSON: `{"model_type":"apertus","llm_config":{"model_type":"qwen3"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools", "thinking"}, + caps: []string{"completion", "tools"}, }, { name: "nested model type", configJSON: `{"model_type":"wrapper","llm_config":{"model_type":"apertus"}}`, parser: "apertus", renderer: "apertus", - caps: []string{"completion", "tools", "thinking"}, + caps: []string{"completion", "tools"}, }, { name: "near matches are not Apertus", From 2fd492fa25fa20c4f623f87db3c43952a4b9dd54 Mon Sep 17 00:00:00 2001 From: Philipp Date: Thu, 27 Aug 2026 07:43:40 +0000 Subject: [PATCH 40/58] server: preserve Apertus structured output content Co-authored-by: Codex --- model/parsers/apertus.go | 32 +++++++- model/parsers/apertus_test.go | 120 ++++++++++++++++++++++++++++ model/parsers/parsers.go | 7 ++ server/routes.go | 12 ++- server/routes_generate_test.go | 140 +++++++++++++++++++++++++++++++++ 5 files changed, 305 insertions(+), 6 deletions(-) diff --git a/model/parsers/apertus.go b/model/parsers/apertus.go index 2415dd6b82a..5364a4f2090 100644 --- a/model/parsers/apertus.go +++ b/model/parsers/apertus.go @@ -40,9 +40,14 @@ type ApertusParser struct { callIndex int pendingBare bool thinking bool + format bool } func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool { + return p.InitWithFormat(tools, lastMessage, thinkValue, nil) +} + +func (p *ApertusParser) InitWithFormat(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue, format json.RawMessage) []api.Tool { p.state = apertusContent p.acc.Reset() p.allowedTool = make(map[string]struct{}, len(tools)) @@ -50,6 +55,7 @@ func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkVa p.callIndex = 0 p.pendingBare = false p.thinking = thinkValue != nil && thinkValue.Bool() + p.format = apertusResponseFormatActive(format) for _, tool := range tools { name := tool.Function.Name if !apertusIdentifier(name) { @@ -65,6 +71,13 @@ func (p *ApertusParser) Init(tools []api.Tool, lastMessage *api.Message, thinkVa return tools } +func apertusResponseFormatActive(format json.RawMessage) bool { + if len(format) == 0 || bytes.Equal(format, []byte("null")) || bytes.Equal(format, []byte(`""`)) { + return false + } + return bytes.Equal(format, []byte(`"json"`)) || format[0] == '{' && json.Valid(format) +} + func (p *ApertusParser) Add(s string, done bool) (content, thinking string, calls []api.ToolCall, err error) { if p.initErr != nil { return "", "", nil, p.initErr @@ -109,6 +122,11 @@ func (p *ApertusParser) Add(s string, done bool) (content, thinking string, call p.acc.Reset() p.pendingBare = false return out.String(), thought.String(), parsed, nil + } else if p.format && json.Valid([]byte(cleaned)) { + p.acc.Reset() + p.pendingBare = false + out.WriteString(cleaned) + return out.String(), thought.String(), calls, nil } else if !isSoftApertusToolParseError(parseErr) { return "", "", nil, parseErr } @@ -237,7 +255,7 @@ func (p *ApertusParser) pendingFragmentWouldExceed(s string) bool { if idx := strings.Index(probe, apertusToolOpenTag); idx >= 0 { return p.acc.Len()+len(s)-idx-len(apertusToolOpenTag) > maxApertusToolCallBytes } - return apertusPendingBarePrefix(probe) && p.acc.Len()+len(s) > maxApertusToolCallBytes + return (apertusPendingBarePrefix(probe) || p.looksLikeToolCallStart(probe)) && p.acc.Len()+len(s) > maxApertusToolCallBytes } func apertusPendingBarePrefix(s string) bool { @@ -301,9 +319,9 @@ func (p *ApertusParser) parseToolCalls(raw string) ([]api.ToolCall, error) { } } } - calls = append(calls, api.ToolCall{Function: api.ToolCallFunction{Index: p.callIndex, Name: name, Arguments: args}}) - p.callIndex++ + calls = append(calls, api.ToolCall{Function: api.ToolCallFunction{Index: p.callIndex + len(calls), Name: name, Arguments: args}}) } + p.callIndex += len(calls) return calls, nil } @@ -359,7 +377,13 @@ func (p *ApertusParser) looksLikeToolCall(s string) bool { func (p *ApertusParser) looksLikeToolCallStart(s string) bool { s = strings.TrimSpace(s) - return len(p.allowedTool) > 0 && s != "" && (strings.HasPrefix("[{", s) || strings.HasPrefix("{", s) || strings.HasPrefix(s, "[{") || strings.HasPrefix(s, "{")) + if len(p.allowedTool) == 0 || s == "" { + return false + } + if p.format && (strings.HasPrefix(s, "[") || strings.HasPrefix(s, "{")) { + return true + } + return strings.HasPrefix("[{", s) || strings.HasPrefix("{", s) || strings.HasPrefix(s, "[{") || strings.HasPrefix(s, "{") } func apertusIdentifier(s string) bool { diff --git a/model/parsers/apertus_test.go b/model/parsers/apertus_test.go index 8ec0e668b84..ee94e00479a 100644 --- a/model/parsers/apertus_test.go +++ b/model/parsers/apertus_test.go @@ -1,12 +1,132 @@ package parsers import ( + "encoding/json" "strings" "testing" "github.com/ollama/ollama/api" ) +func TestApertusParserFormatFallback(t *testing.T) { + format := json.RawMessage(`{"type":"object"}`) + for _, response := range []string{ + `{"answer":"forty-two"}`, + `[{"answer":"forty-two"}]`, + } { + t.Run(response, func(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, format) + content, _, calls, err := p.Add(response, true) + if err != nil { + t.Fatal(err) + } + if content != response || len(calls) != 0 { + t.Fatalf("content=%q calls=%#v", content, calls) + } + }) + } +} + +func TestApertusParserFormatDoesNotRelaxToolCallErrors(t *testing.T) { + format := json.RawMessage(`{"type":"object"}`) + + t.Run("normal unknown bare call", func(t *testing.T) { + p := &ApertusParser{} + p.Init([]api.Tool{apertusParserTool("known")}, nil, nil) + if _, _, _, err := p.Add(`{"unknown":{}}`, true); err == nil { + t.Fatal("unknown bare call accepted without a response format") + } + }) + + t.Run("tagged unknown call", func(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, format) + if _, _, _, err := p.Add(apertusToolOpenTag+`{"unknown":{}}`+apertusToolCloseTag, true); err == nil { + t.Fatal("tagged unknown call accepted with a response format") + } + }) + + t.Run("declared bare call", func(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, format) + content, _, calls, err := p.Add(`{"known":{"value":1}}`, true) + if err != nil { + t.Fatal(err) + } + if content != "" || len(calls) != 1 || calls[0].Function.Index != 0 || calls[0].Function.Name != "known" { + t.Fatalf("content=%q calls=%#v", content, calls) + } + }) + + t.Run("format is request scoped", func(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, format) + p.Init([]api.Tool{apertusParserTool("known")}, nil, nil) + if _, _, _, err := p.Add(`{"unknown":{}}`, true); err == nil { + t.Fatal("format from a previous request remained active") + } + }) + + for _, inactive := range []json.RawMessage{nil, json.RawMessage(`null`), json.RawMessage(`""`), json.RawMessage(`{"type":`)} { + t.Run("inactive format "+string(inactive), func(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, inactive) + if _, _, _, err := p.Add(`{"unknown":{}}`, true); err == nil { + t.Fatalf("unknown bare call accepted with inactive format %q", inactive) + } + }) + } +} + +func TestApertusParserFormatFallbackCallIndexIsTransactional(t *testing.T) { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, json.RawMessage(`{"type":"array"}`)) + + mixed := `[{"known":{}},{"unknown":{}}]` + content, _, calls, err := p.Add(mixed, true) + if err != nil { + t.Fatal(err) + } + if content != mixed || len(calls) != 0 { + t.Fatalf("content=%q calls=%#v", content, calls) + } + + content, _, calls, err = p.Add(`{"known":{}}`, true) + if err != nil { + t.Fatal(err) + } + if content != "" || len(calls) != 1 || calls[0].Function.Index != 0 { + t.Fatalf("content=%q calls=%#v", content, calls) + } +} + +func TestApertusParserFormatFallbackStreamingBoundaries(t *testing.T) { + format := json.RawMessage(`{"type":"object"}`) + for _, response := range []string{ + `{"answer":"forty-two"}`, + `[{"answer":"forty-two"}]`, + } { + t.Run(response, func(t *testing.T) { + for split := range len(response) + 1 { + p := &ApertusParser{} + p.InitWithFormat([]api.Tool{apertusParserTool("known")}, nil, nil, format) + content1, _, calls1, err := p.Add(response[:split], false) + if err != nil { + t.Fatalf("split %d first chunk: %v", split, err) + } + content2, _, calls2, err := p.Add(response[split:], true) + if err != nil { + t.Fatalf("split %d final chunk: %v", split, err) + } + if content1+content2 != response || len(calls1)+len(calls2) != 0 { + t.Fatalf("split %d: content=%q calls=%d", split, content1+content2, len(calls1)+len(calls2)) + } + } + }) + } +} + func TestApertusParserGrammarAndStreamingBoundaries(t *testing.T) { input := `prefix <|tools_prefix|>[{"get_weather":{"city":"Bern"}}]<|tools_suffix|><|assistant_end|>` for i := range len(input) + 1 { diff --git a/model/parsers/parsers.go b/model/parsers/parsers.go index cd9062d414a..a4281d03d95 100644 --- a/model/parsers/parsers.go +++ b/model/parsers/parsers.go @@ -1,6 +1,7 @@ package parsers import ( + "encoding/json" "strings" "unicode" "unicode/utf8" @@ -23,6 +24,12 @@ type Parser interface { HasThinkingSupport() bool } +// RequestFormatParser is implemented by parsers whose output interpretation +// depends on the structured response format for the current request. +type RequestFormatParser interface { + InitWithFormat(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue, format json.RawMessage) []api.Tool +} + type ParserConstructor func() Parser type ParserRegistry struct { diff --git a/server/routes.go b/server/routes.go index cddd82785b6..cc0dc120914 100644 --- a/server/routes.go +++ b/server/routes.go @@ -441,7 +441,11 @@ func (s *Server) GenerateHandler(c *gin.Context) { builtinParser = parsers.ParserForName(m.Config.Parser) if builtinParser != nil { // no tools or last message for generate endpoint - builtinParser.Init(nil, nil, req.Think) + if formatParser, ok := builtinParser.(parsers.RequestFormatParser); ok { + formatParser.InitWithFormat(nil, nil, req.Think, req.Format) + } else { + builtinParser.Init(nil, nil, req.Think) + } } } @@ -2688,7 +2692,11 @@ func (s *Server) ChatHandler(c *gin.Context) { lastMessage = &msgs[len(msgs)-1] } // Initialize parser and get processed tools - processedTools = builtinParser.Init(req.Tools, lastMessage, req.Think) + if formatParser, ok := builtinParser.(parsers.RequestFormatParser); ok { + processedTools = formatParser.InitWithFormat(req.Tools, lastMessage, req.Think, req.Format) + } else { + processedTools = builtinParser.Init(req.Tools, lastMessage, req.Think) + } } } diff --git a/server/routes_generate_test.go b/server/routes_generate_test.go index 5de22f78b0c..67b2cc2e740 100644 --- a/server/routes_generate_test.go +++ b/server/routes_generate_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -2938,6 +2939,145 @@ func TestChatFormatWithThinkFalse(t *testing.T) { } } +func TestChatApertusToolsAndFormat(t *testing.T) { + gin.SetMode(gin.TestMode) + + mock := &mockRunner{} + s := newServerWithMockRunner(t, mock) + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(1), + "llama.context_length": uint32(8192), + "llama.embedding_length": uint32(4096), + "llama.attention.head_count": uint32(32), + "llama.attention.head_count_kv": uint32(8), + "tokenizer.ggml.tokens": []string{""}, + "tokenizer.ggml.scores": []float32{0}, + "tokenizer.ggml.token_type": []int32{0}, + }, []*ggml.Tensor{ + {Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_gate.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_up.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_k.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_q.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_v.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + }) + + modelName := "test-apertus-tools-and-format" + w := createRequest(t, s.CreateHandler, api.CreateRequest{ + Model: modelName, + Files: map[string]string{"file.gguf": digest}, + Parser: "apertus", + Template: `{{- range .Messages }}{{ .Role }}: {{ .Content }}{{ end }}`, + Stream: &stream, + }) + if w.Code != http.StatusOK { + t.Fatalf("create: expected status 200, got %d: %s", w.Code, w.Body.String()) + } + + tool := api.Tool{Type: "function", Function: api.ToolFunction{ + Name: "get_weather", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: testPropsMap(map[string]api.ToolProperty{ + "city": {Type: api.PropertyType{"string"}}, + }), + }, + }} + + responses := []struct { + name string + format json.RawMessage + response string + }{ + { + name: "object", + format: json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}`), + response: `{"answer":"forty-two"}`, + }, + { + name: "array", + format: json.RawMessage(`{"type":"array","items":{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}}`), + response: `[{"answer":"forty-two"}]`, + }, + } + + for _, tt := range responses { + for _, streaming := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/stream=%t", tt.name, streaming), func(t *testing.T) { + completionCount := 0 + var forwardedFormat json.RawMessage + mock.CompletionFn = func(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error { + completionCount++ + forwardedFormat = append(forwardedFormat[:0], req.Format...) + for i := range len(tt.response) { + fn(llm.CompletionResponse{Content: tt.response[i : i+1]}) + } + fn(llm.CompletionResponse{ + Done: true, + DoneReason: llm.DoneReasonStop, + PromptEvalCount: 1, + PromptEvalDuration: 1, + EvalCount: 1, + EvalDuration: 1, + }) + return nil + } + + think := false + w := createRequest(t, s.ChatHandler, api.ChatRequest{ + Model: modelName, + Messages: []api.Message{{Role: "user", Content: "Respond in the requested format."}}, + Tools: []api.Tool{tool}, + Think: &api.ThinkValue{Value: think}, + Stream: &streaming, + Format: tt.format, + }) + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + if completionCount != 1 { + t.Fatalf("completion calls=%d, want 1", completionCount) + } + if !bytes.Equal(forwardedFormat, tt.format) { + t.Fatalf("forwarded format=%q, want %q", forwardedFormat, tt.format) + } + + decoder := json.NewDecoder(w.Body) + var content strings.Builder + var events []api.ChatResponse + for { + var event api.ChatResponse + if err := decoder.Decode(&event); err == io.EOF { + break + } else if err != nil { + t.Fatal(err) + } + events = append(events, event) + content.WriteString(event.Message.Content) + if len(event.Message.ToolCalls) != 0 { + t.Fatalf("unexpected tool calls: %#v", event.Message.ToolCalls) + } + } + if len(events) != 1 { + t.Fatalf("response events=%d, want 1", len(events)) + } + if content.String() != tt.response { + t.Fatalf("content=%q, want %q", content.String(), tt.response) + } + if !events[0].Done { + t.Fatal("response was not terminal") + } + }) + } + } +} + func TestGenerateUnload(t *testing.T) { gin.SetMode(gin.TestMode) From a413b35e59f9f6790daab07fce81a308b6d13235 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 20:40:52 +0000 Subject: [PATCH 41/58] gemma4: add MLX image embedding path Parse Gemma 4 multimodal config and tokenizer tokens, load the vision tower plus embed_vision projector, preprocess ordered images through the official media contract, and scatter lazy projected vision features into the language embedding stream. Co-authored-by: Codex --- x/models/gemma4/gemma4.go | 101 +++- x/models/gemma4/gemma4_test.go | 61 +++ x/models/gemma4/vision.go | 916 +++++++++++++++++++++++++++++++++ x/models/gemma4/vision_test.go | 298 +++++++++++ 4 files changed, 1373 insertions(+), 3 deletions(-) create mode 100644 x/models/gemma4/vision.go create mode 100644 x/models/gemma4/vision_test.go diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index f390289b329..edebc08e6ba 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -67,6 +67,10 @@ type TextConfig struct { ExpertIntermediateSize int32 `json:"moe_intermediate_size"` RopeParameters map[string]*RopeParams `json:"rope_parameters"` ImageTokenIDValue int32 `json:"image_token_id"` + AudioTokenIDValue int32 `json:"-"` + BOITokenIDValue int32 `json:"-"` + EOITokenIDValue int32 `json:"-"` + VisionSoftTokens int32 `json:"-"` // Quantization parameters. QuantGroupSize int `json:"-"` @@ -374,6 +378,8 @@ type Model struct { Layers []*DecoderLayer Norm *nn.RMSNorm LMHead nn.LinearLayer + Vision *VisionModel + EmbedVision *MultimodalEmbedder // PLE model-level components (nil if no PLE). EmbedTokensPerLayer nn.EmbeddingLayer @@ -386,9 +392,11 @@ type Model struct { tok *tokenizer.Tokenizer *TextConfig + VisionConfig *VisionConfig SuppressLogitBias *mlx.Array weightPrefix string + mediaTokens gemma4MediaTokens } func parseTextConfig(configData []byte) (TextConfig, error) { @@ -407,6 +415,23 @@ func parseTextConfig(configData []byte) (TextConfig, error) { if wrapped.TextConfig != nil { cfg = *wrapped.TextConfig } + var top struct { + ImageTokenID int32 `json:"image_token_id"` + AudioTokenID int32 `json:"audio_token_id"` + BOITokenID int32 `json:"boi_token_id"` + EOITokenID int32 `json:"eoi_token_id"` + VisionSoftTokensPerImage int32 `json:"vision_soft_tokens_per_image"` + } + if err := json.Unmarshal(configData, &top); err != nil { + return TextConfig{}, fmt.Errorf("parse top-level multimodal config: %w", err) + } + if cfg.ImageTokenIDValue == 0 { + cfg.ImageTokenIDValue = top.ImageTokenID + } + cfg.AudioTokenIDValue = top.AudioTokenID + cfg.BOITokenIDValue = top.BOITokenID + cfg.EOITokenIDValue = top.EOITokenID + cfg.VisionSoftTokens = top.VisionSoftTokensPerImage // Apply defaults. if cfg.HeadDim == 0 { @@ -433,6 +458,24 @@ func parseTextConfig(configData []byte) (TextConfig, error) { if cfg.MaxPositionEmbeddings == 0 { cfg.MaxPositionEmbeddings = 131072 } + if cfg.ImageTokenIDValue == 0 { + cfg.ImageTokenIDValue = 258880 + } + if cfg.AudioTokenIDValue == 0 { + cfg.AudioTokenIDValue = 258881 + } + if cfg.BOITokenIDValue == 0 { + cfg.BOITokenIDValue = 255999 + } + if cfg.EOITokenIDValue == 0 { + cfg.EOITokenIDValue = 258882 + } + if cfg.VisionSoftTokens == 0 { + cfg.VisionSoftTokens = 280 + } + if err := validateGemma4ImageTokenConfig(&cfg); err != nil { + return TextConfig{}, err + } // Gemma 4 uses scaling=1.0 (no 1/sqrt(head_dim) scaling); the Q/K norms // handle magnitude control. This differs from Gemma 3 which uses @@ -513,6 +556,31 @@ func parseTextConfig(configData []byte) (TextConfig, error) { return cfg, nil } +func validateGemma4ImageTokenConfig(cfg *TextConfig) error { + if cfg.VisionSoftTokens <= 0 || cfg.VisionSoftTokens > maxGemma4VisionSoftTokens { + return fmt.Errorf("invalid Gemma4 vision soft-token count %d", cfg.VisionSoftTokens) + } + tokens := []struct { + name string + id int32 + }{ + {"boi_token_id", cfg.BOITokenIDValue}, + {"image_token_id", cfg.ImageTokenIDValue}, + {"eoi_token_id", cfg.EOITokenIDValue}, + } + seen := make(map[int32]string, len(tokens)) + for _, token := range tokens { + if token.id < 0 || token.id >= cfg.VocabSize { + return fmt.Errorf("invalid Gemma4 %s %d for vocab size %d", token.name, token.id, cfg.VocabSize) + } + if other, ok := seen[token.id]; ok { + return fmt.Errorf("Gemma4 %s duplicates %s at token %d", token.name, other, token.id) + } + seen[token.id] = token.name + } + return nil +} + func parseSuppressTokens(configData []byte) []int32 { var cfg generationConfig if err := json.Unmarshal(configData, &cfg); err != nil { @@ -610,6 +678,10 @@ func newModel(root *model.Root) (base.Model, error) { if err != nil { return nil, err } + visionConfig, err := parseVisionConfig(configData) + if err != nil { + return nil, err + } if qt := root.QuantType(); qt != "" { cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt) @@ -632,8 +704,10 @@ func newModel(root *model.Root) (base.Model, error) { tokConfig.GenerationConfigJSON = genConfigData suppressTokens = parseSuppressTokens(genConfigData) } + mediaTokens := defaultGemma4MediaTokens() if tokConfigData, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { tokConfig.TokenizerConfigJSON = tokConfigData + mediaTokens = parseGemma4MediaTokens(tokConfigData, mediaTokens) } tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) @@ -644,7 +718,9 @@ func newModel(root *model.Root) (base.Model, error) { m := &Model{ Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), TextConfig: &cfg, + VisionConfig: visionConfig, tok: tok, + mediaTokens: mediaTokens, SuppressLogitBias: makeSuppressLogitBias(suppressTokens, cfg.VocabSize), } @@ -695,6 +771,19 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { m.LMHead = m.EmbedTokens.AsLinear() } + if m.VisionConfig != nil && hasGemma4VisionWeights(tensors) { + vision, err := loadVisionModel(tensors, m.VisionConfig, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + embedVision, err := loadMultimodalEmbedder(tensors, "embed_vision", m.VisionConfig.RMSNormEps, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + m.Vision = vision + m.EmbedVision = embedVision + } + // PLE model-level weights. if m.HiddenSizePerLayer > 0 { pleEmbed := model.MakeEmbeddingLayer(tensors, prefix+"embed_tokens_per_layer", m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) @@ -983,13 +1072,19 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden 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) - h = mlx.MulScalar(h, m.EmbedScale) + h := m.TokenEmbeddings(b.InputIDs) + if len(b.Media) > 0 { + h = m.scatterMedia(h, b) + } // Compute PLE inputs if configured. var perLayerInputs *mlx.Array if m.HiddenSizePerLayer > 0 && m.EmbedTokensPerLayer != nil { - perLayerInputs = m.computePLEInputs(b.InputIDs, h) + pleTokens := b.InputIDs + if len(b.Media) > 0 { + pleTokens = gemma4PLETokens(pleTokens, b) + } + perLayerInputs = m.computePLEInputs(pleTokens, h) } // KV sharing: each donor layer stores its KVHistory here so later diff --git a/x/models/gemma4/gemma4_test.go b/x/models/gemma4/gemma4_test.go index 7a85dca407e..27f190434e9 100644 --- a/x/models/gemma4/gemma4_test.go +++ b/x/models/gemma4/gemma4_test.go @@ -24,6 +24,67 @@ func TestParseSuppressTokens(t *testing.T) { } } +func TestParseTextConfigTopLevelMediaFields(t *testing.T) { + data := []byte(`{ + "image_token_id": 258880, + "audio_token_id": 258881, + "boi_token_id": 255999, + "eoi_token_id": 258882, + "vision_soft_tokens_per_image": 280, + "text_config": { + "hidden_size": 2560, + "num_hidden_layers": 42, + "intermediate_size": 10240, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 256, + "global_head_dim": 512, + "vocab_size": 262144, + "rms_norm_eps": 1e-6 + } + }`) + + cfg, err := parseTextConfig(data) + if err != nil { + t.Fatalf("parseTextConfig() error = %v", err) + } + if cfg.ImageTokenIDValue != 258880 { + t.Fatalf("ImageTokenIDValue = %d, want 258880", cfg.ImageTokenIDValue) + } + if cfg.AudioTokenIDValue != 258881 { + t.Fatalf("AudioTokenIDValue = %d, want 258881", cfg.AudioTokenIDValue) + } + if cfg.BOITokenIDValue != 255999 { + t.Fatalf("BOITokenIDValue = %d, want 255999", cfg.BOITokenIDValue) + } + if cfg.EOITokenIDValue != 258882 { + t.Fatalf("EOITokenIDValue = %d, want 258882", cfg.EOITokenIDValue) + } + if cfg.VisionSoftTokens != 280 { + t.Fatalf("VisionSoftTokens = %d, want 280", cfg.VisionSoftTokens) + } +} + +func TestParseTextConfigRejectsInvalidImageMarkers(t *testing.T) { + tests := []struct { + name string + json string + }{ + {"duplicate begin and image", `{"boi_token_id":258880,"image_token_id":258880}`}, + {"duplicate image and end", `{"image_token_id":258880,"eoi_token_id":258880}`}, + {"negative image", `{"image_token_id":-1}`}, + {"image outside vocab", `{"image_token_id":262144}`}, + {"unbounded soft tokens", `{"vision_soft_tokens_per_image":16385}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if _, err := parseTextConfig([]byte(tt.json)); err == nil { + t.Fatal("parseTextConfig() error = nil") + } + }) + } +} + func TestParseTextConfigE2B(t *testing.T) { skipIfNoMLX(t) data := []byte(`{ diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go new file mode 100644 index 00000000000..9d6ca71cc7a --- /dev/null +++ b/x/models/gemma4/vision.go @@ -0,0 +1,916 @@ +package gemma4 + +// Portions of the Gemma 4 vision preprocessing and embedding flow are adapted +// from MLX-VLM's MIT-licensed Gemma 4 implementation. + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "image" + _ "image/jpeg" + _ "image/png" + "math" + + xdraw "golang.org/x/image/draw" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "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" +) + +const ( + defaultGemma4BOIToken = "<|image>" + defaultGemma4ImageToken = "<|image|>" + defaultGemma4EOIToken = "" + + maxGemma4ImageBytes = 32 << 20 + maxGemma4ImageDimension = 16_384 + maxGemma4ImagePixels = 64 << 20 + maxGemma4ResizePixels = 1 << 20 + maxGemma4VisionHiddenSize = 32_768 + maxGemma4VisionIntermediate = 262_144 + maxGemma4VisionLayers = 512 + maxGemma4VisionHeads = 512 + maxGemma4VisionHeadDim = 1_024 + maxGemma4VisionSoftTokens = 16_384 + maxGemma4PositionTableEntries = 1 << 20 + maxGemma4PositionValues = 1 << 26 + maxIntValue = int64(^uint(0) >> 1) +) + +type VisionRopeParameters struct { + RopeTheta float32 `json:"rope_theta"` + RopeType string `json:"rope_type"` +} + +type VisionConfig struct { + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + RMSNormEps float32 `json:"rms_norm_eps"` + DefaultOutputLength int32 `json:"default_output_length"` + PatchSize int32 `json:"patch_size"` + PositionEmbeddingSize int32 `json:"position_embedding_size"` + PoolingKernelSize int32 `json:"pooling_kernel_size"` + UseClippedLinears bool `json:"use_clipped_linears"` + Standardize bool `json:"standardize"` + RopeParameters VisionRopeParameters `json:"rope_parameters"` +} + +type gemma4MediaTokens struct { + BOI string + Image string + EOI string +} + +type gemma4ImageInput struct { + Pixels []float32 + Width int + Height int + PatchWidth int + PatchHeight int + SoftTokens int +} + +type gemma4MediaPayload struct { + Image gemma4ImageInput + ImageStart int + ImageEnd int +} + +type ClippableLinear struct { + Linear nn.LinearLayer + InputMin *mlx.Array + InputMax *mlx.Array + OutputMin *mlx.Array + OutputMax *mlx.Array +} + +type VisionAttention struct { + QProj *ClippableLinear + KProj *ClippableLinear + VProj *ClippableLinear + OProj *ClippableLinear + + QNorm *nn.RMSNorm + KNorm *nn.RMSNorm +} + +type VisionMLP struct { + GateProj *ClippableLinear + UpProj *ClippableLinear + DownProj *ClippableLinear +} + +type VisionLayer struct { + Attention *VisionAttention + MLP *VisionMLP + + InputNorm *nn.RMSNorm + PostAttnNorm *nn.RMSNorm + PreFFNorm *nn.RMSNorm + PostFFNorm *nn.RMSNorm +} + +type VisionPatchEmbedder struct { + InputProj nn.LinearLayer + PositionEmbeddingTable *mlx.Array + PatchSize int32 + PositionEmbeddingSize int32 +} + +type VisionModel struct { + Config *VisionConfig + PatchEmbedder *VisionPatchEmbedder + Layers []*VisionLayer + StdBias *mlx.Array + StdScale *mlx.Array +} + +type MultimodalEmbedder struct { + Projection nn.LinearLayer + Eps float32 +} + +type visionPositionArrays struct { + X *mlx.Array + Y *mlx.Array + RopeCos *mlx.Array + RopeSin *mlx.Array +} + +func parseVisionConfig(configData []byte) (*VisionConfig, error) { + var wrapped struct { + VisionConfig *VisionConfig `json:"vision_config"` + } + if err := json.Unmarshal(configData, &wrapped); err != nil { + return nil, fmt.Errorf("parse vision config: %w", err) + } + if wrapped.VisionConfig == nil { + return nil, nil + } + cfg := *wrapped.VisionConfig + if cfg.HiddenSize == 0 { + cfg.HiddenSize = 768 + } + if cfg.IntermediateSize == 0 { + cfg.IntermediateSize = 3072 + } + if cfg.NumHiddenLayers == 0 { + cfg.NumHiddenLayers = 16 + } + if cfg.NumAttentionHeads == 0 { + cfg.NumAttentionHeads = 12 + } + if cfg.NumKeyValueHeads == 0 { + cfg.NumKeyValueHeads = cfg.NumAttentionHeads + } + if cfg.HeadDim == 0 { + cfg.HeadDim = 64 + } + if cfg.RMSNormEps == 0 { + cfg.RMSNormEps = 1e-6 + } + if cfg.DefaultOutputLength == 0 { + cfg.DefaultOutputLength = 280 + } + if cfg.PatchSize == 0 { + cfg.PatchSize = 16 + } + if cfg.PositionEmbeddingSize == 0 { + cfg.PositionEmbeddingSize = 10240 + } + if cfg.PoolingKernelSize == 0 { + cfg.PoolingKernelSize = 3 + } + if cfg.RopeParameters.RopeTheta == 0 { + cfg.RopeParameters.RopeTheta = 100 + } + if err := validateVisionConfig(&cfg); err != nil { + return nil, err + } + return &cfg, nil +} + +func validateVisionConfig(cfg *VisionConfig) error { + if cfg == nil { + return errors.New("missing Gemma4 vision config") + } + positiveBounded := func(name string, value, limit int32) error { + if value <= 0 || value > limit { + return fmt.Errorf("invalid Gemma4 vision %s %d (limit %d)", name, value, limit) + } + return nil + } + for _, field := range []struct { + name string + value, limit int32 + }{ + {"hidden_size", cfg.HiddenSize, maxGemma4VisionHiddenSize}, + {"intermediate_size", cfg.IntermediateSize, maxGemma4VisionIntermediate}, + {"num_hidden_layers", cfg.NumHiddenLayers, maxGemma4VisionLayers}, + {"num_attention_heads", cfg.NumAttentionHeads, maxGemma4VisionHeads}, + {"num_key_value_heads", cfg.NumKeyValueHeads, maxGemma4VisionHeads}, + {"head_dim", cfg.HeadDim, maxGemma4VisionHeadDim}, + {"default_output_length", cfg.DefaultOutputLength, maxGemma4VisionSoftTokens}, + {"patch_size", cfg.PatchSize, maxGemma4ImageDimension}, + {"position_embedding_size", cfg.PositionEmbeddingSize, maxGemma4PositionTableEntries}, + {"pooling_kernel_size", cfg.PoolingKernelSize, maxGemma4ImageDimension}, + } { + if err := positiveBounded(field.name, field.value, field.limit); err != nil { + return err + } + } + if cfg.NumKeyValueHeads > cfg.NumAttentionHeads || cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return fmt.Errorf("invalid Gemma4 vision attention heads %d/%d", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) + } + attentionWidth := int64(cfg.NumAttentionHeads) * int64(cfg.HeadDim) + if attentionWidth != int64(cfg.HiddenSize) { + return fmt.Errorf("Gemma4 vision attention width %d does not match hidden_size %d", attentionWidth, cfg.HiddenSize) + } + if cfg.RMSNormEps <= 0 || math.IsNaN(float64(cfg.RMSNormEps)) || math.IsInf(float64(cfg.RMSNormEps), 0) { + return fmt.Errorf("invalid Gemma4 vision rms_norm_eps %v", cfg.RMSNormEps) + } + if cfg.RopeParameters.RopeTheta <= 0 || math.IsNaN(float64(cfg.RopeParameters.RopeTheta)) || math.IsInf(float64(cfg.RopeParameters.RopeTheta), 0) { + return fmt.Errorf("invalid Gemma4 vision rope_theta %v", cfg.RopeParameters.RopeTheta) + } + + if _, ok := checkedPositiveProduct(maxGemma4ResizePixels, + int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize), int64(cfg.PoolingKernelSize), int64(cfg.PatchSize), int64(cfg.PatchSize)); !ok { + return fmt.Errorf("Gemma4 vision resize budget exceeds limit %d", maxGemma4ResizePixels) + } + if _, ok := checkedPositiveProduct(maxGemma4PositionValues, + int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize), int64(cfg.PoolingKernelSize), int64(cfg.HeadDim)); !ok { + return fmt.Errorf("Gemma4 vision position allocation exceeds limit %d", maxGemma4PositionValues) + } + if _, ok := checkedPositiveProduct(maxGemma4ImageDimension, int64(cfg.PoolingKernelSize), int64(cfg.PatchSize)); !ok { + return fmt.Errorf("invalid Gemma4 pooled patch size (limit %d)", maxGemma4ImageDimension) + } + maxPatchSide, ok := checkedPositiveProduct(maxGemma4PositionTableEntries, int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize)) + if !ok || maxPatchSide > int64(cfg.PositionEmbeddingSize) { + return fmt.Errorf("Gemma4 vision patch side %d exceeds position table size %d", maxPatchSide, cfg.PositionEmbeddingSize) + } + return nil +} + +func checkedPositiveProduct(limit int64, values ...int64) (int64, bool) { + if limit <= 0 { + return 0, false + } + product := int64(1) + for _, value := range values { + if value <= 0 || product > limit/value { + return 0, false + } + product *= value + } + return product, true +} + +func defaultGemma4MediaTokens() gemma4MediaTokens { + return gemma4MediaTokens{ + BOI: defaultGemma4BOIToken, + Image: defaultGemma4ImageToken, + EOI: defaultGemma4EOIToken, + } +} + +func parseGemma4MediaTokens(data []byte, fallback gemma4MediaTokens) gemma4MediaTokens { + var cfg struct { + BOIToken string `json:"boi_token"` + ImageToken string `json:"image_token"` + EOIToken string `json:"eoi_token"` + } + if err := json.Unmarshal(data, &cfg); err != nil { + return fallback + } + if cfg.BOIToken != "" { + fallback.BOI = cfg.BOIToken + } + if cfg.ImageToken != "" { + fallback.Image = cfg.ImageToken + } + if cfg.EOIToken != "" { + fallback.EOI = cfg.EOIToken + } + return fallback +} + +func hasGemma4VisionWeights(tensors map[string]*mlx.Array) bool { + return firstNonNil(tensors, + "vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.input_proj.weight", + ) != nil +} + +func resolveVisionPrefix(tensors map[string]*mlx.Array) string { + if tensors["vision_tower.patch_embedder.input_proj.weight"] != nil { + return "" + } + if tensors["model.vision_tower.patch_embedder.input_proj.weight"] != nil { + return "model." + } + return "" +} + +func loadVisionModel(tensors map[string]*mlx.Array, cfg *VisionConfig, groupSize, bits int, mode string, tq map[string]*model.TensorQuantInfo) (*VisionModel, error) { + if err := validateVisionConfig(cfg); err != nil { + return nil, err + } + prefix := resolveVisionPrefix(tensors) + linears := model.NewLinearFactory(tensors, groupSize, bits, mode, tq) + + patchProj := linears.Make(prefix + "vision_tower.patch_embedder.input_proj") + if patchProj == nil { + return nil, fmt.Errorf("missing vision patch projection") + } + posTable := tensors[prefix+"vision_tower.patch_embedder.position_embedding_table"] + if posTable == nil { + return nil, fmt.Errorf("missing vision position embedding table") + } + + v := &VisionModel{ + Config: cfg, + PatchEmbedder: &VisionPatchEmbedder{ + InputProj: patchProj, + PositionEmbeddingTable: posTable, + PatchSize: cfg.PatchSize, + PositionEmbeddingSize: cfg.PositionEmbeddingSize, + }, + Layers: make([]*VisionLayer, cfg.NumHiddenLayers), + } + + for i := range cfg.NumHiddenLayers { + layerPrefix := fmt.Sprintf("%svision_tower.encoder.layers.%d", prefix, i) + layer := &VisionLayer{ + Attention: &VisionAttention{ + QProj: makeClippableLinear(tensors, linears, layerPrefix+".self_attn.q_proj", cfg.UseClippedLinears), + KProj: makeClippableLinear(tensors, linears, layerPrefix+".self_attn.k_proj", cfg.UseClippedLinears), + VProj: makeClippableLinear(tensors, linears, layerPrefix+".self_attn.v_proj", cfg.UseClippedLinears), + OProj: makeClippableLinear(tensors, linears, layerPrefix+".self_attn.o_proj", cfg.UseClippedLinears), + QNorm: nn.NewRMSNorm(tensors[layerPrefix+".self_attn.q_norm.weight"], cfg.RMSNormEps), + KNorm: nn.NewRMSNorm(tensors[layerPrefix+".self_attn.k_norm.weight"], cfg.RMSNormEps), + }, + MLP: &VisionMLP{ + GateProj: makeClippableLinear(tensors, linears, layerPrefix+".mlp.gate_proj", cfg.UseClippedLinears), + UpProj: makeClippableLinear(tensors, linears, layerPrefix+".mlp.up_proj", cfg.UseClippedLinears), + DownProj: makeClippableLinear(tensors, linears, layerPrefix+".mlp.down_proj", cfg.UseClippedLinears), + }, + InputNorm: nn.NewRMSNorm(tensors[layerPrefix+".input_layernorm.weight"], cfg.RMSNormEps), + PostAttnNorm: nn.NewRMSNorm(tensors[layerPrefix+".post_attention_layernorm.weight"], cfg.RMSNormEps), + PreFFNorm: nn.NewRMSNorm(tensors[layerPrefix+".pre_feedforward_layernorm.weight"], cfg.RMSNormEps), + PostFFNorm: nn.NewRMSNorm(tensors[layerPrefix+".post_feedforward_layernorm.weight"], cfg.RMSNormEps), + } + if layer.Attention.QProj == nil || layer.Attention.KProj == nil || layer.Attention.VProj == nil || layer.Attention.OProj == nil { + return nil, fmt.Errorf("vision layer %d: missing attention projection", i) + } + if layer.Attention.QNorm.Weight == nil || layer.Attention.KNorm.Weight == nil { + return nil, fmt.Errorf("vision layer %d: missing attention norm", i) + } + if layer.MLP.GateProj == nil || layer.MLP.UpProj == nil || layer.MLP.DownProj == nil { + return nil, fmt.Errorf("vision layer %d: missing mlp projection", i) + } + if layer.InputNorm.Weight == nil || layer.PostAttnNorm.Weight == nil || layer.PreFFNorm.Weight == nil || layer.PostFFNorm.Weight == nil { + return nil, fmt.Errorf("vision layer %d: missing block norm", i) + } + v.Layers[i] = layer + } + + if cfg.Standardize { + v.StdBias = tensors[prefix+"vision_tower.std_bias"] + v.StdScale = tensors[prefix+"vision_tower.std_scale"] + if err := validateVisionStandardizationTensors(cfg, v.StdBias != nil, v.StdScale != nil); err != nil { + return nil, err + } + } + + return v, nil +} + +func validateVisionStandardizationTensors(cfg *VisionConfig, hasBias, hasScale bool) error { + if cfg == nil || !cfg.Standardize { + return nil + } + if !hasBias || !hasScale { + return fmt.Errorf("missing Gemma4 vision standardization tensors: bias=%t scale=%t", hasBias, hasScale) + } + return nil +} + +func loadMultimodalEmbedder(tensors map[string]*mlx.Array, path string, eps float32, groupSize, bits int, mode string, tq map[string]*model.TensorQuantInfo) (*MultimodalEmbedder, error) { + linears := model.NewLinearFactory(tensors, groupSize, bits, mode, tq) + proj := linears.Make(path + ".embedding_projection") + if proj == nil { + proj = linears.Make("model." + path + ".embedding_projection") + } + if proj == nil { + return nil, fmt.Errorf("missing %s embedding projection", path) + } + return &MultimodalEmbedder{Projection: proj, Eps: eps}, nil +} + +func makeClippableLinear(tensors map[string]*mlx.Array, linears model.LinearFactory, path string, useClip bool) *ClippableLinear { + linear := linears.Make(path + ".linear") + if linear == nil { + linear = linears.Make(path) + } + if linear == nil { + return nil + } + out := &ClippableLinear{Linear: linear} + if useClip { + out.InputMin = tensors[path+".input_min"] + out.InputMax = tensors[path+".input_max"] + out.OutputMin = tensors[path+".output_min"] + out.OutputMax = tensors[path+".output_max"] + } + return out +} + +func (l *ClippableLinear) Forward(x *mlx.Array) *mlx.Array { + if l.InputMin != nil && l.InputMax != nil { + x = mlx.Clip(x, l.InputMin, l.InputMax) + } + out := l.Linear.Forward(x) + if l.OutputMin != nil && l.OutputMax != nil { + out = mlx.Clip(out, l.OutputMin, l.OutputMax) + } + return out +} + +func (m *MultimodalEmbedder) Forward(x *mlx.Array) *mlx.Array { + return m.Projection.Forward(mlx.RMSNormFn(x, nil, m.Eps)) +} + +// PrepareMedia implements the runner's media contract. Each image is expanded +// in stream order and remains a separate cache-identity item. +func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { + prepared := &base.PreparedRequest{} + for source, seg := range segments { + if seg.Data == nil { + prepared.Tokens = append(prepared.Tokens, seg.Tokens...) + continue + } + if m.VisionConfig == nil || m.Vision == nil || m.EmbedVision == nil { + return nil, fmt.Errorf("this model does not support %s input", seg.Kind) + } + if seg.Kind != "image" { + return nil, fmt.Errorf("gemma4 does not support %s input", seg.Kind) + } + + img, err := preprocessGemma4Image(seg.Data, m.VisionConfig, int(m.VisionSoftTokens)) + if err != nil { + return nil, err + } + start := len(prepared.Tokens) + prepared.Tokens = append(prepared.Tokens, m.BOITokenIDValue) + imageStart := len(prepared.Tokens) - start + for range img.SoftTokens { + prepared.Tokens = append(prepared.Tokens, m.ImageTokenIDValue) + } + imageEnd := len(prepared.Tokens) - start + prepared.Tokens = append(prepared.Tokens, m.EOITokenIDValue) + + geom := *img + pixels := geom.Pixels + geom.Pixels = nil + prepared.Items = append(prepared.Items, base.PreparedItem{ + Range: [2]int{start, len(prepared.Tokens)}, + Source: source, + MediaData: pixels, + Dims: []int{1, 3, geom.Height, geom.Width}, + Opaque: gemma4MediaPayload{ + Image: geom, + ImageStart: imageStart, + ImageEnd: imageEnd, + }, + }) + } + return prepared, nil +} + +// EncodeMedia builds the lazy image feature graph. The runner owns and frees +// MediaData, so pixels are always read from data rather than Opaque. +func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { + payload := item.Opaque.(gemma4MediaPayload) + pixels := mlx.Reshape(data, 1, 3, int32(payload.Image.Height), int32(payload.Image.Width)) + features := m.EmbedVision.Forward(m.Vision.Forward(pixels, &payload.Image)) + return mlx.Squeeze(features, 0) +} + +func gemma4ImageRun(item batch.MediaItem) (start, end int) { + payload := item.Opaque.(gemma4MediaPayload) + return item.Pos + payload.ImageStart, item.Pos + payload.ImageEnd +} + +func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch) *mlx.Array { + for _, item := range b.Media { + if item.Features == nil { + continue + } + start, end := gemma4ImageRun(item) + base := int(b.SeqOffsets[item.Seq]) + lo := max(start, base) + hi := min(end, base+int(b.SeqQueryLens[item.Seq])) + if hi <= lo { + continue + } + features := item.Features.Slice(mlx.Slice(lo-start, hi-start), mlx.Slice()) + features = mlx.Reshape(features.AsType(h.DType()), 1, int32(hi-lo), m.HiddenSize) + h = h.SliceUpdate(features, mlx.Slice(item.Seq, item.Seq+1), mlx.Slice(lo-base, hi-base), mlx.Slice()) + } + return h +} + +// pleTokens masks feature-bearing image tokens from PLE exactly where the +// same prepared items are scattered into the token embeddings. +func gemma4PLETokens(tokens *mlx.Array, b *batch.Batch) *mlx.Array { + for _, item := range b.Media { + start, end := gemma4ImageRun(item) + base := int(b.SeqOffsets[item.Seq]) + lo := max(start, base) + hi := min(end, base+int(b.SeqQueryLens[item.Seq])) + if hi <= lo { + continue + } + zeros := mlx.Zeros(mlx.DTypeInt32, 1, hi-lo) + tokens = tokens.SliceUpdate(zeros, mlx.Slice(item.Seq, item.Seq+1), mlx.Slice(lo-base, hi-base)) + } + return tokens +} + +func preprocessGemma4Image(data []byte, cfg *VisionConfig, maxSoftTokens int) (*gemma4ImageInput, error) { + if err := validateVisionConfig(cfg); err != nil { + return nil, err + } + if err := validateGemma4ImageDataSize(len(data)); err != nil { + return nil, err + } + imageConfig, _, err := image.DecodeConfig(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("decode Gemma4 image config: %w", err) + } + if err := validateGemma4ImageDimensions(imageConfig.Width, imageConfig.Height); err != nil { + return nil, err + } + img, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("decode Gemma4 image: %w", err) + } + b := img.Bounds() + width, height := b.Dx(), b.Dy() + if err := validateGemma4ImageDimensions(width, height); err != nil { + return nil, err + } + + patchSize := int(cfg.PatchSize) + pooling := int(cfg.PoolingKernelSize) + if patchSize <= 0 || pooling <= 0 { + return nil, fmt.Errorf("invalid Gemma4 vision patch configuration") + } + if maxSoftTokens <= 0 { + maxSoftTokens = int(cfg.DefaultOutputLength) + } + if maxSoftTokens > int(cfg.DefaultOutputLength) || maxSoftTokens > maxGemma4VisionSoftTokens { + return nil, fmt.Errorf("invalid Gemma4 image soft-token limit %d", maxSoftTokens) + } + maxPatches64, ok := checkedPositiveProduct(maxIntValue, int64(maxSoftTokens), int64(pooling), int64(pooling)) + if !ok { + return nil, errors.New("Gemma4 image patch budget exceeds platform limits") + } + maxPatches := int(maxPatches64) + targetW, targetH, err := gemma4ResizeDimensions(width, height, patchSize, maxPatches, pooling) + if err != nil { + return nil, err + } + + resized := img + if targetW != width || targetH != height { + dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH)) + xdraw.CatmullRom.Scale(dst, dst.Bounds(), img, b, xdraw.Over, nil) + resized = dst + } + + pixels, err := imageToCHWFloat32(resized) + if err != nil { + return nil, err + } + patchW := targetW / patchSize + patchH := targetH / patchSize + patchCount, ok := checkedPositiveProduct(maxIntValue, int64(patchW), int64(patchH)) + if !ok { + return nil, errors.New("Gemma4 image patch count exceeds platform limits") + } + if _, ok := checkedPositiveProduct(maxGemma4PositionValues, patchCount, int64(cfg.HeadDim)); !ok { + return nil, fmt.Errorf("Gemma4 image position allocation exceeds limit %d", maxGemma4PositionValues) + } + softTokens := int(patchCount / (int64(pooling) * int64(pooling))) + if softTokens <= 0 || softTokens > maxSoftTokens { + return nil, fmt.Errorf("Gemma4 image produced %d soft tokens, limit %d", softTokens, maxSoftTokens) + } + + return &gemma4ImageInput{ + Pixels: pixels, + Width: targetW, + Height: targetH, + PatchWidth: patchW, + PatchHeight: patchH, + SoftTokens: softTokens, + }, nil +} + +func validateGemma4ImageDataSize(size int) error { + if size <= 0 { + return errors.New("Gemma4 image is empty") + } + if size > maxGemma4ImageBytes { + return fmt.Errorf("Gemma4 image is %d bytes, limit %d", size, maxGemma4ImageBytes) + } + return nil +} + +func validateGemma4ImageDimensions(width, height int) error { + if width <= 0 || height <= 0 { + return fmt.Errorf("invalid Gemma4 image dimensions %dx%d", width, height) + } + if width > maxGemma4ImageDimension || height > maxGemma4ImageDimension { + return fmt.Errorf("Gemma4 image dimensions %dx%d exceed limit %dx%d", width, height, maxGemma4ImageDimension, maxGemma4ImageDimension) + } + pixels := int64(width) * int64(height) + if pixels > maxGemma4ImagePixels { + return fmt.Errorf("Gemma4 image has %d pixels, limit %d", pixels, maxGemma4ImagePixels) + } + return nil +} + +func gemma4ResizeDimensions(width, height, patchSize, maxPatches, poolingKernelSize int) (int, int, error) { + if width <= 0 || height <= 0 || patchSize <= 0 || maxPatches <= 0 || poolingKernelSize <= 0 { + return 0, 0, errors.New("invalid Gemma4 resize parameters") + } + if err := validateGemma4ImageDimensions(width, height); err != nil { + return 0, 0, err + } + targetPixels, ok := checkedPositiveProduct(maxGemma4ResizePixels, int64(maxPatches), int64(patchSize), int64(patchSize)) + if !ok { + return 0, 0, fmt.Errorf("Gemma4 resize target exceeds %d pixels", maxGemma4ResizePixels) + } + sourcePixels := int64(height) * int64(width) + targetPx := float64(targetPixels) + factor := math.Sqrt(targetPx / float64(sourcePixels)) + sideMult64, ok := checkedPositiveProduct(maxGemma4ImageDimension, int64(poolingKernelSize), int64(patchSize)) + if !ok { + return 0, 0, errors.New("invalid Gemma4 pooled patch size") + } + sideMult := int(sideMult64) + + targetH := int(math.Floor(factor*float64(height)/float64(sideMult))) * sideMult + targetW := int(math.Floor(factor*float64(width)/float64(sideMult))) * sideMult + + if targetH == 0 && targetW == 0 { + return 0, 0, errors.New("attempting to resize Gemma4 image to 0 x 0") + } + + poolingArea, ok := checkedPositiveProduct(maxIntValue, int64(poolingKernelSize), int64(poolingKernelSize)) + if !ok { + return 0, 0, errors.New("invalid Gemma4 pooling area") + } + maxSideLength64, ok := checkedPositiveProduct(maxGemma4ImageDimension, int64(maxPatches)/poolingArea, sideMult64) + if !ok { + return 0, 0, errors.New("invalid Gemma4 maximum resize side") + } + maxSideLength := int(maxSideLength64) + if targetH == 0 { + targetH = sideMult + targetW = min(int(math.Floor(float64(width)/float64(height)))*sideMult, maxSideLength) + } + if targetW == 0 { + targetW = sideMult + targetH = min(int(math.Floor(float64(height)/float64(width)))*sideMult, maxSideLength) + } + if targetW <= 0 || targetH <= 0 { + return 0, 0, fmt.Errorf("invalid Gemma4 resize target %dx%d", targetW, targetH) + } + if err := validateGemma4ImageDimensions(targetW, targetH); err != nil { + return 0, 0, fmt.Errorf("invalid Gemma4 resize target: %w", err) + } + if int64(targetW)*int64(targetH) > targetPixels { + return 0, 0, fmt.Errorf("Gemma4 resize target %dx%d exceeds %d-pixel patch budget", targetW, targetH, targetPixels) + } + return targetW, targetH, nil +} + +func imageToCHWFloat32(img image.Image) ([]float32, error) { + bounds := img.Bounds() + width, height := bounds.Dx(), bounds.Dy() + if err := validateGemma4ImageDimensions(width, height); err != nil { + return nil, err + } + plane64 := int64(height) * int64(width) + values64 := plane64 * 3 + if values64 <= 0 || values64 > maxIntValue { + return nil, errors.New("Gemma4 pixel allocation exceeds platform limits") + } + plane := int(plane64) + out := make([]float32, int(values64)) + for y := range height { + for x := range width { + r, g, blue, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA() + i := y*width + x + out[i] = float32(r) / 65535 + out[plane+i] = float32(g) / 65535 + out[2*plane+i] = float32(blue) / 65535 + } + } + return out, nil +} + +func (m *VisionModel) Forward(pixels *mlx.Array, img *gemma4ImageInput) *mlx.Array { + positions := m.positionArrays(img) + h := m.PatchEmbedder.Forward(pixels, positions) + for _, layer := range m.Layers { + h = layer.Forward(h, positions, m.Config) + } + h = m.pool(h, int32(img.PatchHeight), int32(img.PatchWidth)) + if m.Config.Standardize { + h = mlx.Mul(mlx.Sub(h, m.StdBias), m.StdScale) + } + return h +} + +func (m *VisionModel) positionArrays(img *gemma4ImageInput) visionPositionArrays { + L := img.PatchHeight * img.PatchWidth + xs := make([]int32, L) + ys := make([]int32, L) + cosVals := make([]float32, L*int(m.Config.HeadDim)) + sinVals := make([]float32, L*int(m.Config.HeadDim)) + idx := 0 + for y := range img.PatchHeight { + for x := range img.PatchWidth { + xs[idx] = int32(x) + ys[idx] = int32(y) + fillVisionRoPE(idx, x, y, int(m.Config.HeadDim), float64(m.Config.RopeParameters.RopeTheta), cosVals, sinVals) + idx++ + } + } + return visionPositionArrays{ + X: mlx.FromValues(xs, 1, L), + Y: mlx.FromValues(ys, 1, L), + RopeCos: mlx.FromValues(cosVals, 1, L, 1, int(m.Config.HeadDim)), + RopeSin: mlx.FromValues(sinVals, 1, L, 1, int(m.Config.HeadDim)), + } +} + +func fillVisionRoPE(row, x, y, headDim int, base float64, cosVals, sinVals []float32) { + ndim := 2 + channelsPerDim := 2 * (headDim / (2 * ndim)) + offset := row * headDim + for i := range headDim { + cosVals[offset+i] = 1 + } + if channelsPerDim == 0 { + return + } + half := channelsPerDim / 2 + positions := []int{x, y} + for dim, pos := range positions { + start := dim * channelsPerDim + for j := range half { + timescale := math.Pow(base, float64(2*j)/float64(channelsPerDim)) + angle := float64(pos) / timescale + c, s := float32(math.Cos(angle)), float32(math.Sin(angle)) + cosVals[offset+start+j] = c + cosVals[offset+start+half+j] = c + sinVals[offset+start+j] = s + sinVals[offset+start+half+j] = s + } + } +} + +func (p *VisionPatchEmbedder) Forward(pixelValues *mlx.Array, positions visionPositionArrays) *mlx.Array { + dims := pixelValues.Dims() + B, C, H, W := int32(dims[0]), int32(dims[1]), int32(dims[2]), int32(dims[3]) + patch := p.PatchSize + patchH := H / patch + patchW := W / patch + patches := mlx.Reshape(pixelValues, B, C, patchH, patch, patchW, patch) + patches = mlx.Transpose(patches, 0, 2, 4, 3, 5, 1) + patches = mlx.Reshape(patches, B, patchH*patchW, C*patch*patch) + patches = mlx.MulScalar(mlx.AddScalar(patches, -0.5), 2) + hidden := p.InputProj.Forward(patches) + return mlx.Add(hidden, p.positionEmbeddings(positions).AsType(hidden.DType())) +} + +func (p *VisionPatchEmbedder) positionEmbeddings(positions visionPositionArrays) *mlx.Array { + tableX := mlx.Squeeze(mlx.SliceStartStop(p.PositionEmbeddingTable, + []int32{0, 0, 0}, + []int32{1, p.PositionEmbeddingSize, int32(p.PositionEmbeddingTable.Dim(2))}, + ), 0) + tableY := mlx.Squeeze(mlx.SliceStartStop(p.PositionEmbeddingTable, + []int32{1, 0, 0}, + []int32{2, p.PositionEmbeddingSize, int32(p.PositionEmbeddingTable.Dim(2))}, + ), 0) + return mlx.Add(tableX.TakeAxis(positions.X, 0), tableY.TakeAxis(positions.Y, 0)) +} + +func (l *VisionLayer) Forward(x *mlx.Array, positions visionPositionArrays, cfg *VisionConfig) *mlx.Array { + normed := l.InputNorm.Forward(x, cfg.RMSNormEps) + attn := l.Attention.Forward(normed, positions, cfg) + attn = l.PostAttnNorm.Forward(attn, cfg.RMSNormEps) + h := mlx.Add(x, attn) + normed = l.PreFFNorm.Forward(h, cfg.RMSNormEps) + mlp := l.MLP.Forward(normed, cfg) + mlp = l.PostFFNorm.Forward(mlp, cfg.RMSNormEps) + return mlx.Add(h, mlp) +} + +func (a *VisionAttention) Forward(x *mlx.Array, positions visionPositionArrays, cfg *VisionConfig) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + + q := a.QProj.Forward(x) + q = mlx.Reshape(q, B, L, cfg.NumAttentionHeads, cfg.HeadDim) + k := a.KProj.Forward(x) + k = mlx.Reshape(k, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + v := a.VProj.Forward(x) + v = mlx.Reshape(v, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + + q = a.QNorm.Forward(q, cfg.RMSNormEps) + k = a.KNorm.Forward(k, cfg.RMSNormEps) + v = mlx.RMSNormFn(v, nil, cfg.RMSNormEps) + + q = applyVisionRoPE(q, positions) + k = applyVisionRoPE(k, positions) + + q = mlx.Transpose(q, 0, 2, 1, 3) + k = mlx.Transpose(k, 0, 2, 1, 3) + v = mlx.Transpose(v, 0, 2, 1, 3) + + vb := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, int(B), int(L)), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{L}, + } + out := nn.ScaledDotProductAttention(vb, q, 1.0, nn.WithKV(k, v, vb.SeqQueryLens)) + out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) + return a.OProj.Forward(out) +} + +func applyVisionRoPE(x *mlx.Array, positions visionPositionArrays) *mlx.Array { + rotated := rotateVisionHalf(x) + return mlx.Add(mlx.Mul(x, positions.RopeCos.AsType(x.DType())), mlx.Mul(rotated, positions.RopeSin.AsType(x.DType()))) +} + +func rotateVisionHalf(x *mlx.Array) *mlx.Array { + dims := x.Dims() + B, L, H, D := int32(dims[0]), int32(dims[1]), int32(dims[2]), int32(dims[3]) + ndim := int32(2) + channels := 2 * (D / (2 * ndim)) + if channels == 0 { + return x + } + parts := make([]*mlx.Array, 0, ndim+1) + for dim := range ndim { + start := dim * channels + mid := start + channels/2 + end := start + channels + x1 := mlx.SliceStartStop(x, []int32{0, 0, 0, start}, []int32{B, L, H, mid}) + x2 := mlx.SliceStartStop(x, []int32{0, 0, 0, mid}, []int32{B, L, H, end}) + parts = append(parts, mlx.Concatenate([]*mlx.Array{mlx.Neg(x2), x1}, -1)) + } + if tail := channels * ndim; tail < D { + parts = append(parts, mlx.SliceStartStop(x, []int32{0, 0, 0, tail}, []int32{B, L, H, D})) + } + return mlx.Concatenate(parts, -1) +} + +func (m *VisionMLP) Forward(x *mlx.Array, cfg *VisionConfig) *mlx.Array { + gate := m.GateProj.Forward(x) + up := m.UpProj.Forward(x) + return m.DownProj.Forward(mlx.GeGLU(gate, up)) +} + +func (m *VisionModel) pool(hidden *mlx.Array, patchH, patchW int32) *mlx.Array { + k := m.Config.PoolingKernelSize + B := int32(hidden.Dim(0)) + D := int32(hidden.Dim(2)) + if k <= 1 { + return mlx.MulScalar(hidden, float32(math.Sqrt(float64(m.Config.HiddenSize)))) + } + outH := patchH / k + outW := patchW / k + x := mlx.Reshape(hidden, B, patchH, patchW, D) + x = mlx.Reshape(x, B, outH, k, outW, k, D) + x = mlx.Mean(x, 4, false) + x = mlx.Mean(x, 2, false) + x = mlx.Reshape(x, B, outH*outW, D) + return mlx.MulScalar(x, float32(math.Sqrt(float64(m.Config.HiddenSize)))) +} + +var _ base.MediaModel = (*Model)(nil) diff --git a/x/models/gemma4/vision_test.go b/x/models/gemma4/vision_test.go new file mode 100644 index 00000000000..f269fb750ad --- /dev/null +++ b/x/models/gemma4/vision_test.go @@ -0,0 +1,298 @@ +package gemma4 + +import ( + "bytes" + "image" + "image/color" + "image/png" + "math" + "slices" + "strings" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/model/base" +) + +func testVisionConfig(t *testing.T, outputLength int32) *VisionConfig { + t.Helper() + cfg, err := parseVisionConfig([]byte(`{"vision_config":{"default_output_length":1}}`)) + if err != nil { + t.Fatal(err) + } + cfg.DefaultOutputLength = outputLength + if err := validateVisionConfig(cfg); err != nil { + t.Fatal(err) + } + return cfg +} + +func TestParseVisionConfigDefaults(t *testing.T) { + cfg, err := parseVisionConfig([]byte(`{"vision_config":{}}`)) + if err != nil { + t.Fatalf("parseVisionConfig() error = %v", err) + } + if cfg == nil { + t.Fatal("parseVisionConfig() = nil, want config") + } + if cfg.HiddenSize != 768 || cfg.IntermediateSize != 3072 || cfg.NumHiddenLayers != 16 { + t.Fatalf("unexpected default dimensions: hidden=%d intermediate=%d layers=%d", cfg.HiddenSize, cfg.IntermediateSize, cfg.NumHiddenLayers) + } + if cfg.PatchSize != 16 || cfg.PoolingKernelSize != 3 || cfg.DefaultOutputLength != 280 { + t.Fatalf("unexpected image defaults: patch=%d pooling=%d output=%d", cfg.PatchSize, cfg.PoolingKernelSize, cfg.DefaultOutputLength) + } + if cfg.RopeParameters.RopeTheta != 100 { + t.Fatalf("vision rope theta = %v, want 100", cfg.RopeParameters.RopeTheta) + } +} + +func TestParseVisionConfigRejectsUnsafeDimensions(t *testing.T) { + tests := []struct { + name string + json string + want string + }{ + {"negative hidden", `{"vision_config":{"hidden_size":-1}}`, "hidden_size"}, + {"too many layers", `{"vision_config":{"num_hidden_layers":513}}`, "num_hidden_layers"}, + {"attention width mismatch", `{"vision_config":{"head_dim":65}}`, "attention width"}, + {"invalid kv heads", `{"vision_config":{"num_key_value_heads":7}}`, "attention heads"}, + {"unbounded soft tokens", `{"vision_config":{"default_output_length":16385}}`, "default_output_length"}, + {"unbounded pooled patch", `{"vision_config":{"patch_size":16384,"pooling_kernel_size":2}}`, "resize budget"}, + {"undersized position table", `{"vision_config":{"position_embedding_size":100}}`, "position table"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := parseVisionConfig([]byte(tt.json)) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("parseVisionConfig() error = %v, want %q", err, tt.want) + } + }) + } +} + +func TestVisionStandardizationTensorRequirements(t *testing.T) { + cfg := testVisionConfig(t, 1) + if err := validateVisionStandardizationTensors(cfg, false, false); err != nil { + t.Fatalf("non-standardized config error = %v", err) + } + cfg.Standardize = true + for _, tt := range []struct { + name string + hasBias, hasScale bool + wantErr bool + }{ + {"complete", true, true, false}, + {"missing bias", false, true, true}, + {"missing scale", true, false, true}, + {"missing pair", false, false, true}, + } { + t.Run(tt.name, func(t *testing.T) { + err := validateVisionStandardizationTensors(cfg, tt.hasBias, tt.hasScale) + if (err != nil) != tt.wantErr { + t.Fatalf("validateVisionStandardizationTensors() error = %v, wantErr %t", err, tt.wantErr) + } + }) + } +} + +func TestGemma4ImageBounds(t *testing.T) { + if err := validateGemma4ImageDataSize(maxGemma4ImageBytes); err != nil { + t.Fatalf("exact byte limit error = %v", err) + } + if err := validateGemma4ImageDataSize(maxGemma4ImageBytes + 1); err == nil { + t.Fatal("over byte limit error = nil") + } + if err := validateGemma4ImageDimensions(maxGemma4ImageDimension, 1); err != nil { + t.Fatalf("exact dimension limit error = %v", err) + } + if err := validateGemma4ImageDimensions(maxGemma4ImageDimension+1, 1); err == nil { + t.Fatal("over dimension limit error = nil") + } + if err := validateGemma4ImageDimensions(8192, 8192); err != nil { + t.Fatalf("exact pixel limit error = %v", err) + } + if err := validateGemma4ImageDimensions(8192, 8193); err == nil { + t.Fatal("over pixel limit error = nil") + } +} + +func TestGemma4ResizeDimensions(t *testing.T) { + gotW, gotH, err := gemma4ResizeDimensions(1024, 768, 16, 280*9, 3) + if err != nil { + t.Fatalf("gemma4ResizeDimensions() error = %v", err) + } + if gotW != 912 || gotH != 672 { + t.Fatalf("gemma4ResizeDimensions() = %dx%d, want 912x672", gotW, gotH) + } + if gotW%48 != 0 || gotH%48 != 0 { + t.Fatalf("resize dimensions must be multiples of pooled patch size, got %dx%d", gotW, gotH) + } +} + +func TestGemma4ResizeDimensionsRejectsUnsafeArithmetic(t *testing.T) { + if _, _, err := gemma4ResizeDimensions(1024, 1024, 16, maxGemma4ResizePixels/(16*16), 4); err != nil { + t.Fatalf("exact resize work limit error = %v", err) + } + if _, _, err := gemma4ResizeDimensions(1, 1, 16, maxGemma4ResizePixels/(16*16)+1, 4); err == nil { + t.Fatal("over resize work limit error = nil") + } + if _, _, err := gemma4ResizeDimensions(1, 1, math.MaxInt, math.MaxInt, math.MaxInt); err == nil { + t.Fatal("overflowing resize parameters error = nil") + } +} + +func TestImageToCHWFloat32UsesBoundsAndChannelOrder(t *testing.T) { + img := image.NewRGBA(image.Rect(10, 20, 12, 22)) + img.SetRGBA(10, 20, color.RGBA{R: 255, A: 255}) + img.SetRGBA(11, 20, color.RGBA{G: 128, A: 255}) + img.SetRGBA(10, 21, color.RGBA{B: 64, A: 255}) + img.SetRGBA(11, 21, color.RGBA{R: 32, G: 16, B: 8, A: 255}) + + got, err := imageToCHWFloat32(img) + if err != nil { + t.Fatalf("imageToCHWFloat32() error = %v", err) + } + want := []float32{ + 1, 0, 0, 32.0 / 255.0, + 0, 128.0 / 255.0, 0, 16.0 / 255.0, + 0, 0, 64.0 / 255.0, 8.0 / 255.0, + } + if len(got) != len(want) { + t.Fatalf("len(imageToCHWFloat32()) = %d, want %d", len(got), len(want)) + } + for i := range want { + if math.Abs(float64(got[i]-want[i])) > 1e-6 { + t.Fatalf("pixel[%d] = %f, want %f", i, got[i], want[i]) + } + } +} + +func TestPreprocessGemma4ImageSoftTokenBudget(t *testing.T) { + src := image.NewRGBA(image.Rect(0, 0, 8, 4)) + src.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) + + var buf bytes.Buffer + if err := png.Encode(&buf, src); err != nil { + t.Fatalf("png.Encode() error = %v", err) + } + + img, err := preprocessGemma4Image(buf.Bytes(), testVisionConfig(t, 1), 1) + if err != nil { + t.Fatalf("preprocessGemma4Image() error = %v", err) + } + if img.Width != 48 || img.Height != 48 { + t.Fatalf("preprocessed dimensions = %dx%d, want 48x48", img.Width, img.Height) + } + if img.PatchWidth != 3 || img.PatchHeight != 3 || img.SoftTokens != 1 { + t.Fatalf("patch layout = %dx%d soft=%d, want 3x3 soft=1", img.PatchWidth, img.PatchHeight, img.SoftTokens) + } + if len(img.Pixels) != 3*48*48 { + t.Fatalf("pixel count = %d, want %d", len(img.Pixels), 3*48*48) + } +} + +func TestPrepareMediaPreservesOrderedImageItems(t *testing.T) { + pngData := func(c color.RGBA) []byte { + t.Helper() + img := image.NewRGBA(image.Rect(0, 0, 8, 4)) + img.SetRGBA(0, 0, c) + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + t.Fatalf("png.Encode() error = %v", err) + } + return buf.Bytes() + } + + m := &Model{ + TextConfig: &TextConfig{ + ImageTokenIDValue: 10, + BOITokenIDValue: 11, + EOITokenIDValue: 12, + VisionSoftTokens: 1, + }, + VisionConfig: testVisionConfig(t, 1), + Vision: &VisionModel{}, + EmbedVision: &MultimodalEmbedder{}, + } + segments := []base.Segment{ + {Tokens: []int32{1, 2}}, + {Kind: "image", Data: pngData(color.RGBA{R: 255, A: 255})}, + {Tokens: []int32{3}}, + {Kind: "image", Data: pngData(color.RGBA{B: 255, A: 255})}, + } + got, err := m.PrepareMedia(segments) + if err != nil { + t.Fatalf("PrepareMedia() error = %v", err) + } + wantTokens := []int32{1, 2, 11, 10, 12, 3, 11, 10, 12} + if !slices.Equal(got.Tokens, wantTokens) { + t.Fatalf("PrepareMedia().Tokens = %v, want %v", got.Tokens, wantTokens) + } + if len(got.Items) != 2 { + t.Fatalf("len(PrepareMedia().Items) = %d, want 2", len(got.Items)) + } + for i, want := range []struct { + range_ [2]int + source int + }{{[2]int{2, 5}, 1}, {[2]int{6, 9}, 3}} { + item := got.Items[i] + if item.Range != want.range_ || item.Source != want.source { + t.Fatalf("item %d range/source = %v/%d, want %v/%d", i, item.Range, item.Source, want.range_, want.source) + } + if item.Causal { + t.Fatalf("item %d is causal, want whole bidirectional image expansion", i) + } + if !slices.Equal(item.Dims, []int{1, 3, 48, 48}) || len(item.MediaData) != 3*48*48 { + t.Fatalf("item %d media shape/data = %v/%d, want [1 3 48 48]/%d", i, item.Dims, len(item.MediaData), 3*48*48) + } + payload := item.Opaque.(gemma4MediaPayload) + if payload.ImageStart != 1 || payload.ImageEnd != 2 || payload.Image.Pixels != nil { + t.Fatalf("item %d payload = start %d end %d pixels %d", i, payload.ImageStart, payload.ImageEnd, len(payload.Image.Pixels)) + } + media := batch.MediaItem{Pos: item.Range[0], Opaque: item.Opaque} + start, end := gemma4ImageRun(media) + if start != item.Range[0]+1 || end != item.Range[1]-1 { + t.Fatalf("item %d scatter/PLE run = [%d,%d), want current image span [%d,%d)", i, start, end, item.Range[0]+1, item.Range[1]-1) + } + } +} + +func TestPrepareMediaSequentialRequestsAreIsolated(t *testing.T) { + pngData := func(c color.RGBA) []byte { + t.Helper() + img := image.NewRGBA(image.Rect(0, 0, 8, 4)) + img.SetRGBA(0, 0, c) + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + t.Fatal(err) + } + return buf.Bytes() + } + m := &Model{ + TextConfig: &TextConfig{ImageTokenIDValue: 10, BOITokenIDValue: 11, EOITokenIDValue: 12, VisionSoftTokens: 1}, + VisionConfig: testVisionConfig(t, 1), + Vision: &VisionModel{}, + EmbedVision: &MultimodalEmbedder{}, + } + first, err := m.PrepareMedia([]base.Segment{{Tokens: []int32{1}}, {Kind: "image", Data: pngData(color.RGBA{R: 255, A: 255})}}) + if err != nil { + t.Fatal(err) + } + second, err := m.PrepareMedia([]base.Segment{{Tokens: []int32{2, 3}}, {Kind: "image", Data: pngData(color.RGBA{B: 255, A: 255})}}) + if err != nil { + t.Fatal(err) + } + if !slices.Equal(first.Tokens, []int32{1, 11, 10, 12}) || !slices.Equal(second.Tokens, []int32{2, 3, 11, 10, 12}) { + t.Fatalf("sequential tokens = %v / %v", first.Tokens, second.Tokens) + } + if len(first.Items) != 1 || len(second.Items) != 1 || &first.Items[0].MediaData[0] == &second.Items[0].MediaData[0] { + t.Fatal("sequential prepares reused request/item storage") + } + firstTokens := slices.Clone(first.Tokens) + firstPixels := slices.Clone(first.Items[0].MediaData) + second.Tokens[0] = 99 + second.Items[0].MediaData[0] = 0.5 + if !slices.Equal(first.Tokens, firstTokens) || !slices.Equal(first.Items[0].MediaData, firstPixels) { + t.Fatal("mutating the later prepared request changed the earlier request") + } +} From 84684e4604acd6340c254cfbdaf37b32d4ac1131 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 21:02:58 +0000 Subject: [PATCH 42/58] gemma4: expose MLX vision capability Advertise Gemma 4 vision now that the MLX model implements the official media contract, while retaining audio suppression until its runtime lands. Verify the centralized import policy preserves Gemma media tensors at source precision. Co-authored-by: Codex --- server/images.go | 4 ---- server/images_test.go | 9 ++++++--- server/model_list_cache.go | 7 ++++++- x/create/gemma4_test.go | 40 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 52 insertions(+), 8 deletions(-) diff --git a/server/images.go b/server/images.go index e4b81d564c6..26afb21118c 100644 --- a/server/images.go +++ b/server/images.go @@ -480,10 +480,6 @@ func (m *Model) filterUnsupportedCapabilities(capabilities []model.Capability, m } func suppressVisionCapability(m *Model) bool { - if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" { - return true - } - // The current MLX Nemotron path is text-only. Do not advertise vision for // safetensors manifests until the runner can load and serve that modality. return isNemotron3NanoSafetensors(m) diff --git a/server/images_test.go b/server/images_test.go index 677a59b9784..eec8642ebb8 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -545,7 +545,7 @@ func TestModelCapabilities(t *testing.T) { expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking}, }, { - name: "gemma4 small safetensors suppresses vision and audio", + name: "gemma4 small safetensors exposes vision and suppresses audio", model: Model{ Config: model.ConfigV2{ ModelFormat: "safetensors", @@ -554,9 +554,10 @@ func TestModelCapabilities(t *testing.T) { }, Template: chatTemplate, }, + expectedCaps: []model.Capability{model.CapabilityVision}, }, { - name: "gemma4 large safetensors suppresses vision and audio", + name: "gemma4 large safetensors exposes vision and suppresses audio", model: Model{ Config: model.ConfigV2{ ModelFormat: "safetensors", @@ -565,9 +566,10 @@ func TestModelCapabilities(t *testing.T) { }, Template: chatTemplate, }, + expectedCaps: []model.Capability{model.CapabilityVision}, }, { - name: "default gemma4 safetensors suppresses vision and audio", + name: "default gemma4 safetensors exposes vision and suppresses audio", model: Model{ Config: model.ConfigV2{ ModelFormat: "safetensors", @@ -576,6 +578,7 @@ func TestModelCapabilities(t *testing.T) { }, Template: chatTemplate, }, + expectedCaps: []model.Capability{model.CapabilityVision}, }, } diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 8b3a3346112..9355e8fc067 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -397,7 +397,12 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum } func filterUnsupportedModelListCapabilities(capabilities []model.Capability, cfg model.ConfigV2) []model.Capability { - if cfg.ModelFormat == "safetensors" && (isGemma4Renderer(cfg.Renderer) || isNemotron3NanoSafetensorsConfig(cfg)) { + if cfg.ModelFormat == "safetensors" && isGemma4Renderer(cfg.Renderer) { + capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { + return c == model.CapabilityAudio + }) + } + if isNemotron3NanoSafetensorsConfig(cfg) { capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { return c == model.CapabilityVision || c == model.CapabilityAudio }) diff --git a/x/create/gemma4_test.go b/x/create/gemma4_test.go index 4edd2656a07..edaf89fd7ff 100644 --- a/x/create/gemma4_test.go +++ b/x/create/gemma4_test.go @@ -181,6 +181,46 @@ func TestGemma4QuantizationType(t *testing.T) { } } +func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { + policy := gemma4ImportTransform{numLayers: 2} + inv := newInventory(sourceModelConfig{}, map[string]string{ + "model.embed_tokens.weight": "BF16", + "model.layers.0.self_attn.q_proj.weight": "BF16", + "model.vision_tower.patch_embedder.input_proj.weight": "BF16", + "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight": "BF16", + "model.embed_vision.embedding_projection.weight": "BF16", + "model.audio_tower.subsample_conv_projection.input_proj_linear.weight": "BF16", + "model.embed_audio.embedding_projection.weight": "BF16", + }) + + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "nvfp4"}, policy) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + got := make(map[string]TensorSpec) + for _, spec := range specs { + for _, tensor := range spec.Tensors { + got[tensor.Name] = tensor + } + } + for _, name := range []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", + "model.embed_vision.embedding_projection.weight", + "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", + "model.embed_audio.embedding_projection.weight", + } { + tensor, ok := got[name] + if !ok { + t.Fatalf("%s missing from plan; got %v", name, specNames(specs)) + } + if tensor.Quantize != "" { + t.Fatalf("%s Quantize = %q, want source precision", name, tensor.Quantize) + } + } +} + func TestUseMoreBits(t *testing.T) { // 30 layers: first 1/8 = layers 0-2, last 1/8 = layers 27-29 // In between: every 3rd from offset (i - n/8) % 3 == 2 From 8a71076e929ffef6bd8cd31280cc4b411f9f3f13 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 21:12:06 +0000 Subject: [PATCH 43/58] gemma4: gate MLX vision capability on tensors Require both the Gemma 4 vision tower and projector in source inventories and local manifests before advertising MLX vision, while preserving current centralized capability detection and unrelated model-family suppression. Co-authored-by: Codex --- server/images.go | 65 +++++++++++++++ server/images_test.go | 118 ++++++++++++++++++++++++++- server/model_inference_cache.go | 1 + server/model_inference_cache_test.go | 56 +++++++++++++ server/model_list_cache.go | 5 ++ server/model_list_cache_test.go | 45 ++++++++++ x/create/client/create.go | 46 ++++++++++- x/create/client/create_test.go | 99 +++++++++++++++++++++- 8 files changed, 427 insertions(+), 8 deletions(-) diff --git a/server/images.go b/server/images.go index 26afb21118c..faaa48cebd8 100644 --- a/server/images.go +++ b/server/images.go @@ -73,6 +73,7 @@ type Model struct { PreferChatTemplate bool // set when GGUF chat_template should take precedence over Go TEMPLATE AdapterPaths []string ProjectorPaths []string + TensorLayerNames []string System string License []string Digest string @@ -465,6 +466,12 @@ func (m *Model) modelFamilyCapabilities(capabilities []model.Capability) []model } func (m *Model) filterUnsupportedCapabilities(capabilities []model.Capability, modelArch string) []model.Capability { + if suppressGemma4SafetensorsVisionCapability(m) { + capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { + return c == model.CapabilityVision + }) + } + if suppressAudioCapability(m, modelArch) { capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { return c == model.CapabilityAudio @@ -485,6 +492,14 @@ func suppressVisionCapability(m *Model) bool { return isNemotron3NanoSafetensors(m) } +func suppressGemma4SafetensorsVisionCapability(m *Model) bool { + if m == nil || !isLocalGemma4SafetensorsConfig(m.Config) { + return false + } + + return !hasGemma4VisionTensorLayerNames(m.TensorLayerNames) +} + func suppressAudioCapability(m *Model, arch string) bool { if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" { return true @@ -518,6 +533,50 @@ func isNemotron3NanoSafetensorsConfig(cfg model.ConfigV2) bool { slices.Contains(cfg.ModelFamilies, "nemotron_h_omni")) } +func isLocalGemma4SafetensorsConfig(cfg model.ConfigV2) bool { + return cfg.ModelFormat == "safetensors" && + isGemma4Renderer(cfg.Renderer) && + cfg.RemoteHost == "" && + cfg.RemoteModel == "" +} + +var ( + gemma4VisionPatchTensorLayerNames = []string{ + "vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.input_proj.weight", + } + gemma4VisionProjectorTensorLayerNames = []string{ + "embed_vision.embedding_projection.weight", + "model.embed_vision.embedding_projection.weight", + } +) + +func hasGemma4VisionTensorLayers(layers []manifest.Layer) bool { + names := make([]string, 0, len(layers)) + for _, layer := range layers { + if layer.MediaType == manifest.MediaTypeImageTensor { + names = append(names, layer.Name) + } + } + + return hasGemma4VisionTensorLayerNames(names) +} + +func hasGemma4VisionTensorLayerNames(names []string) bool { + hasVisionTower := hasAnyTensorLayerName(names, gemma4VisionPatchTensorLayerNames...) + hasProjector := hasAnyTensorLayerName(names, gemma4VisionProjectorTensorLayerNames...) + return hasVisionTower && hasProjector +} + +func hasAnyTensorLayerName(names []string, candidates ...string) bool { + for _, candidate := range candidates { + if slices.Contains(names, candidate) { + return true + } + } + return false +} + func projectorHasAudio(f *gguf.File) bool { if f.KeyValue("has_audio_encoder").Bool() { return true @@ -583,6 +642,10 @@ func (m *Model) CheckCapabilities(want ...model.Capability) error { } } + if slices.Contains(errs, errCapabilityVision) && suppressGemma4SafetensorsVisionCapability(m) { + return fmt.Errorf("%w. Recreate or pull the model so it includes Gemma 4 vision tensor layers", err) + } + return err } @@ -741,6 +804,8 @@ func GetModel(name string) (*Model, error) { m.AdapterPaths = append(m.AdapterPaths, filename) case "application/vnd.ollama.image.projector": m.ProjectorPaths = append(m.ProjectorPaths, filename) + case manifest.MediaTypeImageTensor: + m.TensorLayerNames = append(m.TensorLayerNames, layer.Name) case "application/vnd.ollama.image.prompt", "application/vnd.ollama.image.template": m.HasGoTemplate = true diff --git a/server/images_test.go b/server/images_test.go index eec8642ebb8..aef2b710f6a 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "net/url" "os" + "slices" "strings" "testing" "time" @@ -552,7 +553,8 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererSmall, Capabilities: []string{"vision", "audio"}, }, - Template: chatTemplate, + TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -564,7 +566,8 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererLarge, Capabilities: []string{"vision", "audio"}, }, - Template: chatTemplate, + TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -576,7 +579,8 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererLegacy, Capabilities: []string{"vision", "audio"}, }, - Template: chatTemplate, + TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -618,6 +622,114 @@ func TestModelCapabilities(t *testing.T) { } } +func TestGemma4SafetensorsVisionCapabilityRequiresTensorLayers(t *testing.T) { + setTestHome(t, t.TempDir()) + + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "vision", "audio"}, + } + + createSafetensorsTestModel(t, "gemma4-text-only", cfg, nil) + m, err := GetModel("gemma4-text-only") + if err != nil { + t.Fatal(err) + } + caps := m.Capabilities() + if !slices.Contains(caps, model.CapabilityCompletion) { + t.Fatalf("capabilities = %v, want completion", caps) + } + if slices.Contains(caps, model.CapabilityVision) { + t.Fatalf("capabilities = %v, did not expect vision", caps) + } + if slices.Contains(caps, model.CapabilityAudio) { + t.Fatalf("capabilities = %v, did not expect audio", caps) + } + + err = m.CheckCapabilities(model.CapabilityVision) + if err == nil || !strings.Contains(err.Error(), "includes Gemma 4 vision tensor layers") { + t.Fatalf("CheckCapabilities(vision) error = %v, want Gemma 4 vision tensor hint", err) + } + + createSafetensorsTestModel(t, "gemma4-vision", cfg, gemma4VisionManifestLayers(t)) + m, err = GetModel("gemma4-vision") + if err != nil { + t.Fatal(err) + } + caps = m.Capabilities() + if !slices.Contains(caps, model.CapabilityVision) { + t.Fatalf("capabilities = %v, want vision", caps) + } + if slices.Contains(caps, model.CapabilityAudio) { + t.Fatalf("capabilities = %v, did not expect audio", caps) + } +} + +func TestGemma4VisionTensorLayerNamesRejectIncompleteAndNearMatches(t *testing.T) { + tests := []struct { + name string + names []string + want bool + }{ + {name: "missing"}, + {name: "tower only", names: []string{"model.vision_tower.patch_embedder.input_proj.weight"}}, + {name: "projector only", names: []string{"model.embed_vision.embedding_projection.weight"}}, + { + name: "near matches", + names: []string{ + "model.vision_tower.patch_embedder.input_proj.weight.extra", + "model.embed_visionary.embedding_projection.weight", + }, + }, + { + name: "complete model prefix", + names: []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.embed_vision.embedding_projection.weight", + }, + want: true, + }, + { + name: "complete bare prefix", + names: []string{ + "vision_tower.patch_embedder.input_proj.weight", + "embed_vision.embedding_projection.weight", + }, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := hasGemma4VisionTensorLayerNames(tt.names); got != tt.want { + t.Fatalf("hasGemma4VisionTensorLayerNames(%v) = %t, want %t", tt.names, got, tt.want) + } + }) + } +} + +func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { + t.Helper() + + data := []byte("fake-gemma4-vision-tensor") + digest := createTestBlob(t, data) + return []manifest.Layer{ + { + MediaType: manifest.MediaTypeImageTensor, + Digest: digest, + Size: int64(len(data)), + Name: "model.vision_tower.patch_embedder.input_proj.weight", + }, + { + MediaType: manifest.MediaTypeImageTensor, + Digest: digest, + Size: int64(len(data)), + Name: "model.embed_vision.embedding_projection.weight", + }, + } +} + func TestModelCheckCapabilities(t *testing.T) { // Create simple model file for tests that don't depend on GGUF content completionModelPath, _ := createBinFile(t, ggml.KV{ diff --git a/server/model_inference_cache.go b/server/model_inference_cache.go index a5edcbd44c3..7636e98640c 100644 --- a/server/model_inference_cache.go +++ b/server/model_inference_cache.go @@ -105,6 +105,7 @@ func cloneInferenceModel(src *Model) *Model { } dst.AdapterPaths = slices.Clone(src.AdapterPaths) dst.ProjectorPaths = slices.Clone(src.ProjectorPaths) + dst.TensorLayerNames = slices.Clone(src.TensorLayerNames) dst.License = slices.Clone(src.License) dst.Options = maps.Clone(src.Options) dst.Messages = slices.Clone(src.Messages) diff --git a/server/model_inference_cache_test.go b/server/model_inference_cache_test.go index fb53202e2a8..f841ef062ca 100644 --- a/server/model_inference_cache_test.go +++ b/server/model_inference_cache_test.go @@ -119,3 +119,59 @@ func TestInferenceModelCacheConcurrentMiss(t *testing.T) { t.Fatalf("load count = %d, want 1", got) } } + +func TestInferenceModelCacheGemma4VisionTensorCapabilities(t *testing.T) { + setTestHome(t, t.TempDir()) + + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "vision"}, + } + createSafetensorsTestModel(t, "gemma4-cache", cfg, gemma4VisionManifestLayers(t)) + + cache := newInferenceModelCache() + loadCount := 0 + cache.loadModel = func(name string) (*Model, error) { + loadCount++ + return GetModel(name) + } + + first, err := cache.Get("gemma4-cache") + if err != nil { + t.Fatal(err) + } + if !first.capabilitiesCached || !slices.Contains(first.Capabilities(), model.CapabilityVision) { + t.Fatalf("cold capabilities = %v, want cached vision", first.Capabilities()) + } + if len(first.TensorLayerNames) == 0 { + t.Fatal("cold model did not retain Gemma 4 tensor layer names") + } + first.TensorLayerNames[0] = "mutated" + + second, err := cache.Get("gemma4-cache") + if err != nil { + t.Fatal(err) + } + if loadCount != 1 { + t.Fatalf("cache hit load count = %d, want 1", loadCount) + } + if !slices.Contains(second.Capabilities(), model.CapabilityVision) { + t.Fatalf("cached capabilities = %v, want vision", second.Capabilities()) + } + if slices.Contains(second.TensorLayerNames, "mutated") { + t.Fatalf("cached tensor layer names were mutated: %v", second.TensorLayerNames) + } + + createSafetensorsTestModel(t, "gemma4-cache", cfg, nil) + third, err := cache.Get("gemma4-cache") + if err != nil { + t.Fatal(err) + } + if loadCount != 2 { + t.Fatalf("invalidated load count = %d, want 2", loadCount) + } + if slices.Contains(third.Capabilities(), model.CapabilityVision) { + t.Fatalf("refreshed capabilities = %v, did not expect vision", third.Capabilities()) + } +} diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 9355e8fc067..46925719669 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -391,6 +391,11 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityVision) } + if isLocalGemma4SafetensorsConfig(cfg) && !hasGemma4VisionTensorLayers(mf.Layers) { + summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { + return c == model.CapabilityVision + }) + } summary.Capabilities = filterUnsupportedModelListCapabilities(summary.Capabilities, cfg) return summary, nil diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index f9fb960cd4e..465d527dfdf 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -133,6 +133,51 @@ func TestModelListCacheRefreshUpdatesEntry(t *testing.T) { } } +func TestModelListSummaryGemma4SafetensorsVisionRequiresTensorLayers(t *testing.T) { + setTestHome(t, t.TempDir()) + + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "vision", "audio"}, + } + + createSafetensorsTestModel(t, "list-gemma4-text-only", cfg, nil) + mf, err := manifest.ParseNamedManifest(model.ParseName("list-gemma4-text-only")) + if err != nil { + t.Fatal(err) + } + summary, err := buildModelListSummary(model.ParseName("list-gemma4-text-only"), mf) + if err != nil { + t.Fatal(err) + } + if !slices.Contains(summary.Capabilities, model.CapabilityCompletion) { + t.Fatalf("capabilities = %v, want completion", summary.Capabilities) + } + if slices.Contains(summary.Capabilities, model.CapabilityVision) { + t.Fatalf("capabilities = %v, did not expect vision", summary.Capabilities) + } + if slices.Contains(summary.Capabilities, model.CapabilityAudio) { + t.Fatalf("capabilities = %v, did not expect audio", summary.Capabilities) + } + + createSafetensorsTestModel(t, "list-gemma4-vision", cfg, gemma4VisionManifestLayers(t)) + mf, err = manifest.ParseNamedManifest(model.ParseName("list-gemma4-vision")) + if err != nil { + t.Fatal(err) + } + summary, err = buildModelListSummary(model.ParseName("list-gemma4-vision"), mf) + if err != nil { + t.Fatal(err) + } + if !slices.Contains(summary.Capabilities, model.CapabilityVision) { + t.Fatalf("capabilities = %v, want vision", summary.Capabilities) + } + if slices.Contains(summary.Capabilities, model.CapabilityAudio) { + t.Fatalf("capabilities = %v, did not expect audio", summary.Capabilities) + } +} + func TestModelListCacheMutationHooks(t *testing.T) { gin.SetMode(gin.TestMode) setTestHome(t, t.TempDir()) diff --git a/x/create/client/create.go b/x/create/client/create.go index e070e50b5d5..c2cd6da8c5a 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -560,8 +560,13 @@ func detectCapabilities(modelDir string) modelCapabilities { _ = json.Unmarshal(data, &cfg) } + vision := cfg.VisionConfig != nil || cfg.HasVision + if vision && isGemma4ModelConfig(cfg.Architectures, cfg.ModelType) { + vision = gemma4ModelDirHasVisionTensors(modelDir) + } + return modelCapabilities{ - vision: cfg.VisionConfig != nil || cfg.HasVision, + vision: vision, audio: cfg.AudioConfig != nil || cfg.SoundConfig != nil, thinking: chatTemplateHasThinkingSupport(readChatTemplate(modelDir)) || alwaysSupportsThinking(cfg.Architectures, cfg.ModelType), @@ -600,6 +605,45 @@ func isApertus1p0ModelDir(modelDir, parserName string) bool { return false } +func isGemma4ModelConfig(architectures []string, modelType string) bool { + for _, arch := range architectures { + if isGemma4ModelIdentifier(arch) { + return true + } + } + return isGemma4ModelIdentifier(modelType) +} + +func isGemma4ModelIdentifier(value string) bool { + switch strings.ToLower(value) { + case "gemma4", "gemma4_unified", + "gemma4forcausallm", "gemma4forconditionalgeneration", + "gemma4unifiedforcausallm", "gemma4unifiedforconditionalgeneration": + return true + default: + return false + } +} + +func gemma4ModelDirHasVisionTensors(modelDir string) bool { + inv, err := create.ReadInventory(modelDir) + if err != nil { + return false + } + + return hasAnySourceTensor(inv, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") && + hasAnySourceTensor(inv, "embed_vision.embedding_projection.weight", "model.embed_vision.embedding_projection.weight") +} + +func hasAnySourceTensor(inv create.Inventory, names ...string) bool { + for _, name := range names { + if inv.Has(name) { + return true + } + } + return false +} + // readChatTemplate returns the model's chat template, preferring the // chat_template field of tokenizer_config.json and falling back to a standalone // chat_template.jinja. It returns "" when neither is present. diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 5be574acc65..09b898140e5 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -2,6 +2,7 @@ package client import ( "encoding/json" + "io" "os" "path/filepath" "slices" @@ -12,6 +13,7 @@ import ( "github.com/ollama/ollama/parser" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/x/create" + "github.com/ollama/ollama/x/safetensors" ) func TestModelfileConfig(t *testing.T) { @@ -422,14 +424,14 @@ func TestInferSafetensorsCapabilities(t *testing.T) { want: []string{"completion", "vision", "thinking"}, }, { - name: "model with audio config", + name: "gemma4 with audio config and missing vision tensors", configJSON: `{ "architectures": ["Gemma4ForConditionalGeneration"], "model_type": "gemma4", "vision_config": {"hidden_size": 1024}, "audio_config": {"num_mel_bins": 128} }`, - want: []string{"completion", "vision", "audio"}, + want: []string{"completion", "audio"}, }, { name: "model with audio but no vision", @@ -545,6 +547,90 @@ func TestApertusMetadataInference(t *testing.T) { } } +func TestInferSafetensorsCapabilitiesGemma4VisionRequiresTensors(t *testing.T) { + configJSON := `{ + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + "vision_config": {"hidden_size": 1024} + }` + + tests := []struct { + name string + tensors []string + want []string + }{ + {name: "missing vision tensors", want: []string{"completion"}}, + { + name: "patch only", + tensors: []string{"model.vision_tower.patch_embedder.input_proj.weight"}, + want: []string{"completion"}, + }, + { + name: "projector only", + tensors: []string{"model.embed_vision.embedding_projection.weight"}, + want: []string{"completion"}, + }, + { + name: "near-match names", + tensors: []string{ + "model.vision_tower.patch_embedder.input_proj.weight.extra", + "model.embed_visionary.embedding_projection.weight", + }, + want: []string{"completion"}, + }, + { + name: "vision tower and projector", + tensors: []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.embed_vision.embedding_projection.weight", + }, + want: []string{"completion", "vision"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + if len(tt.tensors) > 0 { + writeClientSafetensors(t, dir, tt.tensors...) + } + + if got := inferSafetensorsCapabilities(dir, ""); !slices.Equal(got, tt.want) { + t.Fatalf("inferSafetensorsCapabilities() = %#v, want %#v", got, tt.want) + } + }) + } +} + +func TestGemma4ModelConfigRejectsNearMatches(t *testing.T) { + if isGemma4ModelConfig([]string{"NotGemma4ForConditionalGeneration"}, "gemma4-next") { + t.Fatal("near-match Gemma 4 identifiers classified as Gemma 4") + } + if !isGemma4ModelConfig([]string{"Gemma4ForConditionalGeneration"}, "") { + t.Fatal("released Gemma 4 architecture not classified") + } +} + +func writeClientSafetensors(t *testing.T, dir string, names ...string) { + t.Helper() + + tensors := make([]*safetensors.TensorData, 0, len(names)) + for _, name := range names { + tensors = append(tensors, safetensors.NewTensorDataFromBytes(name, "U8", []int32{1}, []byte{0})) + } + + data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader(tensors)) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "model.safetensors"), data, 0o644); err != nil { + t.Fatal(err) + } +} + func TestCreateModelfileLayersIncludesParameters(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) @@ -806,8 +892,8 @@ func TestDetectCapabilities(t *testing.T) { want: modelCapabilities{vision: false, thinking: true}, }, { - name: "vision config", - configJSON: `{"architectures": ["Gemma4ForConditionalGeneration"], "vision_config": {}}`, + name: "non-gemma vision config", + configJSON: `{"architectures": ["SomeVisionModel"], "model_type": "other", "vision_config": {}}`, want: modelCapabilities{vision: true}, }, { @@ -825,6 +911,11 @@ func TestDetectCapabilities(t *testing.T) { configJSON: `{"architectures": ["LlamaForCausalLM"], "model_type": "llama"}`, want: modelCapabilities{}, }, + { + name: "apertus uses parser-level thinking, not config-level detection", + configJSON: `{"architectures": ["ApertusForCausalLM"], "model_type": "apertus"}`, + want: modelCapabilities{}, + }, { name: "invalid config json", configJSON: `not json`, From 05c55764fc939618a966d66992ba0013385e744d Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 22:54:16 +0000 Subject: [PATCH 44/58] gemma4: harden MLX vision integration Validate executable source, installed, server, and runtime vision contracts including tensor payload ranges; bound descriptor extraction; propagate live request cancellation; and preserve accepted Gemma, Glimmer, and Qwen behavior. Co-authored-by: Codex --- docs/third-party/mlx-vlm.md | 35 ++ server/images.go | 262 ++++++++-- server/images_test.go | 332 ++++++++++-- server/model_list_cache.go | 12 +- x/create/client/create.go | 20 +- x/create/client/create_test.go | 244 ++++++++- x/mlxrunner/grammar_test.go | 3 +- x/mlxrunner/media.go | 9 +- x/mlxrunner/model/base/media.go | 12 +- x/mlxrunner/pipeline.go | 50 +- x/mlxrunner/pipeline_test.go | 98 +++- x/mlxrunner/server.go | 48 +- x/mlxrunner/server_test.go | 101 ++++ x/models/gemma4/gemma4.go | 6 +- x/models/gemma4/metadata/vision.go | 669 ++++++++++++++++++++++++ x/models/gemma4/metadata/vision_test.go | 388 ++++++++++++++ x/models/gemma4/vision.go | 279 ++++++---- x/models/gemma4/vision_test.go | 194 ++++++- x/models/glimmer/media.go | 97 +++- x/models/glimmer/media_test.go | 143 ++++- x/models/qwen3_5/process_image.go | 133 ++++- x/models/qwen3_5/vision.go | 48 +- x/models/qwen3_5/vision_test.go | 246 ++++++++- x/models/qwen4_exp/media_test.go | 20 + x/models/qwen4_exp/qwen4_exp.go | 5 +- 25 files changed, 3144 insertions(+), 310 deletions(-) create mode 100644 docs/third-party/mlx-vlm.md create mode 100644 x/mlxrunner/server_test.go create mode 100644 x/models/gemma4/metadata/vision.go create mode 100644 x/models/gemma4/metadata/vision_test.go create mode 100644 x/models/qwen4_exp/media_test.go diff --git a/docs/third-party/mlx-vlm.md b/docs/third-party/mlx-vlm.md new file mode 100644 index 00000000000..1d717eac373 --- /dev/null +++ b/docs/third-party/mlx-vlm.md @@ -0,0 +1,35 @@ +# MLX-VLM + +Portions of `x/models/gemma4/vision.go` are adapted from the Gemma 4 +implementation in MLX-VLM: + +- Repository: https://github.com/Blaizzy/mlx-vlm +- Revision: `61990c9054f2bc7bb8f32541e3238b4a58fe64e5` +- Source paths: `mlx_vlm/models/gemma4/gemma4.py` and + `mlx_vlm/models/gemma4/vision.py` + +The adapted implementation is distributed under the following license. + +```text +MIT License + +Copyright © 2025 Prince Canuma + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +``` diff --git a/server/images.go b/server/images.go index faaa48cebd8..610c860687a 100644 --- a/server/images.go +++ b/server/images.go @@ -4,12 +4,14 @@ import ( "bytes" "context" "crypto/sha256" + "encoding/binary" "encoding/json" "errors" "fmt" "io" "log" "log/slog" + "math" "net" "net/http" "net/url" @@ -17,6 +19,7 @@ import ( "path/filepath" "runtime" "slices" + "sort" "strconv" "strings" "time" @@ -32,6 +35,8 @@ import ( "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/version" "github.com/ollama/ollama/x/mlxrunner/mlx" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" + "github.com/ollama/ollama/x/safetensors" "github.com/ollama/ollama/x/transfer" ) @@ -39,6 +44,12 @@ import ( // manifest yet. They become eligible for the normal mark-and-sweep pass later. const layerPruneGracePeriod = time.Hour +const ( + maxGemma4SafetensorsHeaderSize = 16 << 20 + maxGemma4VisionDescriptors = 8 << 10 + maxGemma4VisionDescriptorWork = 4 << 20 +) + var ( errCapabilities = errors.New("does not support") errCapabilityCompletion = errors.New("completion") @@ -62,23 +73,25 @@ type registryOptions struct { } type Model struct { - Name string `json:"name"` - Config model.ConfigV2 - ShortName string - ModelPath string - DraftPath string - ParentModel string - HasChatTemplate bool - HasGoTemplate bool - PreferChatTemplate bool // set when GGUF chat_template should take precedence over Go TEMPLATE - AdapterPaths []string - ProjectorPaths []string - TensorLayerNames []string - System string - License []string - Digest string - Options map[string]any - Messages []api.Message + Name string `json:"name"` + Config model.ConfigV2 + ShortName string + ModelPath string + DraftPath string + ParentModel string + HasChatTemplate bool + HasGoTemplate bool + PreferChatTemplate bool // set when GGUF chat_template should take precedence over Go TEMPLATE + AdapterPaths []string + ProjectorPaths []string + TensorLayerNames []string + Gemma4VisionConfig *gemma4metadata.ConfigFile `json:"-"` + Gemma4VisionTensors map[string]gemma4metadata.TensorDescriptor `json:"-"` + System string + License []string + Digest string + Options map[string]any + Messages []api.Message Template *template.Template @@ -497,7 +510,7 @@ func suppressGemma4SafetensorsVisionCapability(m *Model) bool { return false } - return !hasGemma4VisionTensorLayerNames(m.TensorLayerNames) + return m.Gemma4VisionConfig == nil || gemma4metadata.ValidateVisionInstalledInventory(*m.Gemma4VisionConfig, m.Gemma4VisionTensors) != nil } func suppressAudioCapability(m *Model, arch string) bool { @@ -540,41 +553,197 @@ func isLocalGemma4SafetensorsConfig(cfg model.ConfigV2) bool { cfg.RemoteModel == "" } -var ( - gemma4VisionPatchTensorLayerNames = []string{ - "vision_tower.patch_embedder.input_proj.weight", - "model.vision_tower.patch_embedder.input_proj.weight", - } - gemma4VisionProjectorTensorLayerNames = []string{ - "embed_vision.embedding_projection.weight", - "model.embed_vision.embedding_projection.weight", - } -) +func hasGemma4VisionTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifest.Layer) bool { + tensors, err := gemma4VisionTensorDescriptors(layers) + return err == nil && gemma4metadata.ValidateVisionInstalledInventory(cfg, tensors) == nil +} -func hasGemma4VisionTensorLayers(layers []manifest.Layer) bool { - names := make([]string, 0, len(layers)) +func gemma4VisionTensorDescriptors(layers []manifest.Layer) (map[string]gemma4metadata.TensorDescriptor, error) { + tensors := make(map[string]gemma4metadata.TensorDescriptor) + descriptorWork := 0 for _, layer := range layers { - if layer.MediaType == manifest.MediaTypeImageTensor { - names = append(names, layer.Name) + if layer.MediaType != manifest.MediaTypeImageTensor { + continue + } + if !strings.Contains(layer.Name, "vision_tower.") && !strings.Contains(layer.Name, "embed_vision.") { + continue + } + filename, err := manifest.BlobsPath(layer.Digest) + if err != nil { + return nil, err + } + ext, err := openGemma4TensorLayer(filename) + if err != nil { + return nil, fmt.Errorf("open tensor layer %s: %w", layer.Name, err) + } + names := ext.ListTensors() + if len(names) > maxGemma4VisionDescriptors-len(tensors) { + ext.Close() + return nil, fmt.Errorf("Gemma4 vision tensor inventory exceeds %d descriptors", maxGemma4VisionDescriptors) + } + for _, name := range names { + tensor, err := ext.GetTensor(name) + if err != nil { + ext.Close() + return nil, err + } + if _, exists := tensors[name]; exists { + ext.Close() + return nil, fmt.Errorf("duplicate tensor %s", name) + } + work := len(name) + len(tensor.Dtype) + len(tensor.Shape)*4 + if work > maxGemma4VisionDescriptorWork-descriptorWork { + ext.Close() + return nil, fmt.Errorf("Gemma4 vision tensor inventory exceeds descriptor work limit %d", maxGemma4VisionDescriptorWork) + } + descriptorWork += work + tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + } + if err := ext.Close(); err != nil { + return nil, err } } + return tensors, nil +} - return hasGemma4VisionTensorLayerNames(names) +func openGemma4TensorLayer(filename string) (extractor *safetensors.TensorExtractor, err error) { + f, err := os.Open(filename) + if err != nil { + return nil, err + } + var sizeBytes [8]byte + if _, err := io.ReadFull(f, sizeBytes[:]); err != nil { + f.Close() + return nil, fmt.Errorf("read safetensors header length: %w", err) + } + info, err := f.Stat() + if err != nil { + f.Close() + return nil, err + } + headerSize := binary.LittleEndian.Uint64(sizeBytes[:]) + if headerSize > maxGemma4SafetensorsHeaderSize { + f.Close() + return nil, fmt.Errorf("safetensors header too large: %d (limit %d)", headerSize, maxGemma4SafetensorsHeaderSize) + } + if info.Size() < 8 || headerSize > uint64(info.Size()-8) { + f.Close() + return nil, fmt.Errorf("truncated safetensors header: declared %d bytes in %d-byte file", headerSize, info.Size()) + } + header := make([]byte, int(headerSize)) + if _, err := io.ReadFull(f, header); err != nil { + f.Close() + return nil, fmt.Errorf("read safetensors header: %w", err) + } + if err := validateGemma4SafetensorsPayload(header, info.Size()-8-int64(headerSize)); err != nil { + f.Close() + return nil, err + } + if err := f.Close(); err != nil { + return nil, err + } + defer func() { + if recovered := recover(); recovered != nil { + extractor = nil + err = fmt.Errorf("invalid safetensors header: %v", recovered) + } + }() + return safetensors.OpenForExtraction(filename) } -func hasGemma4VisionTensorLayerNames(names []string) bool { - hasVisionTower := hasAnyTensorLayerName(names, gemma4VisionPatchTensorLayerNames...) - hasProjector := hasAnyTensorLayerName(names, gemma4VisionProjectorTensorLayerNames...) - return hasVisionTower && hasProjector +type gemma4SafetensorInfo struct { + Dtype string `json:"dtype"` + Shape []int64 `json:"shape"` + DataOffsets [2]int64 `json:"data_offsets"` } -func hasAnyTensorLayerName(names []string, candidates ...string) bool { - for _, candidate := range candidates { - if slices.Contains(names, candidate) { - return true +func validateGemma4SafetensorsPayload(header []byte, dataSize int64) error { + var entries map[string]gemma4SafetensorInfo + if err := json.Unmarshal(header, &entries); err != nil { + return fmt.Errorf("parse safetensors header: %w", err) + } + if len(entries) > maxGemma4VisionDescriptors+1 { + return fmt.Errorf("Gemma4 vision tensor inventory exceeds %d descriptors", maxGemma4VisionDescriptors) + } + type dataRange struct { + name string + start, end int64 + } + ranges := make([]dataRange, 0, len(entries)) + for name, entry := range entries { + if name == "__metadata__" { + continue + } + start, end := entry.DataOffsets[0], entry.DataOffsets[1] + if start < 0 || end < start { + return fmt.Errorf("invalid safetensors data range for %s: [%d,%d]", name, start, end) + } + if end > dataSize { + return fmt.Errorf("safetensors data range for %s ends at %d beyond data size %d", name, end, dataSize) } + expected, err := gemma4SafetensorByteSize(entry.Dtype, entry.Shape) + if err != nil { + return fmt.Errorf("invalid safetensors tensor %s: %w", name, err) + } + if end-start != expected { + return fmt.Errorf("safetensors tensor %s range length %d, want %d", name, end-start, expected) + } + ranges = append(ranges, dataRange{name: name, start: start, end: end}) } - return false + sort.Slice(ranges, func(i, j int) bool { + if ranges[i].start == ranges[j].start { + return ranges[i].end < ranges[j].end + } + return ranges[i].start < ranges[j].start + }) + if len(ranges) == 0 { + if dataSize != 0 { + return fmt.Errorf("safetensors data region has %d unclaimed bytes", dataSize) + } + return nil + } + if ranges[0].start != 0 { + return fmt.Errorf("safetensors data ranges start at %d, want 0", ranges[0].start) + } + for i := 1; i < len(ranges); i++ { + if ranges[i].start != ranges[i-1].end { + return fmt.Errorf("non-contiguous safetensors data ranges for %s and %s", ranges[i-1].name, ranges[i].name) + } + } + if ranges[len(ranges)-1].end != dataSize { + return fmt.Errorf("safetensors data ranges end at %d, want data size %d", ranges[len(ranges)-1].end, dataSize) + } + return nil +} + +func gemma4SafetensorByteSize(dtype string, shape []int64) (int64, error) { + var width int64 + switch strings.ToUpper(dtype) { + case "BOOL", "U8", "I8", "F8_E4M3", "F8_E4M3FN", "F8_E5M2", "F8_E5M2FNUZ": + width = 1 + case "U16", "I16", "F16", "BF16": + width = 2 + case "U32", "I32", "F32": + width = 4 + case "U64", "I64", "F64": + width = 8 + default: + return 0, fmt.Errorf("unsupported dtype %q", dtype) + } + elements := int64(1) + for _, dim := range shape { + if dim < 0 { + return 0, fmt.Errorf("negative shape dimension %d", dim) + } + if dim != 0 && elements > math.MaxInt64/dim { + return 0, fmt.Errorf("shape byte count overflows int64") + } + elements *= dim + } + if elements > math.MaxInt64/width { + return 0, fmt.Errorf("shape byte count overflows int64") + } + return elements * width, nil } func projectorHasAudio(f *gguf.File) bool { @@ -770,6 +939,15 @@ func GetModel(name string) (*Model, error) { return nil, err } } + if isLocalGemma4SafetensorsConfig(m.Config) { + var cfg gemma4metadata.ConfigFile + if err := mf.ReadConfigJSON("config.json", &cfg); err == nil { + m.Gemma4VisionConfig = &cfg + if tensors, err := gemma4VisionTensorDescriptors(mf.Layers); err == nil { + m.Gemma4VisionTensors = tensors + } + } + } modelHasPooling := false ggufChatTemplate := "" diff --git a/server/images_test.go b/server/images_test.go index aef2b710f6a..4a30fa38064 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -2,12 +2,18 @@ package server import ( "crypto/sha256" + "encoding/binary" + "encoding/json" "errors" "fmt" + "io" + "maps" + "math" "net/http" "net/http/httptest" "net/url" "os" + "path/filepath" "slices" "strings" "testing" @@ -18,6 +24,8 @@ import ( "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/template" "github.com/ollama/ollama/types/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" + "github.com/ollama/ollama/x/safetensors" ) func TestPruneLayersSkipsRecentOrphans(t *testing.T) { @@ -553,8 +561,10 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererSmall, Capabilities: []string{"vision", "audio"}, }, - TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, - Template: chatTemplate, + TensorLayerNames: gemma4VisionTensorNames(2), + Gemma4VisionConfig: gemma4VisionConfig(2), + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -566,8 +576,10 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererLarge, Capabilities: []string{"vision", "audio"}, }, - TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, - Template: chatTemplate, + TensorLayerNames: gemma4VisionTensorNames(2), + Gemma4VisionConfig: gemma4VisionConfig(2), + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -579,8 +591,10 @@ func TestModelCapabilities(t *testing.T) { Renderer: gemma4RendererLegacy, Capabilities: []string{"vision", "audio"}, }, - TensorLayerNames: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}, - Template: chatTemplate, + TensorLayerNames: gemma4VisionTensorNames(2), + Gemma4VisionConfig: gemma4VisionConfig(2), + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityVision}, }, @@ -664,9 +678,30 @@ func TestGemma4SafetensorsVisionCapabilityRequiresTensorLayers(t *testing.T) { if slices.Contains(caps, model.CapabilityAudio) { t.Fatalf("capabilities = %v, did not expect audio", caps) } + + partialLayers := slices.DeleteFunc(gemma4VisionManifestLayers(t), func(layer manifest.Layer) bool { + return layer.MediaType == manifest.MediaTypeImageTensor && + layer.Name != "model.vision_tower.patch_embedder.input_proj.weight" && + layer.Name != "model.embed_vision.embedding_projection.weight" + }) + createSafetensorsTestModel(t, "gemma4-partial-vision", cfg, partialLayers) + m, err = GetModel("gemma4-partial-vision") + if err != nil { + t.Fatal(err) + } + if slices.Contains(m.Capabilities(), model.CapabilityVision) { + t.Fatalf("partial vision capabilities = %v, did not expect vision", m.Capabilities()) + } } -func TestGemma4VisionTensorLayerNamesRejectIncompleteAndNearMatches(t *testing.T) { +func TestGemma4VisionTensorValidationRejectsIncompleteAndNearMatches(t *testing.T) { + bareTensorNames := func() []string { + names := gemma4VisionTensorNames(2) + for i := range names { + names[i] = strings.TrimPrefix(names[i], "model.") + } + return names + } tests := []struct { name string names []string @@ -682,28 +717,60 @@ func TestGemma4VisionTensorLayerNamesRejectIncompleteAndNearMatches(t *testing.T "model.embed_visionary.embedding_projection.weight", }, }, - { - name: "complete model prefix", - names: []string{ - "model.vision_tower.patch_embedder.input_proj.weight", - "model.embed_vision.embedding_projection.weight", - }, - want: true, - }, - { - name: "complete bare prefix", - names: []string{ - "vision_tower.patch_embedder.input_proj.weight", - "embed_vision.embedding_projection.weight", - }, - want: true, - }, + {name: "sentinels only", names: []string{"model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight"}}, + {name: "complete model prefix", names: gemma4VisionTensorNames(2), want: true}, + {name: "complete bare prefix", names: bareTensorNames(), want: true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := hasGemma4VisionTensorLayerNames(tt.names); got != tt.want { - t.Fatalf("hasGemma4VisionTensorLayerNames(%v) = %t, want %t", tt.names, got, tt.want) + got := gemma4metadata.ValidateVisionTensors(*gemma4VisionConfig(2), tt.names) == nil + if got != tt.want { + t.Fatalf("ValidateVisionTensors(%v) success = %t, want %t", tt.names, got, tt.want) + } + }) + } +} + +func TestGemma4InstalledVisionCapabilityRejectsMalformedDescriptors(t *testing.T) { + valid := Model{ + Config: model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererLarge, Capabilities: []string{"vision"}}, + Gemma4VisionConfig: gemma4VisionConfig(2), + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + } + if !slices.Contains(valid.Capabilities(), model.CapabilityVision) { + t.Fatal("valid installed descriptors did not expose vision") + } + malformed := valid + malformed.Gemma4VisionTensors = maps.Clone(valid.Gemma4VisionTensors) + name := "model.vision_tower.patch_embedder.input_proj.weight" + descriptor := malformed.Gemma4VisionTensors[name] + descriptor.Shape = []int32{4, 11} + malformed.Gemma4VisionTensors[name] = descriptor + if slices.Contains(malformed.Capabilities(), model.CapabilityVision) { + t.Fatal("complete-name installed inventory with wrong shape exposed vision") + } + packedAlias := valid + packedAlias.Gemma4VisionTensors = maps.Clone(valid.Gemma4VisionTensors) + delete(packedAlias.Gemma4VisionTensors, name) + packedAlias.Gemma4VisionTensors[strings.TrimSuffix(name, ".weight")+".weight_packed"] = gemma4metadata.TensorDescriptor{Dtype: "U8", Shape: []int32{4, 6}} + if slices.Contains(packedAlias.Capabilities(), model.CapabilityVision) { + t.Fatal("source-only packed alias exposed installed vision") + } +} + +func TestGemma4InstalledVisionCapabilityRejectsMissingOrZeroTextWidth(t *testing.T) { + for _, name := range []string{"missing", "zero"} { + t.Run(name, func(t *testing.T) { + cfg := gemma4VisionConfig(2) + cfg.TextConfig.HiddenSize = 0 + m := Model{ + Config: model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererLarge, Capabilities: []string{"vision"}}, + Gemma4VisionConfig: cfg, + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + } + if slices.Contains(m.Capabilities(), model.CapabilityVision) { + t.Fatal("non-executable text width exposed installed vision") } }) } @@ -712,22 +779,215 @@ func TestGemma4VisionTensorLayerNamesRejectIncompleteAndNearMatches(t *testing.T func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { t.Helper() - data := []byte("fake-gemma4-vision-tensor") - digest := createTestBlob(t, data) - return []manifest.Layer{ - { - MediaType: manifest.MediaTypeImageTensor, - Digest: digest, - Size: int64(len(data)), - Name: "model.vision_tower.patch_embedder.input_proj.weight", - }, - { + layers := make([]manifest.Layer, 0, len(gemma4VisionTensorNames(2))+1) + for name, descriptor := range testGemma4VisionTensorDescriptors(2) { + shape := make([]int64, len(descriptor.Shape)) + for i, dim := range descriptor.Shape { + shape[i] = int64(dim) + } + payloadSize, err := gemma4SafetensorByteSize(descriptor.Dtype, shape) + if err != nil { + t.Fatal(err) + } + data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{ + safetensors.NewTensorDataFromBytes(name, descriptor.Dtype, descriptor.Shape, make([]byte, int(payloadSize))), + })) + if err != nil { + t.Fatal(err) + } + digest := createTestBlob(t, data) + layers = append(layers, manifest.Layer{ MediaType: manifest.MediaTypeImageTensor, Digest: digest, Size: int64(len(data)), - Name: "model.embed_vision.embedding_projection.weight", + Name: name, + }) + } + config := []byte(`{"text_config":{"hidden_size":6},"vision_config":{"hidden_size":4,"intermediate_size":8,"num_hidden_layers":2,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":4,"default_output_length":1,"patch_size":2,"position_embedding_size":16,"pooling_kernel_size":1}}`) + configDigest := createTestBlob(t, config) + layers = append(layers, manifest.Layer{ + MediaType: "application/vnd.ollama.image.json", + Digest: configDigest, + Size: int64(len(config)), + Name: "config.json", + }) + if descriptors, err := gemma4VisionTensorDescriptors(layers); err != nil { + t.Fatalf("read Gemma4 descriptor fixture: %v", err) + } else if err := gemma4metadata.ValidateVisionInstalledInventory(*gemma4VisionConfig(2), descriptors); err != nil { + t.Fatalf("validate Gemma4 descriptor fixture: %v", err) + } + return layers +} + +func TestOpenGemma4TensorLayerRejectsUnsafeHeaders(t *testing.T) { + tests := []struct { + name string + data []byte + }{ + {name: "truncated length", data: []byte{1, 2, 3, 4}}, + {name: "truncated header", data: append(binary.LittleEndian.AppendUint64(nil, 32), []byte("{}")...)}, + {name: "oversized header", data: binary.LittleEndian.AppendUint64(nil, maxGemma4SafetensorsHeaderSize+1)}, + {name: "uint64 overflow header", data: binary.LittleEndian.AppendUint64(nil, ^uint64(0))}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "model.safetensors") + if err := os.WriteFile(path, tt.data, 0o644); err != nil { + t.Fatal(err) + } + if ext, err := openGemma4TensorLayer(path); err == nil { + ext.Close() + t.Fatal("unsafe safetensors header accepted") + } + }) + } +} + +func TestOpenGemma4TensorLayerRejectsInvalidPayloadRanges(t *testing.T) { + write := func(t *testing.T, dtype string, shape []int64, offsets [2]int64, payload int) string { + t.Helper() + header, err := json.Marshal(map[string]any{ + "tensor": map[string]any{"dtype": dtype, "shape": shape, "data_offsets": offsets}, + }) + if err != nil { + t.Fatal(err) + } + data := binary.LittleEndian.AppendUint64(nil, uint64(len(header))) + data = append(data, header...) + data = append(data, make([]byte, payload)...) + path := filepath.Join(t.TempDir(), "model.safetensors") + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + return path + } + tests := []struct { + name string + dtype string + shape []int64 + offsets [2]int64 + payload int + }{ + {name: "truncated range", dtype: "F32", shape: []int64{2}, offsets: [2]int64{0, 4}, payload: 4}, + {name: "overlong range", dtype: "F32", shape: []int64{1}, offsets: [2]int64{0, 8}, payload: 8}, + {name: "overlong payload", dtype: "F32", shape: []int64{1}, offsets: [2]int64{0, 4}, payload: 8}, + {name: "shape byte overflow", dtype: "F64", shape: []int64{math.MaxInt64, 2}, offsets: [2]int64{0, 0}}, + {name: "out of file", dtype: "U8", shape: []int64{2}, offsets: [2]int64{0, 2}, payload: 1}, + {name: "negative range", dtype: "U8", shape: []int64{1}, offsets: [2]int64{-1, 0}, payload: 1}, + {name: "reversed range", dtype: "U8", shape: []int64{1}, offsets: [2]int64{1, 0}, payload: 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if ext, err := openGemma4TensorLayer(write(t, tt.dtype, tt.shape, tt.offsets, tt.payload)); err == nil { + ext.Close() + t.Fatal("invalid safetensors payload range accepted") + } + }) + } +} + +func TestGemma4VisionTensorDescriptorsRejectsExcessiveInventory(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + tensors := make([]*safetensors.TensorData, 0, maxGemma4VisionDescriptors+1) + for i := 0; i <= maxGemma4VisionDescriptors; i++ { + tensors = append(tensors, safetensors.NewTensorDataFromBytes(fmt.Sprintf("tensor.%04d", i), "U8", []int32{1}, []byte{0})) + } + data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader(tensors)) + if err != nil { + t.Fatal(err) + } + digest := createTestBlob(t, data) + layers := []manifest.Layer{{ + MediaType: manifest.MediaTypeImageTensor, + Digest: digest, + Size: int64(len(data)), + Name: "model.vision_tower.synthetic.weight", + }} + if _, err := gemma4VisionTensorDescriptors(layers); err == nil || !strings.Contains(err.Error(), "descriptors") { + t.Fatalf("excessive inventory error = %v", err) + } +} + +func gemma4VisionConfig(layers int) *gemma4metadata.ConfigFile { + return &gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: 6}, + VisionConfig: &gemma4metadata.VisionConfig{HiddenSize: 4, IntermediateSize: 8, NumHiddenLayers: layers, NumAttentionHeads: 1, NumKeyValueHeads: 1, HeadDim: 4, DefaultOutputLength: 1, PatchSize: 2, PositionEmbeddingSize: 16, PoolingKernelSize: 1}, + } +} + +func testGemma4VisionTensorDescriptors(layers int) map[string]gemma4metadata.TensorDescriptor { + return testGemma4VisionTensorDescriptorsForGeometry(layers, 4, 8, 6, 2, 16, 4) +} + +func testGemma4VisionTensorDescriptorsForGeometry(layers int, hidden, intermediate, textHidden, patch, positions, headDim int32) map[string]gemma4metadata.TensorDescriptor { + descriptors := map[string]gemma4metadata.TensorDescriptor{ + "model.vision_tower.patch_embedder.input_proj.weight": {Dtype: "F32", Shape: []int32{hidden, 3 * patch * patch}}, + "model.vision_tower.patch_embedder.position_embedding_table": {Dtype: "F32", Shape: []int32{2, positions, hidden}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "F32", Shape: []int32{textHidden, hidden}}, + } + for i := range layers { + layer := fmt.Sprintf("model.vision_tower.encoder.layers.%d", i) + for _, suffix := range []string{".self_attn.q_proj.linear.weight", ".self_attn.k_proj.linear.weight", ".self_attn.v_proj.linear.weight", ".self_attn.o_proj.linear.weight"} { + descriptors[layer+suffix] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{hidden, hidden}} + } + for _, suffix := range []string{".mlp.gate_proj.linear.weight", ".mlp.up_proj.linear.weight"} { + descriptors[layer+suffix] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{intermediate, hidden}} + } + descriptors[layer+".mlp.down_proj.linear.weight"] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{hidden, intermediate}} + for _, suffix := range []string{".self_attn.q_norm.weight", ".self_attn.k_norm.weight"} { + descriptors[layer+suffix] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{headDim}} + } + for _, suffix := range []string{".input_layernorm.weight", ".post_attention_layernorm.weight", ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight"} { + descriptors[layer+suffix] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{hidden}} + } + } + return descriptors +} + +func TestGemma4InstalledVisionCapabilityReleasedUnequalHeadGeometry(t *testing.T) { + cfg := &gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: 2560}, + VisionConfig: &gemma4metadata.VisionConfig{ + HiddenSize: 768, IntermediateSize: 3072, NumHiddenLayers: 1, + NumAttentionHeads: 12, NumKeyValueHeads: 12, HeadDim: 64, + RMSNormEps: 1e-6, DefaultOutputLength: 280, PatchSize: 16, + PositionEmbeddingSize: 10240, PoolingKernelSize: 3, }, } + m := Model{ + Config: model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererLarge, Capabilities: []string{"vision"}}, + Gemma4VisionConfig: cfg, + Gemma4VisionTensors: testGemma4VisionTensorDescriptorsForGeometry(1, 768, 3072, 2560, 16, 10240, 64), + } + if !slices.Contains(m.Capabilities(), model.CapabilityVision) { + t.Fatal("released-compatible unequal hidden/head geometry did not expose vision") + } +} + +func gemma4VisionTensorNames(layers int) []string { + names := []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.position_embedding_table", + "model.embed_vision.embedding_projection.weight", + } + for i := range layers { + layer := fmt.Sprintf("model.vision_tower.encoder.layers.%d", i) + for _, projection := range []string{ + ".self_attn.q_proj.linear.weight", ".self_attn.k_proj.linear.weight", + ".self_attn.v_proj.linear.weight", ".self_attn.o_proj.linear.weight", + ".mlp.gate_proj.linear.weight", ".mlp.up_proj.linear.weight", ".mlp.down_proj.linear.weight", + } { + names = append(names, layer+projection) + } + for _, norm := range []string{ + ".self_attn.q_norm.weight", ".self_attn.k_norm.weight", + ".input_layernorm.weight", ".post_attention_layernorm.weight", + ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight", + } { + names = append(names, layer+norm) + } + } + return names } func TestModelCheckCapabilities(t *testing.T) { diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 46925719669..2c8d8655fdc 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -23,6 +23,7 @@ import ( ollamatemplate "github.com/ollama/ollama/template" "github.com/ollama/ollama/thinking" "github.com/ollama/ollama/types/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" ) type modelListSummary struct { @@ -391,10 +392,13 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum summary.Capabilities = appendModelListCapability(summary.Capabilities, model.CapabilityVision) } - if isLocalGemma4SafetensorsConfig(cfg) && !hasGemma4VisionTensorLayers(mf.Layers) { - summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { - return c == model.CapabilityVision - }) + if isLocalGemma4SafetensorsConfig(cfg) { + var gemma4cfg gemma4metadata.ConfigFile + if err := mf.ReadConfigJSON("config.json", &gemma4cfg); err != nil || !hasGemma4VisionTensorLayers(gemma4cfg, mf.Layers) { + summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { + return c == model.CapabilityVision + }) + } } summary.Capabilities = filterUnsupportedModelListCapabilities(summary.Capabilities, cfg) diff --git a/x/create/client/create.go b/x/create/client/create.go index c2cd6da8c5a..2b7e5cb2f33 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -26,6 +26,7 @@ import ( "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/x/create" imagemanifest "github.com/ollama/ollama/x/imagegen/manifest" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" "github.com/ollama/ollama/x/quant" ) @@ -630,18 +631,15 @@ func gemma4ModelDirHasVisionTensors(modelDir string) bool { if err != nil { return false } - - return hasAnySourceTensor(inv, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") && - hasAnySourceTensor(inv, "embed_vision.embedding_projection.weight", "model.embed_vision.embedding_projection.weight") -} - -func hasAnySourceTensor(inv create.Inventory, names ...string) bool { - for _, name := range names { - if inv.Has(name) { - return true - } + var cfg gemma4metadata.ConfigFile + if err := json.Unmarshal(inv.RawConfig, &cfg); err != nil { + return false } - return false + tensors := make(map[string]gemma4metadata.TensorDescriptor, len(inv.Tensors)) + for name, tensor := range inv.Tensors { + tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + } + return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil } // readChatTemplate returns the model's chat template, preferring the diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 09b898140e5..49d53991b57 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -2,7 +2,9 @@ package client import ( "encoding/json" + "fmt" "io" + "maps" "os" "path/filepath" "slices" @@ -11,8 +13,10 @@ import ( "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/parser" + "github.com/ollama/ollama/progress" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/x/create" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" "github.com/ollama/ollama/x/safetensors" ) @@ -551,7 +555,8 @@ func TestInferSafetensorsCapabilitiesGemma4VisionRequiresTensors(t *testing.T) { configJSON := `{ "architectures": ["Gemma4ForConditionalGeneration"], "model_type": "gemma4", - "vision_config": {"hidden_size": 1024} + "text_config": {"hidden_size": 6}, + "vision_config": {"hidden_size": 4, "intermediate_size": 8, "num_hidden_layers": 2, "num_attention_heads": 1, "num_key_value_heads": 1, "head_dim": 4, "default_output_length": 1, "patch_size": 2, "position_embedding_size": 16, "pooling_kernel_size": 1} }` tests := []struct { @@ -579,12 +584,17 @@ func TestInferSafetensorsCapabilitiesGemma4VisionRequiresTensors(t *testing.T) { want: []string{"completion"}, }, { - name: "vision tower and projector", + name: "partial vision tower and projector", tensors: []string{ "model.vision_tower.patch_embedder.input_proj.weight", "model.embed_vision.embedding_projection.weight", }, - want: []string{"completion", "vision"}, + want: []string{"completion"}, + }, + { + name: "complete vision tower and projector", + tensors: gemma4ClientVisionTensorNames(2), + want: []string{"completion", "vision"}, }, } @@ -605,6 +615,144 @@ func TestInferSafetensorsCapabilitiesGemma4VisionRequiresTensors(t *testing.T) { } } +func TestInferSafetensorsCapabilitiesGemma4PackedSourceRequiresProducerContract(t *testing.T) { + const configJSON = `{ + "architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4", + "text_config":{"hidden_size":24}, + "vision_config":{"hidden_size":16,"intermediate_size":32,"num_hidden_layers":1,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":16,"default_output_length":1,"patch_size":4,"position_embedding_size":16,"pooling_kernel_size":1} + }` + tensors := make(map[string]gemma4metadata.TensorDescriptor) + for _, name := range gemma4ClientVisionTensorNames(1) { + dtype, shape := gemma4ClientVisionDescriptorForDimensions(name, 16, 32, 24, 4, 16, 16) + tensors[name] = gemma4metadata.TensorDescriptor{Dtype: dtype, Shape: shape} + } + base := "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear" + delete(tensors, base+".weight") + tensors[base+".weight_packed"] = gemma4metadata.TensorDescriptor{Dtype: "U8", Shape: []int32{16, 8}} + tensors[base+".weight_scale"] = gemma4metadata.TensorDescriptor{Dtype: "F8_E4M3", Shape: []int32{16, 1}} + tensors[base+".weight_global_scale"] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: nil} + + check := func(t *testing.T, inventory map[string]gemma4metadata.TensorDescriptor, wantVision bool) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, inventory) + got := inferSafetensorsCapabilities(dir, "") + if slices.Contains(got, "vision") != wantVision { + t.Fatalf("capabilities = %v, wantVision %t", got, wantVision) + } + } + t.Run("complete compressed tensors contract", func(t *testing.T) { check(t, tensors, true) }) + t.Run("missing required global scale", func(t *testing.T) { + partial := maps.Clone(tensors) + delete(partial, base+".weight_global_scale") + check(t, partial, false) + }) + t.Run("wrong producer dtypes reject capability and import", func(t *testing.T) { + checkImport := func(t *testing.T, inventory map[string]gemma4metadata.TensorDescriptor, wantErr string) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, inventory) + t.Setenv("OLLAMA_MODELS", t.TempDir()) + p := progress.NewProgress(io.Discard) + defer p.Stop() + err := CreateModel(CreateOptions{ModelName: "wrong-producer-dtype", ModelDir: dir}, p) + if wantErr == "" && err != nil { + t.Fatalf("CreateModel() error = %v, want import with vision capability suppressed", err) + } + if wantErr != "" && (err == nil || !strings.Contains(err.Error(), wantErr)) { + t.Fatalf("CreateModel() error = %v, want %q", err, wantErr) + } + } + for _, tc := range []struct { + name, tensor, importErr string + }{ + {name: "compressed scale", tensor: base + ".weight_scale"}, + {name: "compressed global", tensor: base + ".weight_global_scale", importErr: "expected F32 tensor"}, + } { + t.Run(tc.name, func(t *testing.T) { + wrong := maps.Clone(tensors) + d := wrong[tc.tensor] + d.Dtype = "F16" + wrong[tc.tensor] = d + check(t, wrong, false) + checkImport(t, wrong, tc.importErr) + }) + } + + modelOpt := maps.Clone(tensors) + delete(modelOpt, base+".weight_packed") + delete(modelOpt, base+".weight_global_scale") + modelOpt[base+".weight"] = gemma4metadata.TensorDescriptor{Dtype: "U8", Shape: []int32{16, 8}} + modelOpt[base+".weight_scale_2"] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: nil} + for _, tc := range []struct { + name, tensor, importErr string + }{ + {name: "ModelOpt scale", tensor: base + ".weight_scale"}, + {name: "ModelOpt global", tensor: base + ".weight_scale_2", importErr: "expected F32 tensor"}, + } { + t.Run(tc.name, func(t *testing.T) { + wrong := maps.Clone(modelOpt) + d := wrong[tc.tensor] + d.Dtype = "F16" + wrong[tc.tensor] = d + check(t, wrong, false) + checkImport(t, wrong, tc.importErr) + }) + } + }) +} + +func TestInferSafetensorsCapabilitiesGemma4ReleasedUnequalHeadGeometry(t *testing.T) { + const configJSON = `{ + "architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4", + "text_config":{"hidden_size":2560}, + "vision_config":{"hidden_size":768,"intermediate_size":3072,"num_hidden_layers":1,"num_attention_heads":12,"num_key_value_heads":12,"head_dim":64,"rms_norm_eps":1e-6,"default_output_length":280,"patch_size":16,"position_embedding_size":10240,"pooling_kernel_size":3} + }` + descriptors := make(map[string]gemma4metadata.TensorDescriptor) + for _, name := range gemma4ClientVisionTensorNames(1) { + dtype, shape := gemma4ClientVisionDescriptorForDimensions(name, 768, 3072, 2560, 16, 10240, 64) + descriptors[name] = gemma4metadata.TensorDescriptor{Dtype: dtype, Shape: shape} + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, descriptors) + if got := inferSafetensorsCapabilities(dir, ""); !slices.Contains(got, "vision") { + t.Fatalf("released-compatible unequal hidden/head geometry capabilities = %v", got) + } +} + +func TestInferSafetensorsCapabilitiesGemma4RejectsMissingOrZeroTextWidth(t *testing.T) { + descriptors := make(map[string]gemma4metadata.TensorDescriptor) + for _, name := range gemma4ClientVisionTensorNames(1) { + dtype, shape := gemma4ClientVisionDescriptorForDimensions(name, 16, 32, 24, 4, 16, 16) + descriptors[name] = gemma4metadata.TensorDescriptor{Dtype: dtype, Shape: shape} + } + for _, tt := range []struct { + name, textConfig string + }{ + {name: "missing"}, + {name: "zero", textConfig: `"text_config":{"hidden_size":0},`}, + } { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + config := `{"architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4",` + tt.textConfig + `"vision_config":{"hidden_size":16,"intermediate_size":32,"num_hidden_layers":1,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":16,"default_output_length":1,"patch_size":4,"position_embedding_size":16,"pooling_kernel_size":1}}` + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(config), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, descriptors) + if got := inferSafetensorsCapabilities(dir, ""); slices.Contains(got, "vision") { + t.Fatalf("capabilities = %v, missing/zero text width exposed vision", got) + } + }) + } +} + func TestGemma4ModelConfigRejectsNearMatches(t *testing.T) { if isGemma4ModelConfig([]string{"NotGemma4ForConditionalGeneration"}, "gemma4-next") { t.Fatal("near-match Gemma 4 identifiers classified as Gemma 4") @@ -614,12 +762,39 @@ func TestGemma4ModelConfigRejectsNearMatches(t *testing.T) { } } +func gemma4ClientVisionTensorNames(layers int) []string { + names := []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.position_embedding_table", + "model.embed_vision.embedding_projection.weight", + } + for i := range layers { + layer := fmt.Sprintf("model.vision_tower.encoder.layers.%d", i) + for _, projection := range []string{ + ".self_attn.q_proj.linear.weight", ".self_attn.k_proj.linear.weight", + ".self_attn.v_proj.linear.weight", ".self_attn.o_proj.linear.weight", + ".mlp.gate_proj.linear.weight", ".mlp.up_proj.linear.weight", ".mlp.down_proj.linear.weight", + } { + names = append(names, layer+projection) + } + for _, norm := range []string{ + ".self_attn.q_norm.weight", ".self_attn.k_norm.weight", + ".input_layernorm.weight", ".post_attention_layernorm.weight", + ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight", + } { + names = append(names, layer+norm) + } + } + return names +} + func writeClientSafetensors(t *testing.T, dir string, names ...string) { t.Helper() tensors := make([]*safetensors.TensorData, 0, len(names)) for _, name := range names { - tensors = append(tensors, safetensors.NewTensorDataFromBytes(name, "U8", []int32{1}, []byte{0})) + dtype, shape := gemma4ClientVisionDescriptor(name) + tensors = append(tensors, safetensors.NewTensorDataFromBytes(name, dtype, shape, clientSafetensorData(dtype, shape))) } data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader(tensors)) @@ -631,6 +806,67 @@ func writeClientSafetensors(t *testing.T, dir string, names ...string) { } } +func writeClientSafetensorDescriptors(t *testing.T, dir string, descriptors map[string]gemma4metadata.TensorDescriptor) { + t.Helper() + tensors := make([]*safetensors.TensorData, 0, len(descriptors)) + for name, descriptor := range descriptors { + raw := clientSafetensorData(descriptor.Dtype, descriptor.Shape) + if strings.EqualFold(descriptor.Dtype, "F32") && len(raw) == 4 { + copy(raw, []byte{0, 0, 128, 63}) + } + tensors = append(tensors, safetensors.NewTensorDataFromBytes(name, descriptor.Dtype, descriptor.Shape, raw)) + } + data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader(tensors)) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "model.safetensors"), data, 0o644); err != nil { + t.Fatal(err) + } +} + +func clientSafetensorData(dtype string, shape []int32) []byte { + width := 1 + switch strings.ToUpper(dtype) { + case "F64", "I64", "U64": + width = 8 + case "F32", "I32", "U32": + width = 4 + case "F16", "BF16", "I16", "U16": + width = 2 + } + elements := 1 + for _, dim := range shape { + elements *= int(dim) + } + return make([]byte, elements*width) +} + +func gemma4ClientVisionDescriptor(name string) (string, []int32) { + return gemma4ClientVisionDescriptorForDimensions(name, 4, 8, 6, 2, 16, 4) +} + +func gemma4ClientVisionDescriptorForDimensions(name string, hidden, intermediate, textHidden, patch, positions, headDim int32) (string, []int32) { + switch { + case strings.HasSuffix(name, "patch_embedder.input_proj.weight"): + return "F32", []int32{hidden, 3 * patch * patch} + case strings.HasSuffix(name, "position_embedding_table"): + return "F32", []int32{2, positions, hidden} + case strings.HasSuffix(name, "embed_vision.embedding_projection.weight"): + return "F32", []int32{textHidden, hidden} + case strings.Contains(name, ".mlp.gate_proj."), strings.Contains(name, ".mlp.up_proj."): + return "F32", []int32{intermediate, hidden} + case strings.Contains(name, ".mlp.down_proj."): + return "F32", []int32{hidden, intermediate} + case strings.Contains(name, ".self_attn.q_norm.weight"), strings.Contains(name, ".self_attn.k_norm.weight"): + return "F32", []int32{headDim} + case strings.Contains(name, "layernorm.weight"): + return "F32", []int32{hidden} + default: + return "F32", []int32{hidden, hidden} + } +} + func TestCreateModelfileLayersIncludesParameters(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) diff --git a/x/mlxrunner/grammar_test.go b/x/mlxrunner/grammar_test.go index c1e83aefc52..94b539de73d 100644 --- a/x/mlxrunner/grammar_test.go +++ b/x/mlxrunner/grammar_test.go @@ -1,6 +1,7 @@ package mlxrunner import ( + "context" "encoding/json" "errors" "net/http" @@ -189,7 +190,7 @@ func TestPrepareGrammarUnavailable(t *testing.T) { Prompt: "0", Format: json.RawMessage(`"json"`), }} - err := r.Prepare(request) + err := r.Prepare(context.Background(), request) var statusErr api.StatusError if !errors.As(err, &statusErr) { t.Fatalf("Prepare error = %T %v, want api.StatusError", err, err) diff --git a/x/mlxrunner/media.go b/x/mlxrunner/media.go index c29a1d7f388..91990e91ab3 100644 --- a/x/mlxrunner/media.go +++ b/x/mlxrunner/media.go @@ -1,6 +1,7 @@ package mlxrunner import ( + "context" "encoding/binary" "errors" "fmt" @@ -46,7 +47,7 @@ func foldValue(data []byte, dims []int) uint32 { // use, released when the expansion is fully evaluated. A nil // *requestMedia is a text-only request; every method is nil-safe. type requestMedia struct { - model base.MediaModel + model base.MediaEncoder items []mediaItem inputLen int @@ -65,7 +66,7 @@ func (r *Runner) openMedia(request Request) *requestMedia { return nil } m := &requestMedia{ - model: r.Model.(base.MediaModel), + model: r.Model.(base.MediaEncoder), items: request.MediaItems, inputLen: len(request.Tokens), manifest: make([]batch.MediaItem, len(request.MediaItems)), @@ -172,7 +173,7 @@ func (m *requestMedia) close() { // expandMedia tokenizes the [img-N]-tagged prompt into segments, expands // them in a single PrepareMedia call, and validates the authored items // before keying cache identity on them. -func (r *Runner) expandMedia(mm base.MediaModel, prompt string, media []llm.MediaData) (*base.PreparedRequest, []mediaItem, error) { +func (r *Runner) expandMedia(ctx context.Context, prepare func(context.Context, []base.Segment) (*base.PreparedRequest, error), prompt string, media []llm.MediaData) (*base.PreparedRequest, []mediaItem, error) { matches := imgTagPattern.FindAllStringSubmatch(prompt, -1) parts := imgTagPattern.Split(prompt, -1) @@ -205,7 +206,7 @@ func (r *Runner) expandMedia(mm base.MediaModel, prompt string, media []llm.Medi } } - prepared, err := mm.PrepareMedia(segments) + prepared, err := prepare(ctx, segments) if err != nil { return nil, nil, err } diff --git a/x/mlxrunner/model/base/media.go b/x/mlxrunner/model/base/media.go index 513add66faa..5697a2171d5 100644 --- a/x/mlxrunner/model/base/media.go +++ b/x/mlxrunner/model/base/media.go @@ -1,6 +1,7 @@ package base import ( + "context" // Every model's PrepareMedia decodes through image.Decode; the decoder // set is registered once here so all models accept the same formats. _ "image/gif" @@ -61,13 +62,20 @@ type PreparedRequest struct { Layout any } -// MediaModel is implemented by models that accept media inputs. +// MediaEncoder is the MLX-thread half of a model's media contract. +type MediaEncoder interface { + EncodeMedia(item *PreparedItem, data *mlx.Array) *mlx.Array +} + +// MediaModel is implemented by models that accept media inputs. Its CPU +// preparation receives the live request context so expensive work can stop +// before the request is queued. type MediaModel interface { // PrepareMedia runs once per request on the request goroutine, CPU // only, and returns the expanded stream. It must be deterministic for // given segments: prefix-cache restores splice cached state with // recomputed state. - PrepareMedia(segments []Segment) (*PreparedRequest, error) + PrepareMedia(ctx context.Context, segments []Segment) (*PreparedRequest, error) // EncodeMedia builds one item's lazy feature graph on the MLX thread; // it must not evaluate — the consuming forward's evaluation pulls it. diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index ec2844aff29..eb9d208eec2 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -26,27 +26,39 @@ func prefillChunkSize() int { // Prepare tokenizes the prompt and validates it against the model's // context length. It is safe to call from any goroutine. On success it // populates request.Tokens and adjusts request.Options.NumPredict. -func (r *Runner) Prepare(request *Request) (err error) { +func (r *Runner) Prepare(ctx context.Context, request *Request) (err error) { + request.Grammar = nil + request.Tokens = nil + request.MediaItems = nil + request.Layout = nil + var grammar *grammarCompilation + defer func() { + if err != nil { + grammar.close() + request.Grammar = nil + request.Tokens = nil + request.MediaItems = nil + request.Layout = nil + } + }() + if r.Model == nil { return errors.New("model not loaded") } + if err := ctx.Err(); err != nil { + return err + } // Launched first so the compile overlaps tokenization and media // preparation as well as prefill. - grammar, err := r.grammarEngine.prepare(request.Format) + grammar, err = r.grammarEngine.prepare(request.Format) if err != nil { return err } - request.Grammar = grammar - defer func() { - if err != nil { - request.Grammar.close() - request.Grammar = nil - } - }() var tokens []int32 var items []mediaItem + var layout any if len(request.Media) == 0 { tokens = r.Tokenizer.Encode(request.Prompt, r.Tokenizer.AddBOS()) } else { @@ -58,12 +70,15 @@ func (r *Runner) Prepare(request *Request) (err error) { } return fmt.Errorf("this model does not support %s input", kind) } - prepared, bound, err := r.expandMedia(mm, request.Prompt, request.Media) + prepared, bound, err := r.expandMedia(ctx, mm.PrepareMedia, request.Prompt, request.Media) if err != nil { return err } + if err := ctx.Err(); err != nil { + return err + } tokens, items = prepared.Tokens, bound - request.Layout = prepared.Layout + layout = prepared.Layout } if len(tokens) == 0 { @@ -76,14 +91,21 @@ func (r *Runner) Prepare(request *Request) (err error) { // Cap generation to stay within the model's context length maxGenerate := r.contextLength - len(tokens) - if request.Options.NumPredict <= 0 { - request.Options.NumPredict = maxGenerate + numPredict := request.Options.NumPredict + if numPredict <= 0 { + numPredict = maxGenerate } else { - request.Options.NumPredict = min(request.Options.NumPredict, maxGenerate) + numPredict = min(numPredict, maxGenerate) + } + if err := ctx.Err(); err != nil { + return err } + request.Grammar = grammar request.Tokens = tokens request.MediaItems = items + request.Layout = layout + request.Options.NumPredict = numPredict return nil } diff --git a/x/mlxrunner/pipeline_test.go b/x/mlxrunner/pipeline_test.go index a80fce634e4..49f4d689dba 100644 --- a/x/mlxrunner/pipeline_test.go +++ b/x/mlxrunner/pipeline_test.go @@ -1,11 +1,14 @@ package mlxrunner import ( + "context" + "errors" "fmt" "slices" "strings" "testing" + "github.com/ollama/ollama/api" "github.com/ollama/ollama/llm" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/cache" @@ -35,7 +38,7 @@ func TestPrepareRejectsMediaWithoutSupport(t *testing.T) { }, } - err := r.Prepare(req) + err := r.Prepare(context.Background(), req) if err == nil { t.Fatal("expected error for media on a text-only model") } @@ -51,9 +54,16 @@ type stubMediaModel struct { textOnlyModel expansion []int32 layout any + observe func(context.Context) } -func (m stubMediaModel) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { +func (m stubMediaModel) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { + if m.observe != nil { + m.observe(ctx) + } + if err := ctx.Err(); err != nil { + return nil, err + } prepared := &base.PreparedRequest{Layout: m.layout} for s, seg := range segments { if seg.Data == nil { @@ -79,8 +89,78 @@ func (m stubMediaModel) PrepareMedia(segments []base.Segment) (*base.PreparedReq return prepared, nil } +func TestPreparePassesLiveContextAndPublishesNothingWhenCanceled(t *testing.T) { + type contextKey struct{} + ctx := context.WithValue(context.Background(), contextKey{}, "live") + var observed context.Context + r := mediaTestRunner(t) + r.Model = stubMediaModel{ + expansion: []int32{70, 71}, + observe: func(ctx context.Context) { observed = ctx }, + } + req := &Request{CompletionRequest: CompletionRequest{ + Prompt: "[img-0]", + Media: []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte("img")}}, + }} + if err := r.Prepare(ctx, req); err != nil { + t.Fatal(err) + } + if observed != ctx || observed.Value(contextKey{}) != "live" { + t.Fatal("PrepareMedia did not receive the exact live request context") + } + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + failed := &Request{CompletionRequest: req.CompletionRequest} + failed.Format = []byte(`{"type":"object"}`) + r.grammarEngine = &grammarEngine{} + if err := r.Prepare(canceled, failed); !errors.Is(err, context.Canceled) { + t.Fatalf("Prepare() error = %v, want context.Canceled", err) + } + if failed.Tokens != nil || failed.MediaItems != nil || failed.Layout != nil || failed.Grammar != nil { + t.Fatalf("canceled Prepare published partial request: %+v", failed) + } +} + func (stubMediaModel) EncodeMedia(*base.PreparedItem, *mlx.Array) *mlx.Array { return nil } +type cancellationIgnoringMediaModel struct { + stubMediaModel + cancel context.CancelFunc +} + +func (m cancellationIgnoringMediaModel) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { + prepared, err := m.stubMediaModel.PrepareMedia(ctx, segments) + m.cancel() + return prepared, err +} + +func TestPrepareRejectsSuccessReturnedAfterCancellation(t *testing.T) { + for _, initialNumPredict := range []int{0, 17} { + t.Run(fmt.Sprintf("num_predict_%d", initialNumPredict), func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + r := mediaTestRunner(t) + r.Model = cancellationIgnoringMediaModel{stubMediaModel: stubMediaModel{expansion: []int32{70, 71}, layout: "L"}, cancel: cancel} + r.grammarEngine = &grammarEngine{} + req := &Request{CompletionRequest: CompletionRequest{ + Prompt: "[img-0]", + Media: []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte("img")}}, + Format: []byte(`{"type":"object"}`), + Options: api.Options{NumPredict: initialNumPredict}, + }} + if err := r.Prepare(ctx, req); !errors.Is(err, context.Canceled) { + t.Fatalf("Prepare() error = %v, want context.Canceled", err) + } + if req.Grammar != nil || req.Tokens != nil || req.MediaItems != nil || req.Layout != nil { + t.Fatalf("canceled Prepare published state: %+v", req) + } + if req.Options.NumPredict != initialNumPredict { + t.Fatalf("NumPredict = %d, want unchanged %d", req.Options.NumPredict, initialNumPredict) + } + }) + } +} + func mediaTestRunner(t *testing.T) *Runner { t.Helper() return &Runner{ @@ -98,7 +178,7 @@ func TestPrepareExpandsMediaTags(t *testing.T) { Media: []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte("img")}}, }, } - if err := r.Prepare(req); err != nil { + if err := r.Prepare(context.Background(), req); err != nil { t.Fatal(err) } @@ -131,7 +211,7 @@ type rawMediaModel struct { prepared base.PreparedRequest } -func (m rawMediaModel) PrepareMedia([]base.Segment) (*base.PreparedRequest, error) { +func (m rawMediaModel) PrepareMedia(context.Context, []base.Segment) (*base.PreparedRequest, error) { p := m.prepared return &p, nil } @@ -157,7 +237,7 @@ func TestPrepareValidatesAuthoredItems(t *testing.T) { for _, c := range cases { r := mediaTestRunner(t) r.Model = rawMediaModel{prepared: base.PreparedRequest{Tokens: []int32{0, 5, 5, 5, 1}, Items: c.items}} - err := r.Prepare(&Request{CompletionRequest: CompletionRequest{Prompt: "0[img-0]1", Media: media}}) + err := r.Prepare(context.Background(), &Request{CompletionRequest: CompletionRequest{Prompt: "0[img-0]1", Media: media}}) if c.want == "" { if err != nil { t.Fatalf("%s: unexpected error %v", c.name, err) @@ -174,14 +254,14 @@ func TestPrepareMediaErrors(t *testing.T) { media := []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte("img")}} r := mediaTestRunner(t) - err := r.Prepare(&Request{CompletionRequest: CompletionRequest{Prompt: "0[img-3]1", Media: media}}) + err := r.Prepare(context.Background(), &Request{CompletionRequest: CompletionRequest{Prompt: "0[img-3]1", Media: media}}) if err == nil || !strings.Contains(err.Error(), "invalid image index: 3") { t.Fatalf("missing-ID error = %v", err) } // Unreferenced media is ignored with a warning; the prompt still works. req := &Request{CompletionRequest: CompletionRequest{Prompt: "01", Media: media}} - if err := r.Prepare(req); err != nil { + if err := r.Prepare(context.Background(), req); err != nil { t.Fatal(err) } if len(req.MediaItems) != 0 { @@ -191,7 +271,7 @@ func TestPrepareMediaErrors(t *testing.T) { // Duplicate references are allowed; each occurrence is its own item with // the same fold. req = &Request{CompletionRequest: CompletionRequest{Prompt: "[img-0]0[img-0]", Media: media}} - if err := r.Prepare(req); err != nil { + if err := r.Prepare(context.Background(), req); err != nil { t.Fatal(err) } if len(req.MediaItems) != 2 || req.MediaItems[0].fold != req.MediaItems[1].fold { @@ -200,7 +280,7 @@ func TestPrepareMediaErrors(t *testing.T) { // A zero-length expansion cannot carry identity into the trie keys. r.Model = stubMediaModel{} - err = r.Prepare(&Request{CompletionRequest: CompletionRequest{Prompt: "[img-0]", Media: media}}) + err = r.Prepare(context.Background(), &Request{CompletionRequest: CompletionRequest{Prompt: "[img-0]", Media: media}}) if err == nil || !strings.Contains(err.Error(), "no tokens") { t.Fatalf("zero-expansion error = %v", err) } diff --git a/x/mlxrunner/server.go b/x/mlxrunner/server.go index 7178e486057..52e04c41cbf 100644 --- a/x/mlxrunner/server.go +++ b/x/mlxrunner/server.go @@ -145,7 +145,11 @@ func Execute(args []string) error { TopLogprobs: request.TopLogprobs, } - if err := runner.Prepare(&request); err != nil { + var cancel context.CancelFunc + request.Ctx, cancel = context.WithCancel(r.Context()) + defer cancel() + + if err := runner.prepareAndQueue(request.Ctx, &request); err != nil { var statusErr api.StatusError if errors.As(err, &statusErr) { http.Error(w, statusErr.ErrorMessage, statusErr.StatusCode) @@ -155,18 +159,6 @@ func Execute(args []string) error { return } - var cancel context.CancelFunc - request.Ctx, cancel = context.WithCancel(r.Context()) - defer cancel() - - select { - case <-r.Context().Done(): - // Never queued, so the runner will not close the grammar. - request.Grammar.close() - return - case runner.Requests <- request: - } - w.Header().Set("Content-Type", "application/jsonl") w.WriteHeader(http.StatusOK) enc := json.NewEncoder(w) @@ -236,6 +228,36 @@ func Execute(args []string) error { })) } +func (runner *Runner) prepareAndQueue(ctx context.Context, request *Request) error { + originalNumPredict := request.Options.NumPredict + queued := false + defer func() { + if !queued { + request.Grammar.close() + request.Grammar = nil + request.Tokens = nil + request.MediaItems = nil + request.Layout = nil + request.Options.NumPredict = originalNumPredict + } + }() + + request.Ctx = ctx + if err := runner.Prepare(ctx, request); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + select { + case <-ctx.Done(): + return ctx.Err() + case runner.Requests <- *request: + queued = true + return nil + } +} + type statusRecorder struct { http.ResponseWriter code int diff --git a/x/mlxrunner/server_test.go b/x/mlxrunner/server_test.go new file mode 100644 index 00000000000..799fae92121 --- /dev/null +++ b/x/mlxrunner/server_test.go @@ -0,0 +1,101 @@ +package mlxrunner + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/llm" +) + +type queueBlockingContext struct { + done chan struct{} + reached chan struct{} + once sync.Once + mu sync.Mutex + err error +} + +func newQueueBlockingContext() *queueBlockingContext { + return &queueBlockingContext{done: make(chan struct{}), reached: make(chan struct{})} +} + +func (*queueBlockingContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (*queueBlockingContext) Value(any) any { return nil } + +func (c *queueBlockingContext) Done() <-chan struct{} { + c.once.Do(func() { close(c.reached) }) + return c.done +} + +func (c *queueBlockingContext) Err() error { + c.mu.Lock() + defer c.mu.Unlock() + return c.err +} + +func (c *queueBlockingContext) cancel() { + c.mu.Lock() + c.err = context.Canceled + c.mu.Unlock() + close(c.done) +} + +func TestPrepareAndQueueCanceledRequest(t *testing.T) { + runner := mediaTestRunner(t) + runner.Requests = make(chan Request, 1) + request := &Request{CompletionRequest: CompletionRequest{ + Prompt: "[img-0]", + Media: []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte("img")}}, + }} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if err := runner.prepareAndQueue(ctx, request); !errors.Is(err, context.Canceled) { + t.Fatalf("prepareAndQueue() error = %v, want context.Canceled", err) + } + if len(runner.Requests) != 0 { + t.Fatal("canceled request was queued") + } + if request.Tokens != nil || request.MediaItems != nil || request.Layout != nil { + t.Fatalf("canceled request retained partial preparation: %+v", request) + } +} + +func TestPrepareAndQueueClosesStructuredFormatOnCancellation(t *testing.T) { + for _, initialNumPredict := range []int{0, 17} { + t.Run(fmt.Sprintf("num_predict_%d", initialNumPredict), func(t *testing.T) { + runner := mediaTestRunner(t) + runner.grammarEngine = &grammarEngine{} + runner.Requests = make(chan Request) + request := &Request{CompletionRequest: CompletionRequest{ + Prompt: "0", + Format: json.RawMessage(`{"type":"object"}`), + Options: api.Options{NumPredict: initialNumPredict}, + }} + ctx := newQueueBlockingContext() + errCh := make(chan error, 1) + go func() { errCh <- runner.prepareAndQueue(ctx, request) }() + + <-ctx.reached + ctx.cancel() + if err := <-errCh; !errors.Is(err, context.Canceled) { + t.Fatalf("prepareAndQueue() error = %v, want context.Canceled", err) + } + if request.Grammar != nil || request.Tokens != nil || request.MediaItems != nil || request.Layout != nil { + t.Fatalf("canceled queued request retained published state: %+v", request) + } + if request.Options.NumPredict != initialNumPredict { + t.Fatalf("NumPredict = %d, want restored %d", request.Options.NumPredict, initialNumPredict) + } + if len(runner.Requests) != 0 { + t.Fatal("canceled request was queued") + } + }) + } +} diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index edebc08e6ba..41e7d9a0665 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -771,7 +771,11 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { m.LMHead = m.EmbedTokens.AsLinear() } - if m.VisionConfig != nil && hasGemma4VisionWeights(tensors) { + visionReady, err := validateGemma4VisionWeights(tensors, m.VisionConfig, int(m.HiddenSize), m.TensorQuant) + if err != nil { + return fmt.Errorf("invalid Gemma4 vision tensors: %w", err) + } + if visionReady { vision, err := loadVisionModel(tensors, m.VisionConfig, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) if err != nil { return err diff --git a/x/models/gemma4/metadata/vision.go b/x/models/gemma4/metadata/vision.go new file mode 100644 index 00000000000..d5bfd86422d --- /dev/null +++ b/x/models/gemma4/metadata/vision.go @@ -0,0 +1,669 @@ +// Package metadata contains Gemma 4 model metadata checks that do not depend +// on MLX, allowing import, server, and runtime paths to share one definition. +package metadata + +import ( + "fmt" + "math" + "slices" + "strings" +) + +const ( + defaultVisionLayers = 16 + MaxVisionLayers = 512 + MaxVisionHidden = 32_768 + MaxVisionIntermediate = 262_144 + MaxVisionHeads = 512 + MaxVisionHeadDim = 1_024 + MaxVisionSoftTokens = 16_384 + MaxImageDimension = 16_384 + MaxResizePixels = 1 << 20 + MaxPositionEntries = 1 << 20 + MaxPositionValues = 1 << 26 +) + +type ConfigFile struct { + TextConfig TextConfig `json:"text_config"` + VisionConfig *VisionConfig `json:"vision_config"` + Quantization Quantization `json:"quantization"` + QuantizationConfig Quantization `json:"quantization_config"` +} + +type TextConfig struct { + HiddenSize int `json:"hidden_size"` +} + +type Quantization struct { + Bits int `json:"bits"` + GroupSize int `json:"group_size"` + Mode string `json:"mode"` + Method string `json:"quant_method"` +} + +type VisionConfig struct { + HiddenSize int `json:"hidden_size"` + IntermediateSize int `json:"intermediate_size"` + NumHiddenLayers int `json:"num_hidden_layers"` + NumAttentionHeads int `json:"num_attention_heads"` + NumKeyValueHeads int `json:"num_key_value_heads"` + HeadDim int `json:"head_dim"` + RMSNormEps float64 `json:"rms_norm_eps"` + DefaultOutputLength int `json:"default_output_length"` + PatchSize int `json:"patch_size"` + PositionEmbeddingSize int `json:"position_embedding_size"` + PoolingKernelSize int `json:"pooling_kernel_size"` + UseClippedLinears bool `json:"use_clipped_linears"` + Standardize bool `json:"standardize"` + RopeParameters struct { + RopeTheta float64 `json:"rope_theta"` + } `json:"rope_parameters"` +} + +type Architecture struct { + TextHiddenSize int + HiddenSize int + IntermediateSize int + NumHiddenLayers int + NumAttentionHeads int + NumKeyValueHeads int + HeadDim int + RMSNormEps float64 + DefaultOutputLength int + PatchSize int + PositionEmbeddingSize int + PoolingKernelSize int + UseClippedLinears bool + ClippingBoundsOptional bool + Standardize bool + RopeTheta float64 +} + +type TensorDescriptor struct { + Dtype string + Shape []int32 + QuantType string + GroupSize int +} + +func ProjectVisionArchitecture(cfg ConfigFile) (Architecture, error) { + return projectVisionArchitecture(cfg, true) +} + +func projectVisionArchitecture(cfg ConfigFile, requireTextWidth bool) (Architecture, error) { + if cfg.VisionConfig == nil { + return Architecture{}, fmt.Errorf("missing vision_config") + } + v := *cfg.VisionConfig + if v.HiddenSize == 0 { + v.HiddenSize = 768 + } + if v.IntermediateSize == 0 { + v.IntermediateSize = 3072 + } + if v.NumHiddenLayers == 0 { + v.NumHiddenLayers = defaultVisionLayers + } + if v.NumAttentionHeads == 0 { + v.NumAttentionHeads = 12 + } + if v.NumKeyValueHeads == 0 { + v.NumKeyValueHeads = v.NumAttentionHeads + } + if v.HeadDim == 0 { + v.HeadDim = 64 + } + if v.RMSNormEps == 0 { + v.RMSNormEps = 1e-6 + } + if v.DefaultOutputLength == 0 { + v.DefaultOutputLength = 280 + } + if v.PatchSize == 0 { + v.PatchSize = 16 + } + if v.PositionEmbeddingSize == 0 { + v.PositionEmbeddingSize = 10240 + } + if v.PoolingKernelSize == 0 { + v.PoolingKernelSize = 3 + } + if v.RopeParameters.RopeTheta == 0 { + v.RopeParameters.RopeTheta = 100 + } + + a := Architecture{ + TextHiddenSize: cfg.TextConfig.HiddenSize, + HiddenSize: v.HiddenSize, IntermediateSize: v.IntermediateSize, + NumHiddenLayers: v.NumHiddenLayers, NumAttentionHeads: v.NumAttentionHeads, + NumKeyValueHeads: v.NumKeyValueHeads, HeadDim: v.HeadDim, + RMSNormEps: v.RMSNormEps, DefaultOutputLength: v.DefaultOutputLength, + PatchSize: v.PatchSize, PositionEmbeddingSize: v.PositionEmbeddingSize, + PoolingKernelSize: v.PoolingKernelSize, UseClippedLinears: v.UseClippedLinears, + // Published Gemma 4 towers may declare clippable linears while + // omitting every bound. Mixed/partial bound sets are never permitted. + ClippingBoundsOptional: true, + Standardize: v.Standardize, RopeTheta: v.RopeParameters.RopeTheta, + } + if err := validateArchitecture(a, requireTextWidth); err != nil { + return Architecture{}, err + } + return a, nil +} + +func validateArchitecture(a Architecture, requireTextWidth bool) error { + if a.TextHiddenSize < 0 || a.TextHiddenSize > MaxVisionHidden || (requireTextWidth && a.TextHiddenSize == 0) { + return fmt.Errorf("invalid Gemma4 text hidden_size %d", a.TextHiddenSize) + } + for _, f := range []struct { + name string + value, max int + }{ + {"hidden_size", a.HiddenSize, MaxVisionHidden}, + {"intermediate_size", a.IntermediateSize, MaxVisionIntermediate}, + {"num_hidden_layers", a.NumHiddenLayers, MaxVisionLayers}, + {"num_attention_heads", a.NumAttentionHeads, MaxVisionHeads}, + {"num_key_value_heads", a.NumKeyValueHeads, MaxVisionHeads}, + {"head_dim", a.HeadDim, MaxVisionHeadDim}, + {"default_output_length", a.DefaultOutputLength, MaxVisionSoftTokens}, + {"patch_size", a.PatchSize, MaxImageDimension}, + {"position_embedding_size", a.PositionEmbeddingSize, MaxPositionEntries}, + {"pooling_kernel_size", a.PoolingKernelSize, MaxImageDimension}, + } { + if f.value <= 0 || f.value > f.max { + return fmt.Errorf("invalid Gemma4 vision %s %d (limit %d)", f.name, f.value, f.max) + } + } + if a.NumKeyValueHeads > a.NumAttentionHeads || a.NumAttentionHeads%a.NumKeyValueHeads != 0 { + return fmt.Errorf("invalid Gemma4 vision attention heads %d/%d", a.NumAttentionHeads, a.NumKeyValueHeads) + } + if p, ok := checkedProduct(int64(MaxVisionHidden), int64(a.NumAttentionHeads), int64(a.HeadDim)); !ok || p != int64(a.HiddenSize) { + return fmt.Errorf("Gemma4 vision attention width does not match hidden_size %d", a.HiddenSize) + } + if !finitePositive(a.RMSNormEps) { + return fmt.Errorf("invalid Gemma4 vision rms_norm_eps %v", a.RMSNormEps) + } + if !finitePositive(float64(float32(a.RMSNormEps))) { + return fmt.Errorf("invalid Gemma4 vision rms_norm_eps %v after float32 projection", a.RMSNormEps) + } + if !finitePositive(a.RopeTheta) { + return fmt.Errorf("invalid Gemma4 vision rope_theta %v", a.RopeTheta) + } + if !finitePositive(float64(float32(a.RopeTheta))) { + return fmt.Errorf("invalid Gemma4 vision rope_theta %v after float32 projection", a.RopeTheta) + } + if _, ok := checkedProduct(MaxResizePixels, int64(a.DefaultOutputLength), int64(a.PoolingKernelSize), int64(a.PoolingKernelSize), int64(a.PatchSize), int64(a.PatchSize)); !ok { + return fmt.Errorf("Gemma4 vision resize budget exceeds limit %d", MaxResizePixels) + } + if _, ok := checkedProduct(MaxPositionValues, int64(a.DefaultOutputLength), int64(a.PoolingKernelSize), int64(a.PoolingKernelSize), int64(a.HeadDim)); !ok { + return fmt.Errorf("Gemma4 vision position allocation exceeds limit %d", MaxPositionValues) + } + patchSide, ok := checkedProduct(MaxPositionEntries, int64(a.DefaultOutputLength), int64(a.PoolingKernelSize)) + if !ok || patchSide > int64(a.PositionEmbeddingSize) { + return fmt.Errorf("Gemma4 vision patch side %d exceeds position table size %d", patchSide, a.PositionEmbeddingSize) + } + return nil +} + +func finitePositive(v float64) bool { return v > 0 && !math.IsNaN(v) && !math.IsInf(v, 0) } +func checkedProduct(limit int64, values ...int64) (int64, bool) { + if limit <= 0 { + return 0, false + } + p := int64(1) + for _, v := range values { + if v <= 0 || p > limit/v { + return 0, false + } + p *= v + } + return p, true +} + +// ValidateVisionTensors is the non-executable names-only normalized gate for +// callers that cannot inspect blobs or know the enclosing text width. +// Capability surfaces use the strict descriptor variants below. +func ValidateVisionTensors(cfg ConfigFile, names []string) error { + a, err := projectVisionArchitecture(cfg, false) + if err != nil { + return err + } + return validateNames(a, names, installedMode) +} + +func ValidateVisionSourceTensors(cfg ConfigFile, names []string) error { + a, err := projectVisionArchitecture(cfg, false) + if err != nil { + return err + } + return validateNames(a, names, sourceMode) +} + +func ValidateVisionSourceInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + return validateInventory(cfg, tensors, sourceMode) +} + +func ValidateVisionInstalledInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + return validateInventory(cfg, tensors, installedMode) +} + +func ValidateVisionRuntimeInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + return validateInventory(cfg, tensors, runtimeMode) +} + +type inventoryMode int + +const ( + sourceMode inventoryMode = iota + installedMode + runtimeMode +) + +func validateInventory(cfg ConfigFile, tensors map[string]TensorDescriptor, mode inventoryMode) error { + a, err := ProjectVisionArchitecture(cfg) + if err != nil { + return err + } + if a.TextHiddenSize <= 0 { + return fmt.Errorf("missing Gemma4 text hidden_size") + } + names := make([]string, 0, len(tensors)) + for name := range tensors { + names = append(names, name) + } + if err := validateNames(a, names, mode); err != nil { + return err + } + prefix, _ := visionPrefix(names, mode) + patchDim, _ := checkedProduct(math.MaxInt32, 3, int64(a.PatchSize), int64(a.PatchSize)) + if err := requireLinearDescriptor(tensors, prefix+"vision_tower.patch_embedder.input_proj", []int32{int32(a.HiddenSize), int32(patchDim)}, mode, cfg); err != nil { + return err + } + if err := requireDense(tensors, prefix+"vision_tower.patch_embedder.position_embedding_table", []int32{2, int32(a.PositionEmbeddingSize), int32(a.HiddenSize)}); err != nil { + return err + } + for i := range a.NumHiddenLayers { + layer := fmt.Sprintf("%svision_tower.encoder.layers.%d", prefix, i) + for _, p := range []struct { + suffix string + shape []int32 + }{ + {".self_attn.q_proj", []int32{int32(a.HiddenSize), int32(a.HiddenSize)}}, + {".self_attn.k_proj", []int32{int32(a.NumKeyValueHeads * a.HeadDim), int32(a.HiddenSize)}}, + {".self_attn.v_proj", []int32{int32(a.NumKeyValueHeads * a.HeadDim), int32(a.HiddenSize)}}, + {".self_attn.o_proj", []int32{int32(a.HiddenSize), int32(a.HiddenSize)}}, + {".mlp.gate_proj", []int32{int32(a.IntermediateSize), int32(a.HiddenSize)}}, + {".mlp.up_proj", []int32{int32(a.IntermediateSize), int32(a.HiddenSize)}}, + {".mlp.down_proj", []int32{int32(a.HiddenSize), int32(a.IntermediateSize)}}, + } { + if err := requireLinearDescriptor(tensors, layer+p.suffix, p.shape, mode, cfg); err != nil { + return err + } + if err := validateClipDescriptors(tensors, layer+p.suffix, a.UseClippedLinears, a.ClippingBoundsOptional); err != nil { + return err + } + } + for _, suffix := range []string{".self_attn.q_norm.weight", ".self_attn.k_norm.weight"} { + if err := requireDense(tensors, layer+suffix, []int32{int32(a.HeadDim)}); err != nil { + return err + } + } + for _, suffix := range []string{".input_layernorm.weight", ".post_attention_layernorm.weight", ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight"} { + if err := requireDense(tensors, layer+suffix, []int32{int32(a.HiddenSize)}); err != nil { + return err + } + } + } + if a.Standardize { + if err := requireDense(tensors, prefix+"vision_tower.std_bias", []int32{int32(a.HiddenSize)}); err != nil { + return err + } + if err := requireDense(tensors, prefix+"vision_tower.std_scale", []int32{int32(a.HiddenSize)}); err != nil { + return err + } + } + for _, p := range []string{"embed_vision.embedding_projection", "model.embed_vision.embedding_projection"} { + if linearPresent(tensors, p, mode) { + return requireLinearDescriptor(tensors, p, []int32{int32(a.TextHiddenSize), int32(a.HiddenSize)}, mode, cfg) + } + } + return fmt.Errorf("missing embed_vision.embedding_projection weight") +} + +func validateNames(a Architecture, names []string, mode inventoryMode) error { + prefix, ok := visionPrefix(names, mode) + if !ok { + return fmt.Errorf("missing vision_tower.patch_embedder.input_proj weight") + } + require := func(name string) error { + if !slices.Contains(names, name) { + return fmt.Errorf("missing %s", name) + } + return nil + } + if err := require(prefix + "vision_tower.patch_embedder.position_embedding_table"); err != nil { + return err + } + for i := range a.NumHiddenLayers { + layer := fmt.Sprintf("%svision_tower.encoder.layers.%d", prefix, i) + for _, p := range []string{".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.o_proj", ".mlp.gate_proj", ".mlp.up_proj", ".mlp.down_proj"} { + if !linearNamePresent(names, layer+p, mode) { + return fmt.Errorf("missing %s weight", layer+p) + } + if err := validateClipNames(names, layer+p, a.UseClippedLinears, a.ClippingBoundsOptional); err != nil { + return err + } + } + for _, n := range []string{".self_attn.q_norm.weight", ".self_attn.k_norm.weight", ".input_layernorm.weight", ".post_attention_layernorm.weight", ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight"} { + if err := require(layer + n); err != nil { + return err + } + } + } + if a.Standardize { + if err := require(prefix + "vision_tower.std_bias"); err != nil { + return err + } + if err := require(prefix + "vision_tower.std_scale"); err != nil { + return err + } + } + for _, p := range []string{"embed_vision.embedding_projection", "model.embed_vision.embedding_projection"} { + if linearNamePresent(names, p, mode) { + return nil + } + } + return fmt.Errorf("missing embed_vision.embedding_projection weight") +} + +func visionPrefix(names []string, mode inventoryMode) (string, bool) { + for _, p := range []string{"", "model."} { + if linearNamePresent(names, p+"vision_tower.patch_embedder.input_proj", mode) { + return p, true + } + } + return "", false +} + +func linearNamePresent(names []string, path string, mode inventoryMode) bool { + for _, base := range []string{path, path + ".linear"} { + if slices.Contains(names, base+".weight") { + return true + } + if mode == sourceMode && slices.Contains(names, base+".weight_packed") && slices.Contains(names, base+".weight_scale") && slices.Contains(names, base+".weight_global_scale") { + return true + } + } + return false +} + +func linearPresent(t map[string]TensorDescriptor, p string, mode inventoryMode) bool { + names := make([]string, 0, len(t)) + for n := range t { + names = append(names, n) + } + return linearNamePresent(names, p, mode) +} + +func requireDense(t map[string]TensorDescriptor, name string, shape []int32) error { + d, ok := t[name] + if !ok { + return fmt.Errorf("missing %s", name) + } + if !isFloat(d.Dtype) { + return fmt.Errorf("%s dtype %s is not floating point", name, d.Dtype) + } + if !slices.Equal(d.Shape, shape) { + return fmt.Errorf("%s shape %v, want %v", name, d.Shape, shape) + } + return nil +} + +func requireLinearDescriptor(t map[string]TensorDescriptor, path string, logical []int32, mode inventoryMode, cfg ConfigFile) error { + for _, base := range []string{path, path + ".linear"} { + if d, ok := t[base+".weight"]; ok { + if isFloat(d.Dtype) { + if hasProducerCompanion(t, base) { + return fmt.Errorf("cross-producer companions for dense %s.weight", base) + } + if !slices.Equal(d.Shape, logical) { + return fmt.Errorf("%s.weight shape %v, want %v", base, d.Shape, logical) + } + return nil + } + if mode == sourceMode { + return validatePackedSource(t, base, d, logical, cfg, false) + } + return validateNormalizedQuant(t, base, d, logical, mode, cfg) + } + if mode == sourceMode { + if d, ok := t[base+".weight_packed"]; ok { + return validatePackedSource(t, base, d, logical, cfg, true) + } + } + } + return fmt.Errorf("missing %s weight", path) +} + +func validatePackedSource(t map[string]TensorDescriptor, base string, weight TensorDescriptor, logical []int32, cfg ConfigFile, compressed bool) error { + if compressed { + if !strings.EqualFold(weight.Dtype, "U8") || !packedWeightShape(weight.Shape, logical, 2) { + return fmt.Errorf("invalid compressed NVFP4 weight %s", base) + } + scale, ok := t[base+".weight_scale"] + if !ok { + return fmt.Errorf("incomplete packed weight %s.weight_packed", base) + } + global, ok := t[base+".weight_global_scale"] + if !ok { + return fmt.Errorf("incomplete packed weight %s.weight_packed", base) + } + if !packedScaleShape(scale, logical, 16) || !isE4M3(scale.Dtype) { + return fmt.Errorf("invalid packed scale %s.weight_scale", base) + } + if !isScalar(global.Shape) || !strings.EqualFold(global.Dtype, "F32") { + return fmt.Errorf("invalid packed global scale %s.weight_global_scale", base) + } + for _, suffix := range []string{".scales", ".biases", ".weight_scale_2"} { + if _, cross := t[base+suffix]; cross { + return fmt.Errorf("cross-producer packed companion %s%s", base, suffix) + } + } + return nil + } + if scale, ok := t[base+".scales"]; ok { + for _, suffix := range []string{".weight_scale", ".weight_global_scale", ".weight_scale_2"} { + if _, cross := t[base+suffix]; cross { + return fmt.Errorf("cross-producer packed companion %s%s", base, suffix) + } + } + q := sourceQuant(cfg) + if q.Bits <= 0 || q.GroupSize <= 0 { + return fmt.Errorf("missing MLX packed quantization contract for %s", base) + } + if !strings.EqualFold(weight.Dtype, "U32") || !packedWeightShape(weight.Shape, logical, int32(32/q.Bits)) || !packedScaleShape(scale, logical, int32(q.GroupSize)) || !isScaleDtype(scale.Dtype) { + return fmt.Errorf("invalid MLX packed descriptors for %s", base) + } + if bias, present := t[base+".biases"]; present && (!isFloat(bias.Dtype) || !slices.Equal(bias.Shape, scale.Shape)) { + return fmt.Errorf("invalid MLX packed bias for %s", base) + } + return nil + } + if scale, ok := t[base+".weight_scale"]; ok { + for _, suffix := range []string{".scales", ".biases", ".weight_global_scale"} { + if _, cross := t[base+suffix]; cross { + return fmt.Errorf("cross-producer packed companion %s%s", base, suffix) + } + } + if !strings.EqualFold(weight.Dtype, "U8") || !packedWeightShape(weight.Shape, logical, 2) || !packedScaleShape(scale, logical, 16) || !isE4M3(scale.Dtype) { + return fmt.Errorf("invalid ModelOpt NVFP4 descriptors for %s", base) + } + if g, ok := t[base+".weight_scale_2"]; ok && (!isScalar(g.Shape) || !strings.EqualFold(g.Dtype, "F32")) { + return fmt.Errorf("invalid ModelOpt global scale %s", base) + } + return nil + } + return fmt.Errorf("packed source weight %s has no recognized companions", base) +} + +func validateNormalizedQuant(t map[string]TensorDescriptor, base string, weight TensorDescriptor, logical []int32, mode inventoryMode, cfg ConfigFile) error { + for _, suffix := range []string{".scales", ".weight_packed", ".weight_global_scale", ".weight_scale_2"} { + if _, exists := t[base+suffix]; exists { + return fmt.Errorf("source-only packed companion %s%s in normalized inventory", base, suffix) + } + } + if mode == installedMode { + if _, exists := t[base+".weight_scale"]; exists { + return fmt.Errorf("source-only packed companion %s.weight_scale in installed inventory", base) + } + } + scaleName := base + ".weight.scale" + biasName := base + ".weight.bias" + if mode == runtimeMode { + scaleName = base + ".weight_scale" + biasName = base + ".weight_qbias" + } + scale, ok := t[scaleName] + if !ok { + return fmt.Errorf("missing or invalid normalized scale %s", scaleName) + } + if strings.EqualFold(scale.Dtype, "U8") { + if !strings.EqualFold(weight.Dtype, "U32") || !packedWeightShape(weight.Shape, logical, 8) || !packedScaleShape(scale, logical, 16) { + return fmt.Errorf("invalid normalized NVFP4 descriptors for %s", base) + } + if global, present := t[base+".weight.global_scale"]; present && (!isScalar(global.Shape) || !strings.EqualFold(global.Dtype, "F32")) { + return fmt.Errorf("invalid normalized global scale for %s", base) + } + if _, present := t[biasName]; present { + return fmt.Errorf("unexpected affine bias for normalized NVFP4 %s", base) + } + return nil + } + bits, groupSize, ok := inferAffineContract(weight, scale, logical) + q := sourceQuant(cfg) + if q.Bits > 0 && q.Bits != bits { + ok = false + } + if q.GroupSize > 0 && q.GroupSize != groupSize { + ok = false + } + if weight.GroupSize > 0 && weight.GroupSize != groupSize { + ok = false + } + if !ok || !isFloat(scale.Dtype) || !strings.EqualFold(weight.Dtype, "U32") { + return fmt.Errorf("invalid normalized affine descriptors for %s", base) + } + if bias, present := t[biasName]; present && (!isFloat(bias.Dtype) || !slices.Equal(bias.Shape, scale.Shape)) { + return fmt.Errorf("invalid normalized affine bias for %s", base) + } + if _, present := t[base+".weight.global_scale"]; present { + return fmt.Errorf("unexpected global scale for normalized affine %s", base) + } + return nil +} + +func inferAffineContract(weight, scale TensorDescriptor, logical []int32) (bits, groupSize int, ok bool) { + if len(weight.Shape) != 2 || len(scale.Shape) != 2 || len(logical) != 2 || weight.Shape[0] != logical[0] || scale.Shape[0] != logical[0] || weight.Shape[1] <= 0 || scale.Shape[1] <= 0 { + return 0, 0, false + } + if logical[1]%weight.Shape[1] != 0 || logical[1]%scale.Shape[1] != 0 { + return 0, 0, false + } + perWord := logical[1] / weight.Shape[1] + if perWord <= 0 || 32%perWord != 0 { + return 0, 0, false + } + bits = int(32 / perWord) + groupSize = int(logical[1] / scale.Shape[1]) + if (bits != 4 && bits != 8) || groupSize <= 0 { + return 0, 0, false + } + return bits, groupSize, true +} + +func sourceQuant(cfg ConfigFile) Quantization { + if cfg.QuantizationConfig.Bits != 0 { + return cfg.QuantizationConfig + } + return cfg.Quantization +} + +func hasProducerCompanion(t map[string]TensorDescriptor, b string) bool { + for _, s := range []string{".scales", ".weight_scale", ".weight_global_scale", ".weight_scale_2"} { + if _, ok := t[b+s]; ok { + return true + } + } + return false +} + +func packedWeightShape(got, logical []int32, perWord int32) bool { + return len(got) == 2 && len(logical) == 2 && got[0] == logical[0] && logical[1]%perWord == 0 && got[1] == logical[1]/perWord +} + +func packedScaleShape(d TensorDescriptor, logical []int32, group int32) bool { + return len(d.Shape) == 2 && len(logical) == 2 && d.Shape[0] == logical[0] && logical[1]%group == 0 && d.Shape[1] == logical[1]/group +} + +func isScaleDtype(d string) bool { + switch strings.ToUpper(d) { + case "U8", "F8_E4M3", "F8_E4M3FN", "BF16", "F16", "F32": + return true + } + return false +} + +func isE4M3(d string) bool { + switch strings.ToUpper(d) { + case "F8_E4M3", "F8_E4M3FN": + return true + } + return false +} + +func isFloat(d string) bool { + switch strings.ToUpper(d) { + case "BF16", "F16", "F32", "BFLOAT16", "FLOAT16", "FLOAT32": + return true + } + return false +} + +func isScalar(s []int32) bool { return len(s) == 0 || (len(s) == 1 && s[0] == 1) } + +func validateClipNames(names []string, path string, enabled, optional bool) error { + present := 0 + for _, s := range []string{".input_min", ".input_max", ".output_min", ".output_max"} { + if slices.Contains(names, path+s) { + present++ + } + } + if present != 0 && present != 4 { + return fmt.Errorf("incomplete clipping tensors for %s", path) + } + if !enabled && present != 0 { + return fmt.Errorf("unexpected clipping tensors for %s", path) + } + if enabled && present == 0 && !optional { + return fmt.Errorf("missing clipping tensors for %s", path) + } + return nil +} + +func validateClipDescriptors(t map[string]TensorDescriptor, path string, enabled, optional bool) error { + names := make([]string, 0, len(t)) + for n := range t { + names = append(names, n) + } + if err := validateClipNames(names, path, enabled, optional); err != nil { + return err + } + for _, s := range []string{".input_min", ".input_max", ".output_min", ".output_max"} { + if d, ok := t[path+s]; ok { + if !isFloat(d.Dtype) || !isScalar(d.Shape) { + return fmt.Errorf("invalid clipping tensor %s%s", path, s) + } + } + } + return nil +} diff --git a/x/models/gemma4/metadata/vision_test.go b/x/models/gemma4/metadata/vision_test.go new file mode 100644 index 00000000000..ff4de346d1e --- /dev/null +++ b/x/models/gemma4/metadata/vision_test.go @@ -0,0 +1,388 @@ +package metadata + +import ( + "fmt" + "math" + "slices" + "strconv" + "strings" + "testing" +) + +func completeVisionTensorNames(layers int, standardize bool) []string { + names := []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.position_embedding_table", + "model.embed_vision.embedding_projection.weight", + } + for i := range layers { + layer := "model.vision_tower.encoder.layers." + strconv.Itoa(i) + for _, projection := range []string{ + ".self_attn.q_proj.linear.weight", ".self_attn.k_proj.linear.weight", + ".self_attn.v_proj.linear.weight", ".self_attn.o_proj.linear.weight", + ".mlp.gate_proj.linear.weight", ".mlp.up_proj.linear.weight", ".mlp.down_proj.linear.weight", + } { + names = append(names, layer+projection) + } + for _, norm := range []string{ + ".self_attn.q_norm.weight", ".self_attn.k_norm.weight", + ".input_layernorm.weight", ".post_attention_layernorm.weight", + ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight", + } { + names = append(names, layer+norm) + } + } + if standardize { + names = append(names, "model.vision_tower.std_bias", "model.vision_tower.std_scale") + } + return names +} + +func TestValidateVisionTensors(t *testing.T) { + cfg := ConfigFile{VisionConfig: &VisionConfig{NumHiddenLayers: 2, Standardize: true}} + names := completeVisionTensorNames(2, true) + if err := ValidateVisionTensors(cfg, names); err != nil { + t.Fatalf("ValidateVisionTensors() error = %v", err) + } + + for _, missing := range []string{ + "model.vision_tower.patch_embedder.position_embedding_table", + "model.vision_tower.encoder.layers.1.self_attn.q_norm.weight", + "model.vision_tower.encoder.layers.1.mlp.down_proj.linear.weight", + "model.vision_tower.std_scale", + "model.embed_vision.embedding_projection.weight", + } { + partial := append([]string(nil), names...) + for i, name := range partial { + if name == missing { + partial = append(partial[:i], partial[i+1:]...) + break + } + } + if err := ValidateVisionTensors(cfg, partial); err == nil { + t.Fatalf("missing %s: error = %v", missing, err) + } + } +} + +func TestValidateVisionTensorsBoundsLayerCount(t *testing.T) { + if err := ValidateVisionTensors( + ConfigFile{VisionConfig: &VisionConfig{NumHiddenLayers: MaxVisionLayers}}, + completeVisionTensorNames(MaxVisionLayers, false), + ); err != nil { + t.Fatalf("exact layer limit error = %v", err) + } + + for _, layers := range []int{-1, MaxVisionLayers + 1} { + err := ValidateVisionTensors( + ConfigFile{VisionConfig: &VisionConfig{NumHiddenLayers: layers}}, + []string{ + "model.vision_tower.patch_embedder.input_proj.weight", + "model.vision_tower.patch_embedder.position_embedding_table", + }, + ) + if err == nil || !strings.Contains(err.Error(), "num_hidden_layers") { + t.Fatalf("layer count %d error = %v", layers, err) + } + } +} + +func executableVisionConfig() ConfigFile { + return ConfigFile{ + TextConfig: TextConfig{HiddenSize: 24}, + VisionConfig: &VisionConfig{ + HiddenSize: 16, IntermediateSize: 32, NumHiddenLayers: 1, + NumAttentionHeads: 1, NumKeyValueHeads: 1, HeadDim: 16, + RMSNormEps: 1e-6, DefaultOutputLength: 1, PatchSize: 4, + PositionEmbeddingSize: 16, PoolingKernelSize: 1, + }, + } +} + +func executableVisionDescriptors() map[string]TensorDescriptor { + return visionDescriptorsForGeometry(16, 32, 24, 4, 16, 16) +} + +func visionDescriptorsForGeometry(hidden, intermediate, textHidden, patch, positions, headDim int32) map[string]TensorDescriptor { + d := map[string]TensorDescriptor{ + "model.vision_tower.patch_embedder.input_proj.weight": {Dtype: "F32", Shape: []int32{hidden, 3 * patch * patch}}, + "model.vision_tower.patch_embedder.position_embedding_table": {Dtype: "F32", Shape: []int32{2, positions, hidden}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "F32", Shape: []int32{textHidden, hidden}}, + } + layer := "model.vision_tower.encoder.layers.0" + for _, suffix := range []string{".self_attn.q_proj.linear.weight", ".self_attn.k_proj.linear.weight", ".self_attn.v_proj.linear.weight", ".self_attn.o_proj.linear.weight"} { + d[layer+suffix] = TensorDescriptor{Dtype: "F32", Shape: []int32{hidden, hidden}} + } + for _, suffix := range []string{".mlp.gate_proj.linear.weight", ".mlp.up_proj.linear.weight"} { + d[layer+suffix] = TensorDescriptor{Dtype: "F32", Shape: []int32{intermediate, hidden}} + } + d[layer+".mlp.down_proj.linear.weight"] = TensorDescriptor{Dtype: "F32", Shape: []int32{hidden, intermediate}} + for _, suffix := range []string{".self_attn.q_norm.weight", ".self_attn.k_norm.weight"} { + d[layer+suffix] = TensorDescriptor{Dtype: "F32", Shape: []int32{headDim}} + } + for _, suffix := range []string{".input_layernorm.weight", ".post_attention_layernorm.weight", ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight"} { + d[layer+suffix] = TensorDescriptor{Dtype: "F32", Shape: []int32{hidden}} + } + return d +} + +func cloneDescriptors(in map[string]TensorDescriptor) map[string]TensorDescriptor { + out := make(map[string]TensorDescriptor, len(in)) + for name, descriptor := range in { + descriptor.Shape = slices.Clone(descriptor.Shape) + out[name] = descriptor + } + return out +} + +func TestVisionInventoryProducerAndNormalizationMatrix(t *testing.T) { + const target = "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear" + dense := executableVisionDescriptors() + cfg := executableVisionConfig() + if err := ValidateVisionSourceInventory(cfg, dense); err != nil { + t.Fatalf("dense source error = %v", err) + } + if err := ValidateVisionInstalledInventory(cfg, dense); err != nil { + t.Fatalf("dense installed error = %v", err) + } + + tests := []struct { + name string + config func(*ConfigFile) + add map[string]TensorDescriptor + required string + }{ + { + name: "compressed tensors NVFP4", + add: map[string]TensorDescriptor{ + target + ".weight_packed": {Dtype: "U8", Shape: []int32{16, 8}}, + target + ".weight_scale": {Dtype: "F8_E4M3", Shape: []int32{16, 1}}, + target + ".weight_global_scale": {Dtype: "F32", Shape: nil}, + }, + required: target + ".weight_global_scale", + }, + { + name: "MLX packed", + config: func(cfg *ConfigFile) { + cfg.QuantizationConfig = Quantization{Bits: 4, GroupSize: 16, Mode: "affine"} + }, + add: map[string]TensorDescriptor{ + target + ".weight": {Dtype: "U32", Shape: []int32{16, 2}}, + target + ".scales": {Dtype: "F16", Shape: []int32{16, 1}}, + }, + required: target + ".scales", + }, + { + name: "ModelOpt NVFP4", + add: map[string]TensorDescriptor{ + target + ".weight": {Dtype: "U8", Shape: []int32{16, 8}}, + target + ".weight_scale": {Dtype: "F8_E4M3", Shape: []int32{16, 1}}, + target + ".weight_scale_2": {Dtype: "F32", Shape: []int32{1}}, + }, + required: target + ".weight_scale", + }, + } + + for _, tt := range []struct { + name string + weight string + scale string + global string + }{ + {name: "compressed scale", weight: "weight_packed", scale: "weight_scale", global: "weight_global_scale"}, + {name: "ModelOpt scale", weight: "weight", scale: "weight_scale", global: "weight_scale_2"}, + } { + t.Run("wrong dtype "+tt.name, func(t *testing.T) { + for _, field := range []string{"scale", "global"} { + inventory := cloneDescriptors(dense) + delete(inventory, target+".weight") + inventory[target+"."+tt.weight] = TensorDescriptor{Dtype: "U8", Shape: []int32{16, 8}} + inventory[target+"."+tt.scale] = TensorDescriptor{Dtype: "F8_E4M3", Shape: []int32{16, 1}} + inventory[target+"."+tt.global] = TensorDescriptor{Dtype: "F32", Shape: nil} + bad := inventory[target+"."+map[string]string{"scale": tt.scale, "global": tt.global}[field]] + bad.Dtype = "F16" + inventory[target+"."+map[string]string{"scale": tt.scale, "global": tt.global}[field]] = bad + if err := ValidateVisionSourceInventory(cfg, inventory); err == nil { + t.Fatalf("%s %s with F16 dtype accepted", tt.name, field) + } + } + }) + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + inventory := cloneDescriptors(dense) + delete(inventory, target+".weight") + for name, descriptor := range tt.add { + inventory[name] = descriptor + } + localCfg := cfg + if tt.config != nil { + tt.config(&localCfg) + } + if err := ValidateVisionSourceInventory(localCfg, inventory); err != nil { + t.Fatalf("valid packed source error = %v", err) + } + if err := ValidateVisionInstalledInventory(localCfg, inventory); err == nil { + t.Fatal("source-only packed form accepted as installed") + } + partial := cloneDescriptors(inventory) + delete(partial, tt.required) + if err := ValidateVisionSourceInventory(localCfg, partial); err == nil { + t.Fatalf("missing companion %s accepted", tt.required) + } + }) + } + + normalized := cloneDescriptors(dense) + normalized[target+".weight"] = TensorDescriptor{Dtype: "U32", Shape: []int32{16, 2}} + normalized[target+".weight.scale"] = TensorDescriptor{Dtype: "U8", Shape: []int32{16, 1}} + normalized[target+".weight.global_scale"] = TensorDescriptor{Dtype: "F32", Shape: nil} + if err := ValidateVisionInstalledInventory(cfg, normalized); err != nil { + t.Fatalf("normalized installed error = %v", err) + } + runtime := cloneDescriptors(normalized) + delete(runtime, target+".weight.scale") + runtime[target+".weight_scale"] = TensorDescriptor{Dtype: "U8", Shape: []int32{16, 1}} + if err := ValidateVisionRuntimeInventory(cfg, runtime); err != nil { + t.Fatalf("normalized runtime error = %v", err) + } + for _, missing := range []string{target + ".weight.scale"} { + partial := cloneDescriptors(normalized) + delete(partial, missing) + if err := ValidateVisionInstalledInventory(cfg, partial); err == nil { + t.Fatalf("normalized installed inventory missing %s accepted", missing) + } + } + cross := cloneDescriptors(normalized) + cross[target+".scales"] = TensorDescriptor{Dtype: "F16", Shape: []int32{16, 1}} + if err := ValidateVisionInstalledInventory(cfg, cross); err == nil { + t.Fatal("cross-producer companion accepted") + } + + affineCfg := cfg + affineCfg.QuantizationConfig = Quantization{Bits: 4, GroupSize: 16, Mode: "affine"} + affine := cloneDescriptors(dense) + affine[target+".weight"] = TensorDescriptor{Dtype: "U32", Shape: []int32{16, 2}} + affine[target+".weight.scale"] = TensorDescriptor{Dtype: "F16", Shape: []int32{16, 1}} + affine[target+".weight.bias"] = TensorDescriptor{Dtype: "F16", Shape: []int32{16, 1}} + if err := ValidateVisionInstalledInventory(affineCfg, affine); err != nil { + t.Fatalf("normalized affine installed error = %v", err) + } + affineRuntime := cloneDescriptors(affine) + delete(affineRuntime, target+".weight.scale") + delete(affineRuntime, target+".weight.bias") + affineRuntime[target+".weight_scale"] = TensorDescriptor{Dtype: "F16", Shape: []int32{16, 1}} + affineRuntime[target+".weight_qbias"] = TensorDescriptor{Dtype: "F16", Shape: []int32{16, 1}} + if err := ValidateVisionRuntimeInventory(affineCfg, affineRuntime); err != nil { + t.Fatalf("normalized affine runtime error = %v", err) + } +} + +func TestVisionInventoryReleasedUnequalHeadGeometry(t *testing.T) { + cfg := ConfigFile{ + TextConfig: TextConfig{HiddenSize: 2560}, + VisionConfig: &VisionConfig{ + HiddenSize: 768, IntermediateSize: 3072, NumHiddenLayers: 1, + NumAttentionHeads: 12, NumKeyValueHeads: 12, HeadDim: 64, + RMSNormEps: 1e-6, DefaultOutputLength: 280, PatchSize: 16, + PositionEmbeddingSize: 10240, PoolingKernelSize: 3, + }, + } + descriptors := visionDescriptorsForGeometry(768, 3072, 2560, 16, 10240, 64) + for name, validate := range map[string]func(ConfigFile, map[string]TensorDescriptor) error{ + "source": ValidateVisionSourceInventory, + "installed": ValidateVisionInstalledInventory, + "runtime": ValidateVisionRuntimeInventory, + } { + t.Run(name, func(t *testing.T) { + if err := validate(cfg, descriptors); err != nil { + t.Fatalf("released-compatible unequal hidden/head geometry: %v", err) + } + bad := cloneDescriptors(descriptors) + qnorm := "model.vision_tower.encoder.layers.0.self_attn.q_norm.weight" + bad[qnorm] = TensorDescriptor{Dtype: "F32", Shape: []int32{768}} + if err := validate(cfg, bad); err == nil { + t.Fatal("hidden-sized q_norm accepted") + } + }) + } +} + +func TestVisionInventoryRejectsMissingOrZeroTextWidth(t *testing.T) { + for _, hidden := range []int{0, -1, MaxVisionHidden + 1} { + cfg := executableVisionConfig() + cfg.TextConfig.HiddenSize = hidden + for name, validate := range map[string]func(ConfigFile, map[string]TensorDescriptor) error{ + "source": ValidateVisionSourceInventory, + "installed": ValidateVisionInstalledInventory, + "runtime": ValidateVisionRuntimeInventory, + } { + t.Run(fmt.Sprintf("%s/%d", name, hidden), func(t *testing.T) { + if err := validate(cfg, executableVisionDescriptors()); err == nil || !strings.Contains(err.Error(), "text hidden_size") { + t.Fatalf("text hidden_size %d error = %v", hidden, err) + } + }) + } + } +} + +func TestVisionInventoryDescriptorAndConfigMatrix(t *testing.T) { + cfg := executableVisionConfig() + valid := executableVisionDescriptors() + for _, mutate := range []func(map[string]TensorDescriptor){ + func(d map[string]TensorDescriptor) { + x := d["model.vision_tower.patch_embedder.input_proj.weight"] + x.Shape = []int32{16, 47} + d["model.vision_tower.patch_embedder.input_proj.weight"] = x + }, + func(d map[string]TensorDescriptor) { + x := d["model.vision_tower.patch_embedder.position_embedding_table"] + x.Dtype = "U8" + d["model.vision_tower.patch_embedder.position_embedding_table"] = x + }, + func(d map[string]TensorDescriptor) { + x := d["model.vision_tower.encoder.layers.0.input_layernorm.weight"] + x.Shape = []int32{1, 16} + d["model.vision_tower.encoder.layers.0.input_layernorm.weight"] = x + }, + } { + malformed := cloneDescriptors(valid) + mutate(malformed) + if err := ValidateVisionSourceInventory(cfg, malformed); err == nil { + t.Fatal("malformed complete-name inventory accepted") + } + } + + clippedCfg := cfg + clippedVision := *cfg.VisionConfig + clippedVision.UseClippedLinears = true + clippedCfg.VisionConfig = &clippedVision + partialClip := cloneDescriptors(valid) + partialClip["model.vision_tower.encoder.layers.0.self_attn.q_proj.input_min"] = TensorDescriptor{Dtype: "F32", Shape: nil} + if err := ValidateVisionInstalledInventory(clippedCfg, partialClip); err == nil { + t.Fatal("partial clipping bounds accepted") + } + + invalid := []ConfigFile{} + for _, change := range []func(*VisionConfig){ + func(v *VisionConfig) { v.HiddenSize = -1 }, + func(v *VisionConfig) { v.HiddenSize = MaxVisionHidden + 1 }, + func(v *VisionConfig) { v.NumAttentionHeads = 3 }, + func(v *VisionConfig) { v.RMSNormEps = math.NaN() }, + func(v *VisionConfig) { v.RopeParameters.RopeTheta = math.Inf(1) }, + func(v *VisionConfig) { v.RMSNormEps = math.SmallestNonzeroFloat64 }, + func(v *VisionConfig) { v.RopeParameters.RopeTheta = math.MaxFloat64 }, + func(v *VisionConfig) { v.DefaultOutputLength = MaxVisionSoftTokens + 1 }, + } { + bad := cfg + vision := *cfg.VisionConfig + change(&vision) + bad.VisionConfig = &vision + invalid = append(invalid, bad) + } + for _, bad := range invalid { + if _, err := ProjectVisionArchitecture(bad); err == nil { + t.Fatalf("invalid architecture accepted: %+v", bad.VisionConfig) + } + } +} diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go index 9d6ca71cc7a..de1e7fdd13f 100644 --- a/x/models/gemma4/vision.go +++ b/x/models/gemma4/vision.go @@ -1,16 +1,21 @@ package gemma4 // Portions of the Gemma 4 vision preprocessing and embedding flow are adapted -// from MLX-VLM's MIT-licensed Gemma 4 implementation. +// from MLX-VLM's MIT-licensed Gemma 4 implementation. See +// docs/third-party/mlx-vlm.md for the pinned source revision and license. import ( "bytes" + "context" "encoding/json" "errors" "fmt" "image" + "image/color" + stddraw "image/draw" _ "image/jpeg" _ "image/png" + "io" "math" xdraw "golang.org/x/image/draw" @@ -19,6 +24,7 @@ import ( "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/mlxrunner/model" "github.com/ollama/ollama/x/mlxrunner/model/base" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" "github.com/ollama/ollama/x/models/nn" ) @@ -156,46 +162,27 @@ func parseVisionConfig(configData []byte) (*VisionConfig, error) { if wrapped.VisionConfig == nil { return nil, nil } - cfg := *wrapped.VisionConfig - if cfg.HiddenSize == 0 { - cfg.HiddenSize = 768 - } - if cfg.IntermediateSize == 0 { - cfg.IntermediateSize = 3072 - } - if cfg.NumHiddenLayers == 0 { - cfg.NumHiddenLayers = 16 - } - if cfg.NumAttentionHeads == 0 { - cfg.NumAttentionHeads = 12 - } - if cfg.NumKeyValueHeads == 0 { - cfg.NumKeyValueHeads = cfg.NumAttentionHeads - } - if cfg.HeadDim == 0 { - cfg.HeadDim = 64 - } - if cfg.RMSNormEps == 0 { - cfg.RMSNormEps = 1e-6 - } - if cfg.DefaultOutputLength == 0 { - cfg.DefaultOutputLength = 280 - } - if cfg.PatchSize == 0 { - cfg.PatchSize = 16 - } - if cfg.PositionEmbeddingSize == 0 { - cfg.PositionEmbeddingSize = 10240 + var metadataConfig gemma4metadata.ConfigFile + if err := json.Unmarshal(configData, &metadataConfig); err != nil { + return nil, fmt.Errorf("parse vision metadata: %w", err) } - if cfg.PoolingKernelSize == 0 { - cfg.PoolingKernelSize = 3 - } - if cfg.RopeParameters.RopeTheta == 0 { - cfg.RopeParameters.RopeTheta = 100 - } - if err := validateVisionConfig(&cfg); err != nil { + architecture, err := gemma4metadata.ProjectVisionArchitecture(metadataConfig) + if err != nil { return nil, err } + cfg := *wrapped.VisionConfig + cfg.HiddenSize = int32(architecture.HiddenSize) + cfg.IntermediateSize = int32(architecture.IntermediateSize) + cfg.NumHiddenLayers = int32(architecture.NumHiddenLayers) + cfg.NumAttentionHeads = int32(architecture.NumAttentionHeads) + cfg.NumKeyValueHeads = int32(architecture.NumKeyValueHeads) + cfg.HeadDim = int32(architecture.HeadDim) + cfg.RMSNormEps = float32(architecture.RMSNormEps) + cfg.DefaultOutputLength = int32(architecture.DefaultOutputLength) + cfg.PatchSize = int32(architecture.PatchSize) + cfg.PositionEmbeddingSize = int32(architecture.PositionEmbeddingSize) + cfg.PoolingKernelSize = int32(architecture.PoolingKernelSize) + cfg.RopeParameters.RopeTheta = float32(architecture.RopeTheta) return &cfg, nil } @@ -203,61 +190,10 @@ func validateVisionConfig(cfg *VisionConfig) error { if cfg == nil { return errors.New("missing Gemma4 vision config") } - positiveBounded := func(name string, value, limit int32) error { - if value <= 0 || value > limit { - return fmt.Errorf("invalid Gemma4 vision %s %d (limit %d)", name, value, limit) - } - return nil - } - for _, field := range []struct { - name string - value, limit int32 - }{ - {"hidden_size", cfg.HiddenSize, maxGemma4VisionHiddenSize}, - {"intermediate_size", cfg.IntermediateSize, maxGemma4VisionIntermediate}, - {"num_hidden_layers", cfg.NumHiddenLayers, maxGemma4VisionLayers}, - {"num_attention_heads", cfg.NumAttentionHeads, maxGemma4VisionHeads}, - {"num_key_value_heads", cfg.NumKeyValueHeads, maxGemma4VisionHeads}, - {"head_dim", cfg.HeadDim, maxGemma4VisionHeadDim}, - {"default_output_length", cfg.DefaultOutputLength, maxGemma4VisionSoftTokens}, - {"patch_size", cfg.PatchSize, maxGemma4ImageDimension}, - {"position_embedding_size", cfg.PositionEmbeddingSize, maxGemma4PositionTableEntries}, - {"pooling_kernel_size", cfg.PoolingKernelSize, maxGemma4ImageDimension}, - } { - if err := positiveBounded(field.name, field.value, field.limit); err != nil { - return err - } - } - if cfg.NumKeyValueHeads > cfg.NumAttentionHeads || cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { - return fmt.Errorf("invalid Gemma4 vision attention heads %d/%d", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) - } - attentionWidth := int64(cfg.NumAttentionHeads) * int64(cfg.HeadDim) - if attentionWidth != int64(cfg.HiddenSize) { - return fmt.Errorf("Gemma4 vision attention width %d does not match hidden_size %d", attentionWidth, cfg.HiddenSize) - } - if cfg.RMSNormEps <= 0 || math.IsNaN(float64(cfg.RMSNormEps)) || math.IsInf(float64(cfg.RMSNormEps), 0) { - return fmt.Errorf("invalid Gemma4 vision rms_norm_eps %v", cfg.RMSNormEps) - } - if cfg.RopeParameters.RopeTheta <= 0 || math.IsNaN(float64(cfg.RopeParameters.RopeTheta)) || math.IsInf(float64(cfg.RopeParameters.RopeTheta), 0) { - return fmt.Errorf("invalid Gemma4 vision rope_theta %v", cfg.RopeParameters.RopeTheta) - } - - if _, ok := checkedPositiveProduct(maxGemma4ResizePixels, - int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize), int64(cfg.PoolingKernelSize), int64(cfg.PatchSize), int64(cfg.PatchSize)); !ok { - return fmt.Errorf("Gemma4 vision resize budget exceeds limit %d", maxGemma4ResizePixels) - } - if _, ok := checkedPositiveProduct(maxGemma4PositionValues, - int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize), int64(cfg.PoolingKernelSize), int64(cfg.HeadDim)); !ok { - return fmt.Errorf("Gemma4 vision position allocation exceeds limit %d", maxGemma4PositionValues) - } - if _, ok := checkedPositiveProduct(maxGemma4ImageDimension, int64(cfg.PoolingKernelSize), int64(cfg.PatchSize)); !ok { - return fmt.Errorf("invalid Gemma4 pooled patch size (limit %d)", maxGemma4ImageDimension) - } - maxPatchSide, ok := checkedPositiveProduct(maxGemma4PositionTableEntries, int64(cfg.DefaultOutputLength), int64(cfg.PoolingKernelSize)) - if !ok || maxPatchSide > int64(cfg.PositionEmbeddingSize) { - return fmt.Errorf("Gemma4 vision patch side %d exceeds position table size %d", maxPatchSide, cfg.PositionEmbeddingSize) - } - return nil + // This validates vision-only call sites before the enclosing text model is + // available. Executable inventory validation supplies the real text width. + _, err := gemma4metadata.ProjectVisionArchitecture(metadataConfigFromVision(cfg, 1)) + return err } func checkedPositiveProduct(limit int64, values ...int64) (int64, bool) { @@ -303,11 +239,47 @@ func parseGemma4MediaTokens(data []byte, fallback gemma4MediaTokens) gemma4Media return fallback } -func hasGemma4VisionWeights(tensors map[string]*mlx.Array) bool { - return firstNonNil(tensors, - "vision_tower.patch_embedder.input_proj.weight", - "model.vision_tower.patch_embedder.input_proj.weight", - ) != nil +func metadataConfigFromVision(cfg *VisionConfig, textHidden int) gemma4metadata.ConfigFile { + if cfg == nil { + return gemma4metadata.ConfigFile{} + } + v := &gemma4metadata.VisionConfig{ + HiddenSize: int(cfg.HiddenSize), IntermediateSize: int(cfg.IntermediateSize), NumHiddenLayers: int(cfg.NumHiddenLayers), + NumAttentionHeads: int(cfg.NumAttentionHeads), NumKeyValueHeads: int(cfg.NumKeyValueHeads), HeadDim: int(cfg.HeadDim), + RMSNormEps: float64(cfg.RMSNormEps), DefaultOutputLength: int(cfg.DefaultOutputLength), PatchSize: int(cfg.PatchSize), + PositionEmbeddingSize: int(cfg.PositionEmbeddingSize), PoolingKernelSize: int(cfg.PoolingKernelSize), + UseClippedLinears: cfg.UseClippedLinears, Standardize: cfg.Standardize, + } + v.RopeParameters.RopeTheta = float64(cfg.RopeParameters.RopeTheta) + return gemma4metadata.ConfigFile{TextConfig: gemma4metadata.TextConfig{HiddenSize: textHidden}, VisionConfig: v} +} + +func validateGemma4VisionWeights(tensors map[string]*mlx.Array, cfg *VisionConfig, textHidden int, tq map[string]*model.TensorQuantInfo) (bool, error) { + if firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") == nil { + if firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight_packed", "model.vision_tower.patch_embedder.input_proj.weight_packed") != nil { + return false, fmt.Errorf("runtime contains source-only Gemma4 vision packed sentinel") + } + return false, nil + } + descriptors := make(map[string]gemma4metadata.TensorDescriptor, len(tensors)) + for name, tensor := range tensors { + if tensor != nil { + shape := make([]int32, tensor.NumDims()) + for i, d := range tensor.Dims() { + shape[i] = int32(d) + } + d := gemma4metadata.TensorDescriptor{Dtype: tensor.DType().String(), Shape: shape} + if q := tq[name]; q != nil { + d.QuantType = q.QuantType + d.GroupSize = q.GroupSize + } + descriptors[name] = d + } + } + if err := gemma4metadata.ValidateVisionRuntimeInventory(metadataConfigFromVision(cfg, textHidden), descriptors); err != nil { + return false, err + } + return true, nil } func resolveVisionPrefix(tensors map[string]*mlx.Array) string { @@ -451,9 +423,12 @@ func (m *MultimodalEmbedder) Forward(x *mlx.Array) *mlx.Array { // PrepareMedia implements the runner's media contract. Each image is expanded // in stream order and remains a separate cache-identity item. -func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { +func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { prepared := &base.PreparedRequest{} for source, seg := range segments { + if err := ctx.Err(); err != nil { + return nil, err + } if seg.Data == nil { prepared.Tokens = append(prepared.Tokens, seg.Tokens...) continue @@ -465,7 +440,7 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er return nil, fmt.Errorf("gemma4 does not support %s input", seg.Kind) } - img, err := preprocessGemma4Image(seg.Data, m.VisionConfig, int(m.VisionSoftTokens)) + img, err := preprocessGemma4Image(ctx, seg.Data, m.VisionConfig, int(m.VisionSoftTokens)) if err != nil { return nil, err } @@ -481,7 +456,7 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er geom := *img pixels := geom.Pixels geom.Pixels = nil - prepared.Items = append(prepared.Items, base.PreparedItem{ + item := base.PreparedItem{ Range: [2]int{start, len(prepared.Tokens)}, Source: source, MediaData: pixels, @@ -491,7 +466,14 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er ImageStart: imageStart, ImageEnd: imageEnd, }, - }) + } + if err := ctx.Err(); err != nil { + return nil, err + } + prepared.Items = append(prepared.Items, item) + } + if err := ctx.Err(); err != nil { + return nil, err } return prepared, nil } @@ -546,21 +528,24 @@ func gemma4PLETokens(tokens *mlx.Array, b *batch.Batch) *mlx.Array { return tokens } -func preprocessGemma4Image(data []byte, cfg *VisionConfig, maxSoftTokens int) (*gemma4ImageInput, error) { +func preprocessGemma4Image(ctx context.Context, data []byte, cfg *VisionConfig, maxSoftTokens int) (*gemma4ImageInput, error) { + if err := ctx.Err(); err != nil { + return nil, err + } if err := validateVisionConfig(cfg); err != nil { return nil, err } if err := validateGemma4ImageDataSize(len(data)); err != nil { return nil, err } - imageConfig, _, err := image.DecodeConfig(bytes.NewReader(data)) + imageConfig, _, err := decodeGemma4ImageConfig(ctx, data) if err != nil { return nil, fmt.Errorf("decode Gemma4 image config: %w", err) } if err := validateGemma4ImageDimensions(imageConfig.Width, imageConfig.Height); err != nil { return nil, err } - img, _, err := image.Decode(bytes.NewReader(data)) + img, _, err := decodeGemma4Image(ctx, data) if err != nil { return nil, fmt.Errorf("decode Gemma4 image: %w", err) } @@ -593,12 +578,13 @@ func preprocessGemma4Image(data []byte, cfg *VisionConfig, maxSoftTokens int) (* resized := img if targetW != width || targetH != height { - dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH)) - xdraw.CatmullRom.Scale(dst, dst.Bounds(), img, b, xdraw.Over, nil) - resized = dst + resized, err = resizeGemma4Image(ctx, img, b, targetW, targetH) + if err != nil { + return nil, err + } } - pixels, err := imageToCHWFloat32(resized) + pixels, err := imageToCHWFloat32Context(ctx, resized) if err != nil { return nil, err } @@ -626,6 +612,41 @@ func preprocessGemma4Image(data []byte, cfg *VisionConfig, maxSoftTokens int) (* }, nil } +type gemma4CancellationPanic struct{ err error } + +type cancellableDrawImage struct { + stddraw.Image + contextErr func() error +} + +func (dst cancellableDrawImage) Set(x, y int, c color.Color) { + if err := dst.contextErr(); err != nil { + panic(gemma4CancellationPanic{err: err}) + } + dst.Image.Set(x, y, c) +} + +func resizeGemma4Image(ctx context.Context, src image.Image, bounds image.Rectangle, width, height int) (resized image.Image, err error) { + if err := ctx.Err(); err != nil { + return nil, err + } + dst := image.NewRGBA(image.Rect(0, 0, width, height)) + defer func() { + if recovered := recover(); recovered != nil { + if canceled, ok := recovered.(gemma4CancellationPanic); ok { + resized, err = nil, canceled.err + return + } + panic(recovered) + } + }() + xdraw.CatmullRom.Scale(cancellableDrawImage{Image: dst, contextErr: ctx.Err}, dst.Bounds(), src, bounds, xdraw.Over, nil) + if err := ctx.Err(); err != nil { + return nil, err + } + return dst, nil +} + func validateGemma4ImageDataSize(size int) error { if size <= 0 { return errors.New("Gemma4 image is empty") @@ -636,6 +657,36 @@ func validateGemma4ImageDataSize(size int) error { return nil } +type contextReader struct { + contextErr func() error + r io.Reader +} + +func (r contextReader) Read(p []byte) (int, error) { + if err := r.contextErr(); err != nil { + return 0, err + } + return r.r.Read(p) +} + +func decodeGemma4ImageConfig(ctx context.Context, data []byte) (image.Config, string, error) { + r := func() io.Reader { return contextReader{contextErr: ctx.Err, r: bytes.NewReader(data)} } + cfg, format, err := image.DecodeConfig(r()) + if err != nil && ctx.Err() != nil { + return image.Config{}, "", ctx.Err() + } + return cfg, format, err +} + +func decodeGemma4Image(ctx context.Context, data []byte) (image.Image, string, error) { + r := func() io.Reader { return contextReader{contextErr: ctx.Err, r: bytes.NewReader(data)} } + img, format, err := image.Decode(r()) + if err != nil && ctx.Err() != nil { + return nil, "", ctx.Err() + } + return img, format, err +} + func validateGemma4ImageDimensions(width, height int) error { if width <= 0 || height <= 0 { return fmt.Errorf("invalid Gemma4 image dimensions %dx%d", width, height) @@ -707,6 +758,10 @@ func gemma4ResizeDimensions(width, height, patchSize, maxPatches, poolingKernelS } func imageToCHWFloat32(img image.Image) ([]float32, error) { + return imageToCHWFloat32Context(context.Background(), img) +} + +func imageToCHWFloat32Context(ctx context.Context, img image.Image) ([]float32, error) { bounds := img.Bounds() width, height := bounds.Dx(), bounds.Dy() if err := validateGemma4ImageDimensions(width, height); err != nil { @@ -717,9 +772,15 @@ func imageToCHWFloat32(img image.Image) ([]float32, error) { if values64 <= 0 || values64 > maxIntValue { return nil, errors.New("Gemma4 pixel allocation exceeds platform limits") } + if err := ctx.Err(); err != nil { + return nil, err + } plane := int(plane64) out := make([]float32, int(values64)) for y := range height { + if err := ctx.Err(); err != nil { + return nil, err + } for x := range width { r, g, blue, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA() i := y*width + x diff --git a/x/models/gemma4/vision_test.go b/x/models/gemma4/vision_test.go index f269fb750ad..6f8bf4bad79 100644 --- a/x/models/gemma4/vision_test.go +++ b/x/models/gemma4/vision_test.go @@ -2,6 +2,8 @@ package gemma4 import ( "bytes" + "context" + "errors" "image" "image/color" "image/png" @@ -9,6 +11,7 @@ import ( "slices" "strings" "testing" + "time" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/model/base" @@ -16,7 +19,7 @@ import ( func testVisionConfig(t *testing.T, outputLength int32) *VisionConfig { t.Helper() - cfg, err := parseVisionConfig([]byte(`{"vision_config":{"default_output_length":1}}`)) + cfg, err := parseVisionConfig([]byte(`{"text_config":{"hidden_size":1},"vision_config":{"default_output_length":1}}`)) if err != nil { t.Fatal(err) } @@ -28,7 +31,7 @@ func testVisionConfig(t *testing.T, outputLength int32) *VisionConfig { } func TestParseVisionConfigDefaults(t *testing.T) { - cfg, err := parseVisionConfig([]byte(`{"vision_config":{}}`)) + cfg, err := parseVisionConfig([]byte(`{"text_config":{"hidden_size":1},"vision_config":{}}`)) if err != nil { t.Fatalf("parseVisionConfig() error = %v", err) } @@ -62,7 +65,8 @@ func TestParseVisionConfigRejectsUnsafeDimensions(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - _, err := parseVisionConfig([]byte(tt.json)) + input := strings.Replace(tt.json, `{"vision_config"`, `{"text_config":{"hidden_size":1},"vision_config"`, 1) + _, err := parseVisionConfig([]byte(input)) if err == nil || !strings.Contains(err.Error(), tt.want) { t.Fatalf("parseVisionConfig() error = %v, want %q", err, tt.want) } @@ -70,6 +74,17 @@ func TestParseVisionConfigRejectsUnsafeDimensions(t *testing.T) { } } +func TestParseVisionConfigRejectsMissingOrZeroTextWidth(t *testing.T) { + for _, input := range []string{ + `{"vision_config":{}}`, + `{"text_config":{"hidden_size":0},"vision_config":{}}`, + } { + if _, err := parseVisionConfig([]byte(input)); err == nil || !strings.Contains(err.Error(), "text hidden_size") { + t.Fatalf("parseVisionConfig(%s) error = %v, want text hidden_size", input, err) + } + } +} + func TestVisionStandardizationTensorRequirements(t *testing.T) { cfg := testVisionConfig(t, 1) if err := validateVisionStandardizationTensors(cfg, false, false); err != nil { @@ -116,6 +131,25 @@ func TestGemma4ImageBounds(t *testing.T) { } } +func TestPreprocessGemma4ImageRejectsLimitsAndCancellation(t *testing.T) { + cfg := testVisionConfig(t, 1) + if _, err := preprocessGemma4Image(context.Background(), make([]byte, maxGemma4ImageBytes+1), cfg, 1); err == nil || !strings.Contains(err.Error(), "bytes") { + t.Fatalf("oversized image error = %v", err) + } + if err := validateGemma4ImageDimensions(maxGemma4ImageDimension+1, 1); err == nil { + t.Fatal("oversized dimension error = nil") + } + if err := validateGemma4ImageDimensions(8192, 8193); err == nil { + t.Fatal("oversized pixel count error = nil") + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := preprocessGemma4Image(ctx, []byte("image"), cfg, 1); !errors.Is(err, context.Canceled) { + t.Fatalf("cancelled preprocessing error = %v", err) + } +} + func TestGemma4ResizeDimensions(t *testing.T) { gotW, gotH, err := gemma4ResizeDimensions(1024, 768, 16, 280*9, 3) if err != nil { @@ -176,7 +210,7 @@ func TestPreprocessGemma4ImageSoftTokenBudget(t *testing.T) { t.Fatalf("png.Encode() error = %v", err) } - img, err := preprocessGemma4Image(buf.Bytes(), testVisionConfig(t, 1), 1) + img, err := preprocessGemma4Image(context.Background(), buf.Bytes(), testVisionConfig(t, 1), 1) if err != nil { t.Fatalf("preprocessGemma4Image() error = %v", err) } @@ -220,7 +254,7 @@ func TestPrepareMediaPreservesOrderedImageItems(t *testing.T) { {Tokens: []int32{3}}, {Kind: "image", Data: pngData(color.RGBA{B: 255, A: 255})}, } - got, err := m.PrepareMedia(segments) + got, err := m.PrepareMedia(context.Background(), segments) if err != nil { t.Fatalf("PrepareMedia() error = %v", err) } @@ -274,11 +308,11 @@ func TestPrepareMediaSequentialRequestsAreIsolated(t *testing.T) { Vision: &VisionModel{}, EmbedVision: &MultimodalEmbedder{}, } - first, err := m.PrepareMedia([]base.Segment{{Tokens: []int32{1}}, {Kind: "image", Data: pngData(color.RGBA{R: 255, A: 255})}}) + first, err := m.PrepareMedia(context.Background(), []base.Segment{{Tokens: []int32{1}}, {Kind: "image", Data: pngData(color.RGBA{R: 255, A: 255})}}) if err != nil { t.Fatal(err) } - second, err := m.PrepareMedia([]base.Segment{{Tokens: []int32{2, 3}}, {Kind: "image", Data: pngData(color.RGBA{B: 255, A: 255})}}) + second, err := m.PrepareMedia(context.Background(), []base.Segment{{Tokens: []int32{2, 3}}, {Kind: "image", Data: pngData(color.RGBA{B: 255, A: 255})}}) if err != nil { t.Fatal(err) } @@ -296,3 +330,149 @@ func TestPrepareMediaSequentialRequestsAreIsolated(t *testing.T) { t.Fatal("mutating the later prepared request changed the earlier request") } } + +type countingContext struct{ calls int } + +func (c *countingContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (c *countingContext) Done() <-chan struct{} { return nil } +func (c *countingContext) Err() error { + c.calls++ + return nil +} +func (c *countingContext) Value(any) any { return nil } + +type nthCancelContext struct { + target int + calls int +} + +func (c *nthCancelContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (c *nthCancelContext) Done() <-chan struct{} { return nil } +func (c *nthCancelContext) Value(any) any { return nil } +func (c *nthCancelContext) Err() error { + c.calls++ + if c.calls >= c.target { + return context.Canceled + } + return nil +} + +func TestImageToCHWFloat32CancellationAllocationBoundary(t *testing.T) { + img := image.NewRGBA(image.Rect(0, 0, 64, 64)) + measure := func(target int) float64 { + ctx := &nthCancelContext{target: target} + return testing.AllocsPerRun(100, func() { + ctx.calls = 0 + pixels, err := imageToCHWFloat32Context(ctx, img) + if !errors.Is(err, context.Canceled) || pixels != nil { + t.Fatalf("imageToCHWFloat32Context() = (%v, %v), want (nil, context.Canceled)", pixels, err) + } + }) + } + preAllocation := measure(1) + firstRow := measure(2) + if preAllocation != 0 { + t.Fatalf("pre-allocation cancellation allocated %.1f objects", preAllocation) + } + if firstRow <= preAllocation { + t.Fatalf("first-row cancellation allocations = %.1f, want more than pre-allocation %.1f", firstRow, preAllocation) + } +} + +type blockingCancelContext struct { + target int + calls int + reached chan struct{} + release chan struct{} + done chan struct{} +} + +func newBlockingCancelContext(target int) *blockingCancelContext { + return &blockingCancelContext{target: target, reached: make(chan struct{}), release: make(chan struct{}), done: make(chan struct{})} +} +func (c *blockingCancelContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (c *blockingCancelContext) Done() <-chan struct{} { return c.done } +func (c *blockingCancelContext) Value(any) any { return nil } +func (c *blockingCancelContext) Err() error { + c.calls++ + if c.calls < c.target { + return nil + } + if c.calls == c.target { + close(c.done) + close(c.reached) + <-c.release + } + return context.Canceled +} + +func TestPrepareMediaCancellationDuringProductionStages(t *testing.T) { + img := image.NewRGBA(image.Rect(0, 0, 8, 4)) + var encoded bytes.Buffer + if err := png.Encode(&encoded, img); err != nil { + t.Fatal(err) + } + data := encoded.Bytes() + cfg := testVisionConfig(t, 1) + + configChecks := &countingContext{} + if _, _, err := decodeGemma4ImageConfig(configChecks, data); err != nil { + t.Fatal(err) + } + decodeChecks := &countingContext{} + decoded, _, err := decodeGemma4Image(decodeChecks, data) + if err != nil { + t.Fatal(err) + } + targetW, targetH, err := gemma4ResizeDimensions(8, 4, int(cfg.PatchSize), int(cfg.DefaultOutputLength)*int(cfg.PoolingKernelSize)*int(cfg.PoolingKernelSize), int(cfg.PoolingKernelSize)) + if err != nil { + t.Fatal(err) + } + resizeChecks := &countingContext{} + resized, err := resizeGemma4Image(resizeChecks, decoded, decoded.Bounds(), targetW, targetH) + if err != nil { + t.Fatal(err) + } + + // PrepareMedia checks once per segment and preprocessing checks once before + // entering the three measured production stages. + prefixChecks := 2 + tests := []struct { + name string + target int + }{ + {"decode", prefixChecks + 1}, + {"resize", prefixChecks + configChecks.calls + decodeChecks.calls + 2}, + {"chw", prefixChecks + configChecks.calls + decodeChecks.calls + resizeChecks.calls + 2}, + } + _ = resized // its successful construction proves the calibration follows the live resize path. + + m := &Model{ + TextConfig: &TextConfig{ImageTokenIDValue: 10, BOITokenIDValue: 11, EOITokenIDValue: 12, VisionSoftTokens: 1}, + VisionConfig: cfg, + Vision: &VisionModel{}, + EmbedVision: &MultimodalEmbedder{}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := newBlockingCancelContext(tt.target) + result := make(chan struct { + prepared *base.PreparedRequest + err error + }, 1) + go func() { + prepared, err := m.PrepareMedia(ctx, []base.Segment{{Kind: "image", Data: data}}) + result <- struct { + prepared *base.PreparedRequest + err error + }{prepared, err} + }() + <-ctx.reached + close(ctx.release) + got := <-result + if !errors.Is(got.err, context.Canceled) || got.prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", got.prepared, got.err) + } + }) + } +} diff --git a/x/models/glimmer/media.go b/x/models/glimmer/media.go index 97ea6f2d4b7..f146b568688 100644 --- a/x/models/glimmer/media.go +++ b/x/models/glimmer/media.go @@ -2,11 +2,14 @@ package glimmer import ( "bytes" + "context" "fmt" "image" + "image/color" _ "image/gif" _ "image/jpeg" _ "image/png" + "io" "math" "golang.org/x/image/draw" @@ -33,9 +36,15 @@ type preparedImage struct { // PrepareMedia implements base.MediaModel: splice each image segment's // placeholder expansion — image_start, the patch-token run, image_end — into // the stream, decoding, resizing, and patchifying the image on the CPU. -func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { +func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { + if err := ctx.Err(); err != nil { + return nil, err + } prepared := &base.PreparedRequest{} for s, seg := range segments { + if err := ctx.Err(); err != nil { + return nil, err + } if seg.Data == nil { prepared.Tokens = append(prepared.Tokens, seg.Tokens...) continue @@ -47,14 +56,19 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er return nil, fmt.Errorf("glimmer does not support %s input", seg.Kind) } - patches, geom, err := m.preprocessImage(seg.Data) + patches, geom, err := m.preprocessImage(ctx, seg.Data) if err != nil { return nil, fmt.Errorf("preprocess image: %w", err) } start := len(prepared.Tokens) prepared.Tokens = append(prepared.Tokens, m.ImageStartTokenID) - for range geom.outputTokens { + for i := range geom.outputTokens { + if i%256 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } prepared.Tokens = append(prepared.Tokens, m.PatchTokenID) } prepared.Tokens = append(prepared.Tokens, m.ImageEndTokenID) @@ -71,6 +85,9 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er Causal: true, }) } + if err := ctx.Err(); err != nil { + return nil, err + } return prepared, nil } @@ -155,9 +172,15 @@ func computeImageSize(width, height, patchStride, maxTokens int) (targetWidth, t // preprocessImage decodes and prepares one image on the CPU: budgeted // resize, [-1,1] rescale, and patchify to the tower's layout — one row per // patch, (temporal, RGB, pixel row, pixel column) within it. -func (m *Model) preprocessImage(data []byte) ([]float32, preparedImage, error) { - src, _, err := image.Decode(bytes.NewReader(data)) +func (m *Model) preprocessImage(ctx context.Context, data []byte) ([]float32, preparedImage, error) { + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } + src, _, err := image.Decode(glimmerContextReader{check: ctx.Err, r: bytes.NewReader(data)}) if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, preparedImage{}, ctxErr + } return nil, preparedImage{}, fmt.Errorf("decode: %w", err) } @@ -168,17 +191,25 @@ func (m *Model) preprocessImage(data []byte) ([]float32, preparedImage, error) { return nil, preparedImage{}, fmt.Errorf("invalid image dimensions %dx%d", bounds.Dx(), bounds.Dy()) } - resized := image.NewNRGBA(image.Rect(0, 0, targetW, targetH)) - draw.CatmullRom.Scale(resized, resized.Bounds(), src, bounds, draw.Src, nil) + resized, err := resizeGlimmerImage(ctx, src, bounds, targetW, targetH) + if err != nil { + return nil, preparedImage{}, err + } patchSize := int(m.VisionPatchSize) temporal := int(m.VisionPatchTemporal) gridH, gridW := targetH/patchSize, targetW/patchSize patchDim := temporal * 3 * patchSize * patchSize + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } patches := make([]float32, gridH*gridW*patchDim) at := 0 for patchY := range gridH { + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } for patchX := range gridW { for range temporal { for channel := range 3 { @@ -195,5 +226,57 @@ func (m *Model) preprocessImage(data []byte) ([]float32, preparedImage, error) { } } + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } return patches, preparedImage{gridH: gridH, gridW: gridW, outputTokens: outputTokens}, nil } + +type glimmerContextReader struct { + check func() error + r io.Reader +} + +func (r glimmerContextReader) Read(p []byte) (int, error) { + if err := r.check(); err != nil { + return 0, err + } + return r.r.Read(p) +} + +type glimmerResizeCancellation struct{ marker byte } + +var glimmerResizeCanceled = &glimmerResizeCancellation{marker: 1} + +type glimmerCancellableImage struct { + *image.NRGBA + check func() error +} + +func (dst glimmerCancellableImage) Set(x, y int, c color.Color) { + if dst.check() != nil { + panic(glimmerResizeCanceled) + } + dst.NRGBA.Set(x, y, c) +} + +func resizeGlimmerImage(ctx context.Context, src image.Image, bounds image.Rectangle, width, height int) (resized *image.NRGBA, err error) { + if err := ctx.Err(); err != nil { + return nil, err + } + dst := image.NewNRGBA(image.Rect(0, 0, width, height)) + defer func() { + if recovered := recover(); recovered != nil { + if recovered == glimmerResizeCanceled { + resized, err = nil, ctx.Err() + return + } + panic(recovered) + } + }() + draw.CatmullRom.Scale(glimmerCancellableImage{NRGBA: dst, check: ctx.Err}, dst.Bounds(), src, bounds, draw.Src, nil) + if err := ctx.Err(); err != nil { + return nil, err + } + return dst, nil +} diff --git a/x/models/glimmer/media_test.go b/x/models/glimmer/media_test.go index 380fcf0de52..f22854e13ce 100644 --- a/x/models/glimmer/media_test.go +++ b/x/models/glimmer/media_test.go @@ -2,13 +2,75 @@ package glimmer import ( "bytes" + "context" + "errors" "image" + "image/color" "image/png" + "sync" + "sync/atomic" "testing" + "time" "github.com/ollama/ollama/x/mlxrunner/model/base" ) +type glimmerPanickingImage struct{ payload any } + +func (glimmerPanickingImage) ColorModel() color.Model { return color.NRGBAModel } +func (glimmerPanickingImage) Bounds() image.Rectangle { return image.Rect(0, 0, 4, 4) } +func (p glimmerPanickingImage) At(int, int) color.Color { panic(p.payload) } + +type glimmerTestContext struct{} + +func (glimmerTestContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (glimmerTestContext) Done() <-chan struct{} { return nil } +func (glimmerTestContext) Value(any) any { return nil } + +type glimmerCountingContext struct { + glimmerTestContext + calls atomic.Int32 +} + +func (c *glimmerCountingContext) Err() error { + c.calls.Add(1) + return nil +} + +type glimmerCheckpointContext struct { + glimmerTestContext + blockAt int32 + calls atomic.Int32 + canceled atomic.Bool + reached chan struct{} + release chan struct{} + once sync.Once +} + +func newGlimmerCheckpointContext(blockAt int32) *glimmerCheckpointContext { + return &glimmerCheckpointContext{ + blockAt: blockAt, + reached: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (c *glimmerCheckpointContext) Err() error { + if c.calls.Add(1) == c.blockAt { + c.once.Do(func() { close(c.reached) }) + <-c.release + } + if c.canceled.Load() { + return context.Canceled + } + return nil +} + +func (c *glimmerCheckpointContext) cancel() { + c.canceled.Store(true) + close(c.release) +} + func testVisionModel() *Model { return &Model{ VisionEncoder: &VisionEncoder{}, @@ -35,7 +97,7 @@ func testPNG(t *testing.T, w, h int) []byte { func TestPrepareMediaSplicesExpansion(t *testing.T) { m := testVisionModel() - prepared, err := m.PrepareMedia([]base.Segment{ + prepared, err := m.PrepareMedia(context.Background(), []base.Segment{ {Tokens: []int32{1, 2}}, {Kind: "image", Data: testPNG(t, 56, 56)}, {Tokens: []int32{3}}, @@ -85,14 +147,89 @@ func TestPrepareMediaSplicesExpansion(t *testing.T) { func TestPrepareMediaRejectsUnsupportedKind(t *testing.T) { m := testVisionModel() - _, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: []byte{1}}}) + _, err := m.PrepareMedia(context.Background(), []base.Segment{{Kind: "audio", Data: []byte{1}}}) if err == nil { t.Fatal("expected error for audio input") } text := &Model{Config: &Config{}} - _, err = text.PrepareMedia([]base.Segment{{Kind: "image", Data: []byte{1}}}) + _, err = text.PrepareMedia(context.Background(), []base.Segment{{Kind: "image", Data: []byte{1}}}) if err == nil { t.Fatal("expected error for text-only model") } } + +func TestPrepareMediaPreCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + prepared, err := testVisionModel().PrepareMedia(ctx, []base.Segment{{Tokens: []int32{1}}}) + if !errors.Is(err, context.Canceled) || prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", prepared, err) + } +} + +func TestImageReaderCancellation(t *testing.T) { + ctx := newGlimmerCheckpointContext(1) + result := make(chan error, 1) + go func() { + _, err := glimmerContextReader{check: ctx.Err, r: bytes.NewReader([]byte("image"))}.Read(make([]byte, 1)) + result <- err + }() + <-ctx.reached + ctx.cancel() + if err := <-result; !errors.Is(err, context.Canceled) { + t.Fatalf("Read() error = %v, want context.Canceled", err) + } +} + +func TestResizeCancellation(t *testing.T) { + ctx := newGlimmerCheckpointContext(2) // entry check, then the first destination pixel + result := make(chan error, 1) + go func() { + _, err := resizeGlimmerImage(ctx, image.NewNRGBA(image.Rect(0, 0, 4, 4)), image.Rect(0, 0, 4, 4), 64, 64) + result <- err + }() + <-ctx.reached + ctx.cancel() + if err := <-result; !errors.Is(err, context.Canceled) { + t.Fatalf("resizeGlimmerImage() error = %v, want context.Canceled", err) + } +} + +func TestResizeRepanicsUnrelatedPayload(t *testing.T) { + payload := &struct{ marker byte }{marker: 1} + recovered := func() (got any) { + defer func() { got = recover() }() + _, _ = resizeGlimmerImage(context.Background(), glimmerPanickingImage{payload: payload}, image.Rect(0, 0, 4, 4), 64, 64) + return nil + }() + if recovered != payload { + t.Fatalf("recovered payload = %v, want exact unrelated payload %v", recovered, payload) + } +} + +func TestPrepareMediaFinalCancellation(t *testing.T) { + m := testVisionModel() + segments := []base.Segment{{Kind: "image", Data: testPNG(t, 56, 56)}} + counting := &glimmerCountingContext{} + if _, err := m.PrepareMedia(counting, segments); err != nil { + t.Fatal(err) + } + + ctx := newGlimmerCheckpointContext(counting.calls.Load()) + type result struct { + prepared *base.PreparedRequest + err error + } + done := make(chan result, 1) + go func() { + prepared, err := m.PrepareMedia(ctx, segments) + done <- result{prepared: prepared, err: err} + }() + <-ctx.reached + ctx.cancel() + got := <-done + if !errors.Is(got.err, context.Canceled) || got.prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", got.prepared, got.err) + } +} diff --git a/x/models/qwen3_5/process_image.go b/x/models/qwen3_5/process_image.go index 2796441f8cc..2cf59ed0f73 100644 --- a/x/models/qwen3_5/process_image.go +++ b/x/models/qwen3_5/process_image.go @@ -2,8 +2,11 @@ package qwen3_5 import ( "bytes" + "context" "fmt" "image" + "image/color" + "io" "math" "golang.org/x/image/draw" @@ -66,42 +69,68 @@ func smartResize(height, width, factor int32) (int32, int32, error) { // preprocessImage decodes and prepares one image: aspect-preserving resize, // (x/255-0.5)/0.5 normalization, and patchification into the tower's // block-major row layout. -func (m *Model) preprocessImage(data []byte) (pixels []float32, prep preparedImage, err error) { - img, _, err := image.Decode(bytes.NewReader(data)) +func (m *Model) preprocessImage(ctx context.Context, data []byte) (pixels []float32, prep preparedImage, err error) { + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } + img, _, err := image.Decode(qwenImageContextReader{check: ctx.Err, r: bytes.NewReader(data)}) if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, preparedImage{}, ctxErr + } return nil, preparedImage{}, fmt.Errorf("decode image: %w", err) } v := m.Vision patch, merge, temporal := v.PatchSize, v.SpatialMergeSize, v.TemporalPatchSize bounds := img.Bounds() - img = dropAlpha(img, bounds) + img, err = dropAlpha(ctx, img, bounds) + if err != nil { + return nil, preparedImage{}, err + } targetH, targetW, err := smartResize(int32(bounds.Dy()), int32(bounds.Dx()), patch*merge) if err != nil { return nil, preparedImage{}, err } - resized := image.NewRGBA(image.Rect(0, 0, int(targetW), int(targetH))) - draw.CatmullRom.Scale(resized, resized.Bounds(), img, bounds, draw.Src, nil) + resized, err := resizeQwenImage(ctx, img, bounds, int(targetW), int(targetH)) + if err != nil { + return nil, preparedImage{}, err + } gridH, gridW := targetH/patch, targetW/patch prep = preparedImage{gridH: gridH, gridW: gridW} - pixels = patchifyImage(resized, gridH, gridW, patch, merge, temporal) - prep.ropePos, prep.posIdx, prep.posWt = visionPatchLookups(gridH, gridW, merge, v.NumPositionEmbeddings) + pixels, err = patchifyImage(ctx, resized, gridH, gridW, patch, merge, temporal) + if err != nil { + return nil, preparedImage{}, err + } + prep.ropePos, prep.posIdx, prep.posWt, err = visionPatchLookups(ctx, gridH, gridW, merge, v.NumPositionEmbeddings) + if err != nil { + return nil, preparedImage{}, err + } + if err := ctx.Err(); err != nil { + return nil, preparedImage{}, err + } return pixels, prep, nil } // patchifyImage converts the resized image to the tower's row layout: // block-major over merge blocks, each row [channel][temporal duplicate] // [pixel row][pixel column]. -func patchifyImage(resized *image.RGBA, gridH, gridW, patch, merge, temporal int32) []float32 { +func patchifyImage(ctx context.Context, resized *image.RGBA, gridH, gridW, patch, merge, temporal int32) ([]float32, error) { n := int(gridH * gridW) rowLen := int(3 * temporal * patch * patch) + if err := ctx.Err(); err != nil { + return nil, err + } pixels := make([]float32, n*rowLen) p, t := int(patch), int(temporal) patchArea := p * p row := 0 for bh := range gridH / merge { + if err := ctx.Err(); err != nil { + return nil, err + } for bw := range gridW / merge { for mh := range merge { for mw := range merge { @@ -126,15 +155,21 @@ func patchifyImage(resized *image.RGBA, gridH, gridW, patch, merge, temporal int } } } - return pixels + if err := ctx.Err(); err != nil { + return nil, err + } + return pixels, nil } // visionPatchLookups computes, in the block-major token order, each patch's // (h, w) rotary position and its four bilinear corners into the learned // side*side position grid with their weights. -func visionPatchLookups(gridH, gridW, merge, numPosEmbeddings int32) (ropePos, posIdx []int32, posWt []float32) { +func visionPatchLookups(ctx context.Context, gridH, gridW, merge, numPosEmbeddings int32) (ropePos, posIdx []int32, posWt []float32, err error) { side := int32(math.Sqrt(float64(numPosEmbeddings))) n := int(gridH * gridW) + if err := ctx.Err(); err != nil { + return nil, nil, nil, err + } ropePos = make([]int32, 0, 2*n) posIdx = make([]int32, 0, 4*n) posWt = make([]float32, 0, 4*n) @@ -148,6 +183,9 @@ func visionPatchLookups(gridH, gridW, merge, numPosEmbeddings int32) (ropePos, p } for bh := range gridH / merge { + if err := ctx.Err(); err != nil { + return nil, nil, nil, err + } for bw := range gridW / merge { for mh := range merge { for mw := range merge { @@ -165,19 +203,80 @@ func visionPatchLookups(gridH, gridW, merge, numPosEmbeddings int32) (ropePos, p } } } - return ropePos, posIdx, posWt + if err := ctx.Err(); err != nil { + return nil, nil, nil, err + } + return ropePos, posIdx, posWt, nil } // dropAlpha flattens a non-opaque image to straight RGB: the reference // drops alpha via RGB conversion before resizing, not by compositing. -func dropAlpha(img image.Image, bounds image.Rectangle) image.Image { +func dropAlpha(ctx context.Context, img image.Image, bounds image.Rectangle) (image.Image, error) { if o, ok := img.(interface{ Opaque() bool }); ok && o.Opaque() { - return img + return img, ctx.Err() } flat := image.NewNRGBA(bounds) - draw.Draw(flat, bounds, img, bounds.Min, draw.Src) - for i := 3; i < len(flat.Pix); i += 4 { - flat.Pix[i] = 0xff + for y := bounds.Min.Y; y < bounds.Max.Y; y++ { + if err := ctx.Err(); err != nil { + return nil, err + } + for x := bounds.Min.X; x < bounds.Max.X; x++ { + c := color.NRGBAModel.Convert(img.At(x, y)).(color.NRGBA) + c.A = 0xff + flat.SetNRGBA(x, y, c) + } + } + if err := ctx.Err(); err != nil { + return nil, err + } + return flat, nil +} + +type qwenImageContextReader struct { + check func() error + r io.Reader +} + +func (r qwenImageContextReader) Read(p []byte) (int, error) { + if err := r.check(); err != nil { + return 0, err + } + return r.r.Read(p) +} + +type qwenResizeCancellation struct{ marker byte } + +var qwenResizeCanceled = &qwenResizeCancellation{marker: 1} + +type qwenCancellableImage struct { + *image.RGBA + check func() error +} + +func (dst qwenCancellableImage) Set(x, y int, c color.Color) { + if dst.check() != nil { + panic(qwenResizeCanceled) + } + dst.RGBA.Set(x, y, c) +} + +func resizeQwenImage(ctx context.Context, src image.Image, bounds image.Rectangle, width, height int) (resized *image.RGBA, err error) { + if err := ctx.Err(); err != nil { + return nil, err + } + dst := image.NewRGBA(image.Rect(0, 0, width, height)) + defer func() { + if recovered := recover(); recovered != nil { + if recovered == qwenResizeCanceled { + resized, err = nil, ctx.Err() + return + } + panic(recovered) + } + }() + draw.CatmullRom.Scale(qwenCancellableImage{RGBA: dst, check: ctx.Err}, dst.Bounds(), src, bounds, draw.Src, nil) + if err := ctx.Err(); err != nil { + return nil, err } - return flat + return dst, nil } diff --git a/x/models/qwen3_5/vision.go b/x/models/qwen3_5/vision.go index fe1b45a1e61..8a3b9430d92 100644 --- a/x/models/qwen3_5/vision.go +++ b/x/models/qwen3_5/vision.go @@ -1,6 +1,7 @@ package qwen3_5 import ( + "context" "encoding/json" "fmt" "math" @@ -140,8 +141,8 @@ func NewVisionAdapter(configData []byte, tensors map[string]*mlx.Array, cfg Visi return &VisionAdapter{Model: m}, nil } -func (a *VisionAdapter) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { - return a.Model.PrepareMedia(segments) +func (a *VisionAdapter) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { + return a.Model.PrepareMedia(ctx, segments) } func (a *VisionAdapter) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { @@ -268,7 +269,10 @@ func (m *Model) visionLoaded() bool { return m.VisionTower != nil } // PrepareMedia implements base.MediaModel: preprocess each image segment, // splice its vision_start + pads + vision_end expansion, and precompute the // request's 3-channel rope positions. -func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { +func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { + if err := ctx.Err(); err != nil { + return nil, err + } prepared := &base.PreparedRequest{} // pos walks the reference's multimodal position rule: text advances one @@ -277,18 +281,29 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er var positions []int32 cur := int32(0) maxPos := int32(-1) - text := func(n int) { - for range n { + text := func(n int) error { + for i := range n { + if i%256 == 0 { + if err := ctx.Err(); err != nil { + return err + } + } positions = append(positions, cur, cur, cur) maxPos = max(maxPos, cur) cur++ } + return nil } for s, seg := range segments { + if err := ctx.Err(); err != nil { + return nil, err + } if seg.Data == nil { prepared.Tokens = append(prepared.Tokens, seg.Tokens...) - text(len(seg.Tokens)) + if err := text(len(seg.Tokens)); err != nil { + return nil, err + } continue } if !m.visionLoaded() { @@ -298,17 +313,22 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er return nil, fmt.Errorf("qwen3.5 does not support %s input", seg.Kind) } - pixels, prep, err := m.preprocessImage(seg.Data) + pixels, prep, err := m.preprocessImage(ctx, seg.Data) if err != nil { return nil, err } start := len(prepared.Tokens) prepared.Tokens = append(prepared.Tokens, m.MM.VisionStartTokenID) - text(1) + if err := text(1); err != nil { + return nil, err + } merge := m.Vision.SpatialMergeSize mh, mw := prep.gridH/merge, prep.gridW/merge for bh := range mh { + if err := ctx.Err(); err != nil { + return nil, err + } for bw := range mw { prepared.Tokens = append(prepared.Tokens, m.MM.ImageTokenID) positions = append(positions, cur, cur+bh, cur+bw) @@ -317,7 +337,9 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er } cur += max(mh, mw) prepared.Tokens = append(prepared.Tokens, m.MM.VisionEndTokenID) - text(1) + if err := text(1); err != nil { + return nil, err + } n := int(prep.numPatches()) prepared.Items = append(prepared.Items, base.PreparedItem{ @@ -336,12 +358,20 @@ func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, er L := int32(len(prepared.Tokens)) table := make([]int32, 3*L) for i := range L { + if i%256 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } table[i] = positions[3*i] table[L+i] = positions[3*i+1] table[2*L+i] = positions[3*i+2] } prepared.Layout = &visionLayout{positions: table, promptLen: L, delta: maxPos + 1 - L} } + if err := ctx.Err(); err != nil { + return nil, err + } return prepared, nil } diff --git a/x/models/qwen3_5/vision_test.go b/x/models/qwen3_5/vision_test.go index e96307f6c84..d637131454a 100644 --- a/x/models/qwen3_5/vision_test.go +++ b/x/models/qwen3_5/vision_test.go @@ -2,16 +2,93 @@ package qwen3_5 import ( "bytes" + "context" + "errors" "image" + "image/color" "image/png" "strings" + "sync" + "sync/atomic" "testing" + "time" "github.com/ollama/ollama/x/internal/mlxtest" "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/mlxrunner/model/base" ) +type qwenPanickingImage struct{ payload any } + +func (qwenPanickingImage) ColorModel() color.Model { return color.NRGBAModel } +func (qwenPanickingImage) Bounds() image.Rectangle { return image.Rect(0, 0, 4, 4) } +func (p qwenPanickingImage) At(int, int) color.Color { panic(p.payload) } + +type qwenTestContext struct{} + +func (qwenTestContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (qwenTestContext) Done() <-chan struct{} { return nil } +func (qwenTestContext) Value(any) any { return nil } + +type qwenCountingContext struct { + qwenTestContext + calls atomic.Int32 +} + +func (c *qwenCountingContext) Err() error { + c.calls.Add(1) + return nil +} + +type qwenCheckpointContext struct { + qwenTestContext + blockAt int32 + calls atomic.Int32 + canceled atomic.Bool + reached chan struct{} + release chan struct{} + once sync.Once +} + +func newQwenCheckpointContext(blockAt int32) *qwenCheckpointContext { + return &qwenCheckpointContext{ + blockAt: blockAt, + reached: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (c *qwenCheckpointContext) Err() error { + if c.calls.Add(1) == c.blockAt { + c.once.Do(func() { close(c.reached) }) + <-c.release + } + if c.canceled.Load() { + return context.Canceled + } + return nil +} + +func (c *qwenCheckpointContext) cancel() { + c.canceled.Store(true) + close(c.release) +} + +func testQwenVisionModel() *Model { + cfg := &VisionConfig{PatchSize: 16, SpatialMergeSize: 2, TemporalPatchSize: 2, InChannels: 3, NumPositionEmbeddings: 2304} + return &Model{ + Config: &Config{}, + VisionTower: &VisionTower{}, + Vision: cfg, + MM: multimodalConfig{ + ImageTokenID: 9, + VisionStartTokenID: 7, + VisionEndTokenID: 8, + VisionConfig: cfg, + }, + } +} + func TestVisionAdapterWeightsAreCollectable(t *testing.T) { mlxtest.Setup(t) weight := mlx.FromValue(float32(1)) @@ -71,20 +148,9 @@ func TestParseMultimodalConfig(t *testing.T) { // image grid anchors T at the current position, offsets H/W by row/column, // and advances the position by the grid's longer side. func TestMRopePositionRule(t *testing.T) { - cfg := &VisionConfig{PatchSize: 16, SpatialMergeSize: 2, TemporalPatchSize: 2, InChannels: 3, NumPositionEmbeddings: 2304} - m := &Model{ - Config: &Config{}, - VisionTower: &VisionTower{}, - Vision: cfg, - MM: multimodalConfig{ - ImageTokenID: 9, - VisionStartTokenID: 7, - VisionEndTokenID: 8, - VisionConfig: cfg, - }, - } + m := testQwenVisionModel() - prepared, err := m.PrepareMedia([]base.Segment{{Tokens: []int32{1, 2, 3}}}) + prepared, err := m.PrepareMedia(context.Background(), []base.Segment{{Tokens: []int32{1, 2, 3}}}) if err != nil { t.Fatal(err) } @@ -98,7 +164,7 @@ func TestMRopePositionRule(t *testing.T) { if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 64, 64))); err != nil { t.Fatal(err) } - prepared, err = m.PrepareMedia([]base.Segment{ + prepared, err = m.PrepareMedia(context.Background(), []base.Segment{ {Tokens: []int32{1, 2, 3}}, {Kind: "image", Data: buf.Bytes()}, }) @@ -147,10 +213,29 @@ func TestMRopePositionRule(t *testing.T) { } } +func TestPrepareMediaPreCanceled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + for name, mediaModel := range map[string]base.MediaModel{ + "model": &Model{}, + "adapter": &VisionAdapter{Model: &Model{}}, + } { + t.Run(name, func(t *testing.T) { + prepared, err := mediaModel.PrepareMedia(ctx, []base.Segment{{Tokens: []int32{1}}}) + if !errors.Is(err, context.Canceled) || prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", prepared, err) + } + }) + } +} + func TestVisionPatchLookups(t *testing.T) { // 4x6 pre-merge grid, merge 2: token order is block-major; rope // positions are the pre-merge (h, w) indices. - ropePos, posIdx, posWt := visionPatchLookups(4, 6, 2, 2304) + ropePos, posIdx, posWt, err := visionPatchLookups(context.Background(), 4, 6, 2, 2304) + if err != nil { + t.Fatal(err) + } if len(ropePos) != 2*24 || len(posIdx) != 4*24 || len(posWt) != 4*24 { t.Fatalf("lengths %d %d %d", len(ropePos), len(posIdx), len(posWt)) } @@ -176,3 +261,134 @@ func TestVisionPatchLookups(t *testing.T) { } } } + +func TestImageReaderCancellation(t *testing.T) { + ctx := newQwenCheckpointContext(1) + result := make(chan error, 1) + go func() { + _, err := qwenImageContextReader{check: ctx.Err, r: bytes.NewReader([]byte("image"))}.Read(make([]byte, 1)) + result <- err + }() + <-ctx.reached + ctx.cancel() + if err := <-result; !errors.Is(err, context.Canceled) { + t.Fatalf("Read() error = %v, want context.Canceled", err) + } +} + +func TestDropAlphaCancellation(t *testing.T) { + src := image.NewNRGBA(image.Rect(0, 0, 64, 64)) + src.SetNRGBA(0, 0, color.NRGBA{R: 1, A: 1}) + ctx := newQwenCheckpointContext(1) + type result struct { + image image.Image + err error + } + done := make(chan result, 1) + go func() { + img, err := dropAlpha(ctx, src, src.Bounds()) + done <- result{image: img, err: err} + }() + <-ctx.reached + ctx.cancel() + got := <-done + if !errors.Is(got.err, context.Canceled) || got.image != nil { + t.Fatalf("dropAlpha() = (%v, %v), want (nil, context.Canceled)", got.image, got.err) + } +} + +func TestResizeCancellation(t *testing.T) { + ctx := newQwenCheckpointContext(2) // entry check, then the first destination pixel + result := make(chan error, 1) + go func() { + _, err := resizeQwenImage(ctx, image.NewRGBA(image.Rect(0, 0, 4, 4)), image.Rect(0, 0, 4, 4), 64, 64) + result <- err + }() + <-ctx.reached + ctx.cancel() + if err := <-result; !errors.Is(err, context.Canceled) { + t.Fatalf("resizeQwenImage() error = %v, want context.Canceled", err) + } +} + +func TestResizeRepanicsUnrelatedPayload(t *testing.T) { + payload := &struct{ marker byte }{marker: 1} + recovered := func() (got any) { + defer func() { got = recover() }() + _, _ = resizeQwenImage(context.Background(), qwenPanickingImage{payload: payload}, image.Rect(0, 0, 4, 4), 64, 64) + return nil + }() + if recovered != payload { + t.Fatalf("recovered payload = %v, want exact unrelated payload %v", recovered, payload) + } +} + +func TestPatchifyCancellation(t *testing.T) { + ctx := newQwenCheckpointContext(2) // allocation check, then the first merge-block row + type result struct { + pixels []float32 + err error + } + done := make(chan result, 1) + go func() { + pixels, err := patchifyImage(ctx, image.NewRGBA(image.Rect(0, 0, 64, 64)), 4, 4, 16, 2, 2) + done <- result{pixels: pixels, err: err} + }() + <-ctx.reached + ctx.cancel() + got := <-done + if !errors.Is(got.err, context.Canceled) || got.pixels != nil { + t.Fatalf("patchifyImage() = (%v, %v), want (nil, context.Canceled)", got.pixels, got.err) + } +} + +func TestVisionPatchLookupsCancellation(t *testing.T) { + ctx := newQwenCheckpointContext(2) // allocation check, then the first merge-block row + type result struct { + ropePos []int32 + posIdx []int32 + posWt []float32 + err error + } + done := make(chan result, 1) + go func() { + ropePos, posIdx, posWt, err := visionPatchLookups(ctx, 16, 16, 2, 2304) + done <- result{ropePos: ropePos, posIdx: posIdx, posWt: posWt, err: err} + }() + <-ctx.reached + ctx.cancel() + got := <-done + if !errors.Is(got.err, context.Canceled) || got.ropePos != nil || got.posIdx != nil || got.posWt != nil { + t.Fatalf("visionPatchLookups() returned partial results on cancellation: %+v", got) + } +} + +func TestPrepareMediaFinalCancellation(t *testing.T) { + var buf bytes.Buffer + if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 64, 64))); err != nil { + t.Fatal(err) + } + m := testQwenVisionModel() + segments := []base.Segment{{Kind: "image", Data: buf.Bytes()}} + counting := &qwenCountingContext{} + if _, err := m.PrepareMedia(counting, segments); err != nil { + t.Fatal(err) + } + + ctx := newQwenCheckpointContext(counting.calls.Load()) + type result struct { + prepared *base.PreparedRequest + err error + } + done := make(chan result, 1) + go func() { + prepared, err := m.PrepareMedia(ctx, segments) + done <- result{prepared: prepared, err: err} + }() + <-ctx.reached + ctx.cancel() + got := <-done + if !errors.Is(got.err, context.Canceled) || got.prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", got.prepared, got.err) + } +} diff --git a/x/models/qwen4_exp/media_test.go b/x/models/qwen4_exp/media_test.go new file mode 100644 index 00000000000..d7002ebabf5 --- /dev/null +++ b/x/models/qwen4_exp/media_test.go @@ -0,0 +1,20 @@ +package qwen4_exp + +import ( + "context" + "errors" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/model/base" + "github.com/ollama/ollama/x/models/qwen3_5" +) + +func TestPrepareMediaPassesCancellationToVision(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + m := &Model{Vision: &qwen3_5.VisionAdapter{Model: &qwen3_5.Model{}}} + prepared, err := m.PrepareMedia(ctx, []base.Segment{{Tokens: []int32{1}}}) + if !errors.Is(err, context.Canceled) || prepared != nil { + t.Fatalf("PrepareMedia() = (%v, %v), want (nil, context.Canceled)", prepared, err) + } +} diff --git a/x/models/qwen4_exp/qwen4_exp.go b/x/models/qwen4_exp/qwen4_exp.go index 069ca2eba2a..10294db5af3 100644 --- a/x/models/qwen4_exp/qwen4_exp.go +++ b/x/models/qwen4_exp/qwen4_exp.go @@ -2,6 +2,7 @@ package qwen4_exp import ( + "context" "fmt" "github.com/ollama/ollama/x/mlxrunner/batch" @@ -175,11 +176,11 @@ func (m *Model) Unembed(hidden *mlx.Array) *mlx.Array { func (m *Model) Tokenizer() *tokenizer.Tokenizer { return m.tok } -func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { +func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { if m.Vision == nil { return nil, fmt.Errorf("this model does not support media input") } - return m.Vision.PrepareMedia(segments) + return m.Vision.PrepareMedia(ctx, segments) } func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { From 931eb9ab601ad7e17dc53709942688b3da970927 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 23:21:58 +0000 Subject: [PATCH 45/58] gemma4: support unified MLX vision Load and validate encoder-free unified vision tensors, prepare bounded HWC patch rows and positions, and carry request-local image spans for sliding-layer attention relaxation. Co-authored-by: Codex --- docs/third-party/mlx-vlm.md | 5 +- x/create/client/create.go | 2 + x/create/client/create_test.go | 40 +++++ x/create/gemma4_test.go | 3 + x/models/gemma4/gemma4.go | 50 ++++-- x/models/gemma4/gemma4_test.go | 11 ++ x/models/gemma4/metadata/vision.go | 147 ++++++++++++++++ x/models/gemma4/metadata/vision_test.go | 122 ++++++++++++++ x/models/gemma4/vision.go | 214 ++++++++++++++++++++++-- x/models/gemma4/vision_test.go | 113 +++++++++++++ 10 files changed, 676 insertions(+), 31 deletions(-) diff --git a/docs/third-party/mlx-vlm.md b/docs/third-party/mlx-vlm.md index 1d717eac373..0233657227c 100644 --- a/docs/third-party/mlx-vlm.md +++ b/docs/third-party/mlx-vlm.md @@ -5,8 +5,9 @@ implementation in MLX-VLM: - Repository: https://github.com/Blaizzy/mlx-vlm - Revision: `61990c9054f2bc7bb8f32541e3238b4a58fe64e5` -- Source paths: `mlx_vlm/models/gemma4/gemma4.py` and - `mlx_vlm/models/gemma4/vision.py` +- Source paths: `mlx_vlm/models/gemma4/gemma4.py`, + `mlx_vlm/models/gemma4/vision.py`, and + `mlx_vlm/models/gemma4_unified/gemma4_unified.py` The adapted implementation is distributed under the following license. diff --git a/x/create/client/create.go b/x/create/client/create.go index 2b7e5cb2f33..67d07763e37 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -637,6 +637,8 @@ func gemma4ModelDirHasVisionTensors(modelDir string) bool { } tensors := make(map[string]gemma4metadata.TensorDescriptor, len(inv.Tensors)) for name, tensor := range inv.Tensors { + // Unified vision is encoder-free, so dispatch and readiness depend on + // descriptor shapes rather than the released tower sentinel alone. tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} } return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 49d53991b57..6b9f9fac2ac 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -615,6 +615,46 @@ func TestInferSafetensorsCapabilitiesGemma4VisionRequiresTensors(t *testing.T) { } } +func TestInferSafetensorsCapabilitiesGemma4UnifiedVision(t *testing.T) { + const configJSON = `{ + "architectures":["Gemma4UnifiedForConditionalGeneration"],"model_type":"gemma4_unified", + "text_config":{"hidden_size":5}, + "vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2} + }` + valid := map[string]gemma4metadata.TensorDescriptor{ + "model.vision_embedder.patch_ln1.weight": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_ln1.bias": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_dense.weight": {Dtype: "F32", Shape: []int32{3, 12}}, + "model.vision_embedder.patch_dense.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_embedding": {Dtype: "F32", Shape: []int32{4, 2, 3}}, + "model.vision_embedder.pos_norm.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_norm.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "F32", Shape: []int32{5, 3}}, + } + check := func(t *testing.T, tensors map[string]gemma4metadata.TensorDescriptor, wantVision bool) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, tensors) + got := inferSafetensorsCapabilities(dir, "") + if slices.Contains(got, "vision") != wantVision { + t.Fatalf("capabilities = %v, want vision %t", got, wantVision) + } + } + check(t, valid, true) + partial := maps.Clone(valid) + delete(partial, "model.vision_embedder.pos_norm.bias") + check(t, partial, false) + wrong := maps.Clone(valid) + d := wrong["model.vision_embedder.patch_dense.weight"] + d.Shape = []int32{3, 11} + wrong["model.vision_embedder.patch_dense.weight"] = d + check(t, wrong, false) +} + func TestInferSafetensorsCapabilitiesGemma4PackedSourceRequiresProducerContract(t *testing.T) { const configJSON = `{ "architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4", diff --git a/x/create/gemma4_test.go b/x/create/gemma4_test.go index edaf89fd7ff..aedc7cb8339 100644 --- a/x/create/gemma4_test.go +++ b/x/create/gemma4_test.go @@ -165,6 +165,7 @@ func TestGemma4QuantizationType(t *testing.T) { {"vision q_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight", aligned, "nvfp4", ""}, {"unified vision embedder nvfp4", transform26B, "model.vision_embedder.patch_dense.weight", aligned, "nvfp4", ""}, {"vision projection nvfp4", transform26B, "model.embed_vision.embedding_projection.linear.weight", aligned, "nvfp4", ""}, + {"embed_vision int4", transform26B, "model.embed_vision.embedding_projection.weight", aligned, "int4", ""}, // Audio tower down_proj {"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", ""}, {"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", ""}, @@ -189,6 +190,7 @@ func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { "model.vision_tower.patch_embedder.input_proj.weight": "BF16", "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight": "BF16", "model.embed_vision.embedding_projection.weight": "BF16", + "model.vision_embedder.patch_dense.weight": "BF16", "model.audio_tower.subsample_conv_projection.input_proj_linear.weight": "BF16", "model.embed_audio.embedding_projection.weight": "BF16", }) @@ -210,6 +212,7 @@ func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { "model.embed_vision.embedding_projection.weight", "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", "model.embed_audio.embedding_projection.weight", + "model.vision_embedder.patch_dense.weight", } { tensor, ok := got[name] if !ok { diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index 41e7d9a0665..e20d8f77f72 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -374,12 +374,13 @@ type DecoderLayer struct { // Model is the Gemma 4 model (text + optional vision). type Model struct { - EmbedTokens nn.EmbeddingLayer - Layers []*DecoderLayer - Norm *nn.RMSNorm - LMHead nn.LinearLayer - Vision *VisionModel - EmbedVision *MultimodalEmbedder + EmbedTokens nn.EmbeddingLayer + Layers []*DecoderLayer + Norm *nn.RMSNorm + LMHead nn.LinearLayer + Vision *VisionModel + UnifiedVision *UnifiedVisionEmbedder + EmbedVision *MultimodalEmbedder // PLE model-level components (nil if no PLE). EmbedTokensPerLayer nn.EmbeddingLayer @@ -776,15 +777,23 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { return fmt.Errorf("invalid Gemma4 vision tensors: %w", err) } if visionReady { - vision, err := loadVisionModel(tensors, m.VisionConfig, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) - if err != nil { - return err + if m.VisionConfig.unified() { + vision, err := loadUnifiedVisionEmbedder(tensors, m.VisionConfig, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + m.UnifiedVision = vision + } else { + vision, err := loadVisionModel(tensors, m.VisionConfig, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + m.Vision = vision } embedVision, err := loadMultimodalEmbedder(tensors, "embed_vision", m.VisionConfig.RMSNormEps, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) if err != nil { return err } - m.Vision = vision m.EmbedVision = embedVision } @@ -1375,7 +1384,7 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio // kernel only handles L < 4 (generation). For prefill, we fall back // to explicit matmul+softmax+matmul on CUDA. var k, v *mlx.Array - mask := nn.CausalMask().Intersect(nn.QPaddingMask(b, q.DType())) + mask := gemma4AttentionMask(b, isSliding).Intersect(nn.QPaddingMask(b, q.DType())) if kv.history != nil { k, v = kv.history.K(), kv.history.V() mask = kv.history.Mask(mask) @@ -1405,7 +1414,7 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio out = mlx.Reshape(out, B, cfg.NumAttentionHeads, L, headDim) } else { var opt nn.SDPAOption - mask := nn.CausalMask() + mask := gemma4AttentionMask(b, isSliding) if kv.history != nil { opt = nn.WithKVHistory(kv.history) } else { @@ -1423,6 +1432,23 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio return a.OProj.Forward(out), kv } +func gemma4AttentionMask(b *batch.Batch, isSliding bool) nn.AttentionMask { + mask := nn.CausalMask() + if !isSliding { + return mask + } + for seq, raw := range b.Layout { + layout, ok := raw.(*gemma4MediaLayout) + if !ok || layout == nil { + continue + } + for _, span := range layout.ImageSpans { + mask = mask.Relax(seq, span[0], span[1], span[0], span[1]) + } + } + return mask +} + func (m *MLP) Forward(x *mlx.Array) *mlx.Array { gate := m.GateProj.Forward(x) up := m.UpProj.Forward(x) diff --git a/x/models/gemma4/gemma4_test.go b/x/models/gemma4/gemma4_test.go index 27f190434e9..6a53a7578cc 100644 --- a/x/models/gemma4/gemma4_test.go +++ b/x/models/gemma4/gemma4_test.go @@ -5,9 +5,20 @@ import ( "slices" "testing" + "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" ) +func TestGemma4UnifiedVisionMaskOnlyRelaxesSlidingAttention(t *testing.T) { + b := &batch.Batch{Layout: []any{&gemma4MediaLayout{ImageSpans: [][2]int{{4, 12}}}}} + if gemma4AttentionMask(b, true).IsCausal() { + t.Fatal("sliding vision mask remained pure causal") + } + if !gemma4AttentionMask(b, false).IsCausal() { + t.Fatal("full-attention vision mask must remain causal") + } +} + func TestParseSuppressTokens(t *testing.T) { got := parseSuppressTokens([]byte(`{"suppress_tokens":[258883,258882]}`)) want := []int32{258883, 258882} diff --git a/x/models/gemma4/metadata/vision.go b/x/models/gemma4/metadata/vision.go index d5bfd86422d..75028b9807a 100644 --- a/x/models/gemma4/metadata/vision.go +++ b/x/models/gemma4/metadata/vision.go @@ -24,6 +24,8 @@ const ( ) type ConfigFile struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` TextConfig TextConfig `json:"text_config"` VisionConfig *VisionConfig `json:"vision_config"` Quantization Quantization `json:"quantization"` @@ -42,6 +44,7 @@ type Quantization struct { } type VisionConfig struct { + ModelType string `json:"model_type"` HiddenSize int `json:"hidden_size"` IntermediateSize int `json:"intermediate_size"` NumHiddenLayers int `json:"num_hidden_layers"` @@ -55,12 +58,18 @@ type VisionConfig struct { PoolingKernelSize int `json:"pooling_kernel_size"` UseClippedLinears bool `json:"use_clipped_linears"` Standardize bool `json:"standardize"` + MMEmbedDim int `json:"mm_embed_dim"` + MMPosembSize int `json:"mm_posemb_size"` + ModelPatchSize int `json:"model_patch_size"` + NumSoftTokens int `json:"num_soft_tokens"` + OutputProjDims int `json:"output_proj_dims"` RopeParameters struct { RopeTheta float64 `json:"rope_theta"` } `json:"rope_parameters"` } type Architecture struct { + Unified bool TextHiddenSize int HiddenSize int IntermediateSize int @@ -77,6 +86,9 @@ type Architecture struct { ClippingBoundsOptional bool Standardize bool RopeTheta float64 + MMEmbedDim int + MMPosembSize int + ModelPatchSize int } type TensorDescriptor struct { @@ -95,6 +107,36 @@ func projectVisionArchitecture(cfg ConfigFile, requireTextWidth bool) (Architect return Architecture{}, fmt.Errorf("missing vision_config") } v := *cfg.VisionConfig + if isUnifiedConfig(cfg) { + if v.NumSoftTokens == 0 { + v.NumSoftTokens = 280 + } + if v.PatchSize == 0 { + v.PatchSize = 16 + } + if v.PoolingKernelSize == 0 { + v.PoolingKernelSize = 3 + } + if v.ModelPatchSize == 0 { + v.ModelPatchSize = v.PatchSize * v.PoolingKernelSize + } + if v.MMEmbedDim == 0 { + v.MMEmbedDim = v.OutputProjDims + } + if v.MMPosembSize == 0 { + v.MMPosembSize = 1120 + } + a := Architecture{ + Unified: true, TextHiddenSize: cfg.TextConfig.HiddenSize, + DefaultOutputLength: v.NumSoftTokens, PatchSize: v.PatchSize, + PoolingKernelSize: v.PoolingKernelSize, MMEmbedDim: v.MMEmbedDim, + MMPosembSize: v.MMPosembSize, ModelPatchSize: v.ModelPatchSize, RMSNormEps: 1e-6, + } + if err := validateUnifiedArchitecture(a, requireTextWidth); err != nil { + return Architecture{}, err + } + return a, nil + } if v.HiddenSize == 0 { v.HiddenSize = 768 } @@ -151,6 +193,57 @@ func projectVisionArchitecture(cfg ConfigFile, requireTextWidth bool) (Architect return a, nil } +func isUnifiedConfig(cfg ConfigFile) bool { + if strings.EqualFold(cfg.ModelType, "gemma4_unified") || strings.EqualFold(cfg.VisionConfig.ModelType, "gemma4_unified_vision") { + return true + } + for _, architecture := range cfg.Architectures { + switch strings.ToLower(architecture) { + case "gemma4unifiedforcausallm", "gemma4unifiedforconditionalgeneration": + return true + } + } + return false +} + +func validateUnifiedArchitecture(a Architecture, requireTextWidth bool) error { + if a.TextHiddenSize < 0 || a.TextHiddenSize > MaxVisionHidden || (requireTextWidth && a.TextHiddenSize == 0) { + return fmt.Errorf("invalid Gemma4 text hidden_size %d", a.TextHiddenSize) + } + for _, f := range []struct { + name string + value, max int + }{ + {"mm_embed_dim", a.MMEmbedDim, MaxVisionHidden}, + {"mm_posemb_size", a.MMPosembSize, MaxPositionEntries}, + {"model_patch_size", a.ModelPatchSize, MaxImageDimension}, + {"num_soft_tokens", a.DefaultOutputLength, MaxVisionSoftTokens}, + {"patch_size", a.PatchSize, MaxImageDimension}, + {"pooling_kernel_size", a.PoolingKernelSize, MaxImageDimension}, + } { + if f.value <= 0 || f.value > f.max { + return fmt.Errorf("invalid Gemma4 unified vision %s %d (limit %d)", f.name, f.value, f.max) + } + } + expectedPatch, ok := checkedProduct(MaxImageDimension, int64(a.PatchSize), int64(a.PoolingKernelSize)) + if !ok || expectedPatch != int64(a.ModelPatchSize) { + return fmt.Errorf("unified model_patch_size %d does not match patch_size * pooling_kernel_size (%d)", a.ModelPatchSize, expectedPatch) + } + if _, ok := checkedProduct(math.MaxInt32, 3, int64(a.ModelPatchSize), int64(a.ModelPatchSize)); !ok { + return fmt.Errorf("invalid Gemma4 unified patch dimension") + } + if _, ok := checkedProduct(3*MaxResizePixels, int64(a.DefaultOutputLength), int64(a.ModelPatchSize), int64(a.ModelPatchSize), 3); !ok { + return fmt.Errorf("Gemma4 unified patch allocation exceeds %d values", 3*MaxResizePixels) + } + if _, ok := checkedProduct(MaxPositionValues, int64(a.DefaultOutputLength), 2); !ok { + return fmt.Errorf("Gemma4 unified position allocation exceeds %d values", MaxPositionValues) + } + if a.MMPosembSize < a.DefaultOutputLength { + return fmt.Errorf("Gemma4 unified position table %d is smaller than token count %d", a.MMPosembSize, a.DefaultOutputLength) + } + return nil +} + func validateArchitecture(a Architecture, requireTextWidth bool) error { if a.TextHiddenSize < 0 || a.TextHiddenSize > MaxVisionHidden || (requireTextWidth && a.TextHiddenSize == 0) { return fmt.Errorf("invalid Gemma4 text hidden_size %d", a.TextHiddenSize) @@ -267,6 +360,9 @@ func validateInventory(cfg ConfigFile, tensors map[string]TensorDescriptor, mode if a.TextHiddenSize <= 0 { return fmt.Errorf("missing Gemma4 text hidden_size") } + if a.Unified { + return validateUnifiedInventory(cfg, a, tensors, mode) + } names := make([]string, 0, len(tensors)) for name := range tensors { names = append(names, name) @@ -331,6 +427,9 @@ func validateInventory(cfg ConfigFile, tensors map[string]TensorDescriptor, mode } func validateNames(a Architecture, names []string, mode inventoryMode) error { + if a.Unified { + return validateUnifiedNames(names, mode) + } prefix, ok := visionPrefix(names, mode) if !ok { return fmt.Errorf("missing vision_tower.patch_embedder.input_proj weight") @@ -376,6 +475,54 @@ func validateNames(a Architecture, names []string, mode inventoryMode) error { return fmt.Errorf("missing embed_vision.embedding_projection weight") } +func validateUnifiedNames(names []string, mode inventoryMode) error { + for _, name := range []string{ + "model.vision_embedder.patch_ln1.weight", "model.vision_embedder.patch_ln1.bias", + "model.vision_embedder.patch_dense.bias", "model.vision_embedder.patch_ln2.weight", + "model.vision_embedder.patch_ln2.bias", "model.vision_embedder.pos_embedding", + "model.vision_embedder.pos_norm.weight", "model.vision_embedder.pos_norm.bias", + } { + if !slices.Contains(names, name) { + return fmt.Errorf("missing %s", name) + } + } + for _, path := range []string{"model.vision_embedder.patch_dense", "model.embed_vision.embedding_projection"} { + if !linearNamePresent(names, path, mode) { + return fmt.Errorf("missing %s weight", path) + } + } + return nil +} + +func validateUnifiedInventory(cfg ConfigFile, a Architecture, tensors map[string]TensorDescriptor, mode inventoryMode) error { + names := make([]string, 0, len(tensors)) + for name := range tensors { + names = append(names, name) + } + if err := validateUnifiedNames(names, mode); err != nil { + return err + } + patchDim, _ := checkedProduct(math.MaxInt32, 3, int64(a.ModelPatchSize), int64(a.ModelPatchSize)) + for name, shape := range map[string][]int32{ + "model.vision_embedder.patch_ln1.weight": {int32(patchDim)}, + "model.vision_embedder.patch_ln1.bias": {int32(patchDim)}, + "model.vision_embedder.patch_dense.bias": {int32(a.MMEmbedDim)}, + "model.vision_embedder.patch_ln2.weight": {int32(a.MMEmbedDim)}, + "model.vision_embedder.patch_ln2.bias": {int32(a.MMEmbedDim)}, + "model.vision_embedder.pos_embedding": {int32(a.MMPosembSize), 2, int32(a.MMEmbedDim)}, + "model.vision_embedder.pos_norm.weight": {int32(a.MMEmbedDim)}, + "model.vision_embedder.pos_norm.bias": {int32(a.MMEmbedDim)}, + } { + if err := requireDense(tensors, name, shape); err != nil { + return err + } + } + if err := requireLinearDescriptor(tensors, "model.vision_embedder.patch_dense", []int32{int32(a.MMEmbedDim), int32(patchDim)}, mode, cfg); err != nil { + return err + } + return requireLinearDescriptor(tensors, "model.embed_vision.embedding_projection", []int32{int32(a.TextHiddenSize), int32(a.MMEmbedDim)}, mode, cfg) +} + func visionPrefix(names []string, mode inventoryMode) (string, bool) { for _, p := range []string{"", "model."} { if linearNamePresent(names, p+"vision_tower.patch_embedder.input_proj", mode) { diff --git a/x/models/gemma4/metadata/vision_test.go b/x/models/gemma4/metadata/vision_test.go index ff4de346d1e..7959303256e 100644 --- a/x/models/gemma4/metadata/vision_test.go +++ b/x/models/gemma4/metadata/vision_test.go @@ -386,3 +386,125 @@ func TestVisionInventoryDescriptorAndConfigMatrix(t *testing.T) { } } } + +func unifiedVisionConfig() ConfigFile { + return ConfigFile{ + Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}, + ModelType: "gemma4_unified", TextConfig: TextConfig{HiddenSize: 5}, + VisionConfig: &VisionConfig{ModelType: "gemma4_unified_vision", MMEmbedDim: 3, MMPosembSize: 4, ModelPatchSize: 2, NumSoftTokens: 2, PatchSize: 1, PoolingKernelSize: 2}, + } +} + +func unifiedVisionDescriptors() map[string]TensorDescriptor { + return map[string]TensorDescriptor{ + "model.vision_embedder.patch_ln1.weight": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_ln1.bias": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_dense.weight": {Dtype: "F32", Shape: []int32{3, 12}}, + "model.vision_embedder.patch_dense.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_embedding": {Dtype: "F32", Shape: []int32{4, 2, 3}}, + "model.vision_embedder.pos_norm.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_norm.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "F32", Shape: []int32{5, 3}}, + } +} + +func TestUnifiedVisionInventoryDispatchAndBoundaries(t *testing.T) { + cfg, valid := unifiedVisionConfig(), unifiedVisionDescriptors() + for name, validate := range map[string]func(ConfigFile, map[string]TensorDescriptor) error{ + "source": ValidateVisionSourceInventory, "installed": ValidateVisionInstalledInventory, "runtime": ValidateVisionRuntimeInventory, + } { + t.Run(name, func(t *testing.T) { + if err := validate(cfg, valid); err != nil { + t.Fatalf("complete unified inventory: %v", err) + } + for _, tensor := range []string{"model.vision_embedder.patch_ln1.bias", "model.vision_embedder.patch_dense.weight", "model.embed_vision.embedding_projection.weight"} { + partial := cloneDescriptors(valid) + delete(partial, tensor) + if err := validate(cfg, partial); err == nil { + t.Fatalf("missing %s accepted", tensor) + } + } + bad := cloneDescriptors(valid) + d := bad["model.vision_embedder.pos_embedding"] + d.Shape = []int32{4, 3, 3} + bad["model.vision_embedder.pos_embedding"] = d + if err := validate(cfg, bad); err == nil { + t.Fatal("invalid position shape accepted") + } + }) + } + + for _, mutate := range []func(*ConfigFile){ + func(c *ConfigFile) { c.TextConfig.HiddenSize = 0 }, + func(c *ConfigFile) { c.VisionConfig.MMEmbedDim = MaxVisionHidden + 1 }, + func(c *ConfigFile) { c.VisionConfig.MMPosembSize = 1 }, + func(c *ConfigFile) { c.VisionConfig.ModelPatchSize = 3 }, + func(c *ConfigFile) { c.VisionConfig.NumSoftTokens = MaxVisionSoftTokens + 1 }, + } { + bad := cfg + v := *cfg.VisionConfig + bad.VisionConfig = &v + mutate(&bad) + if _, err := ProjectVisionArchitecture(bad); err == nil { + t.Fatalf("invalid unified architecture accepted: %+v", bad.VisionConfig) + } + } + + boundary := cfg + boundaryVision := *cfg.VisionConfig + boundaryVision.ModelPatchSize = 8 + boundaryVision.PatchSize = 8 + boundaryVision.PoolingKernelSize = 1 + boundaryVision.NumSoftTokens = MaxVisionSoftTokens + boundaryVision.MMPosembSize = MaxVisionSoftTokens + boundary.VisionConfig = &boundaryVision + if _, err := ProjectVisionArchitecture(boundary); err != nil { + t.Fatalf("exact unified patch allocation boundary: %v", err) + } + over := boundary + overVision := boundaryVision + overVision.ModelPatchSize, overVision.PatchSize = 9, 9 + over.VisionConfig = &overVision + if _, err := ProjectVisionArchitecture(over); err == nil || !strings.Contains(err.Error(), "patch allocation") { + t.Fatalf("over unified patch allocation error = %v", err) + } + + near := cfg + near.ModelType = "gemma4_unified_extra" + near.Architectures = []string{"NotGemma4UnifiedForConditionalGeneration"} + v := *cfg.VisionConfig + v.ModelType = "gemma4_unified_vision_extra" + near.VisionConfig = &v + a, err := ProjectVisionArchitecture(near) + if err != nil { + t.Fatal(err) + } + if a.Unified { + t.Fatal("near-match unified identifiers dispatched as unified") + } + + packedCfg := cfg + packedVision := *cfg.VisionConfig + packedVision.ModelPatchSize, packedVision.PatchSize = 4, 2 + packedCfg.VisionConfig = &packedVision + packed := cloneDescriptors(valid) + packed["model.vision_embedder.patch_ln1.weight"] = TensorDescriptor{Dtype: "F32", Shape: []int32{48}} + packed["model.vision_embedder.patch_ln1.bias"] = TensorDescriptor{Dtype: "F32", Shape: []int32{48}} + const dense = "model.vision_embedder.patch_dense" + delete(packed, dense+".weight") + packed[dense+".weight_packed"] = TensorDescriptor{Dtype: "U8", Shape: []int32{3, 24}} + packed[dense+".weight_scale"] = TensorDescriptor{Dtype: "F8_E4M3", Shape: []int32{3, 3}} + packed[dense+".weight_global_scale"] = TensorDescriptor{Dtype: "F32"} + if err := ValidateVisionSourceInventory(packedCfg, packed); err != nil { + t.Fatalf("packed unified source: %v", err) + } + if err := ValidateVisionInstalledInventory(packedCfg, packed); err == nil { + t.Fatal("source-only packed unified inventory accepted as installed") + } + delete(packed, dense+".weight_global_scale") + if err := ValidateVisionSourceInventory(packedCfg, packed); err == nil { + t.Fatal("incomplete packed unified inventory accepted") + } +} diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go index de1e7fdd13f..f8699537c13 100644 --- a/x/models/gemma4/vision.go +++ b/x/models/gemma4/vision.go @@ -17,6 +17,7 @@ import ( _ "image/png" "io" "math" + "slices" xdraw "golang.org/x/image/draw" @@ -54,6 +55,8 @@ type VisionRopeParameters struct { } type VisionConfig struct { + Unified bool `json:"-"` + ModelType string `json:"model_type"` HiddenSize int32 `json:"hidden_size"` IntermediateSize int32 `json:"intermediate_size"` NumHiddenLayers int32 `json:"num_hidden_layers"` @@ -67,9 +70,18 @@ type VisionConfig struct { PoolingKernelSize int32 `json:"pooling_kernel_size"` UseClippedLinears bool `json:"use_clipped_linears"` Standardize bool `json:"standardize"` + MMEmbedDim int32 `json:"mm_embed_dim"` + MMPosembSize int32 `json:"mm_posemb_size"` + ModelPatchSize int32 `json:"model_patch_size"` + NumSoftTokens int32 `json:"num_soft_tokens"` + OutputProjDims int32 `json:"output_proj_dims"` RopeParameters VisionRopeParameters `json:"rope_parameters"` } +func (c *VisionConfig) unified() bool { + return c != nil && c.Unified +} + type gemma4MediaTokens struct { BOI string Image string @@ -83,6 +95,12 @@ type gemma4ImageInput struct { PatchWidth int PatchHeight int SoftTokens int + Patches []float32 + Positions []int32 +} + +type gemma4MediaLayout struct { + ImageSpans [][2]int } type gemma4MediaPayload struct { @@ -140,6 +158,15 @@ type VisionModel struct { StdScale *mlx.Array } +type UnifiedVisionEmbedder struct { + PatchLN1 *nn.LayerNorm + PatchDense nn.LinearLayer + PatchLN2 *nn.LayerNorm + PosEmbedding *mlx.Array + PosNorm *nn.LayerNorm + PatchDim int32 +} + type MultimodalEmbedder struct { Projection nn.LinearLayer Eps float32 @@ -171,6 +198,7 @@ func parseVisionConfig(configData []byte) (*VisionConfig, error) { return nil, err } cfg := *wrapped.VisionConfig + cfg.Unified = architecture.Unified cfg.HiddenSize = int32(architecture.HiddenSize) cfg.IntermediateSize = int32(architecture.IntermediateSize) cfg.NumHiddenLayers = int32(architecture.NumHiddenLayers) @@ -183,6 +211,13 @@ func parseVisionConfig(configData []byte) (*VisionConfig, error) { cfg.PositionEmbeddingSize = int32(architecture.PositionEmbeddingSize) cfg.PoolingKernelSize = int32(architecture.PoolingKernelSize) cfg.RopeParameters.RopeTheta = float32(architecture.RopeTheta) + cfg.MMEmbedDim = int32(architecture.MMEmbedDim) + cfg.MMPosembSize = int32(architecture.MMPosembSize) + cfg.ModelPatchSize = int32(architecture.ModelPatchSize) + if architecture.Unified { + cfg.DefaultOutputLength = int32(architecture.DefaultOutputLength) + cfg.RMSNormEps = float32(architecture.RMSNormEps) + } return &cfg, nil } @@ -244,19 +279,31 @@ func metadataConfigFromVision(cfg *VisionConfig, textHidden int) gemma4metadata. return gemma4metadata.ConfigFile{} } v := &gemma4metadata.VisionConfig{ + ModelType: cfg.ModelType, HiddenSize: int(cfg.HiddenSize), IntermediateSize: int(cfg.IntermediateSize), NumHiddenLayers: int(cfg.NumHiddenLayers), NumAttentionHeads: int(cfg.NumAttentionHeads), NumKeyValueHeads: int(cfg.NumKeyValueHeads), HeadDim: int(cfg.HeadDim), RMSNormEps: float64(cfg.RMSNormEps), DefaultOutputLength: int(cfg.DefaultOutputLength), PatchSize: int(cfg.PatchSize), PositionEmbeddingSize: int(cfg.PositionEmbeddingSize), PoolingKernelSize: int(cfg.PoolingKernelSize), UseClippedLinears: cfg.UseClippedLinears, Standardize: cfg.Standardize, + MMEmbedDim: int(cfg.MMEmbedDim), MMPosembSize: int(cfg.MMPosembSize), ModelPatchSize: int(cfg.ModelPatchSize), + NumSoftTokens: int(cfg.DefaultOutputLength), OutputProjDims: int(cfg.OutputProjDims), + } + if cfg.Unified { + v.ModelType = "gemma4_unified_vision" } v.RopeParameters.RopeTheta = float64(cfg.RopeParameters.RopeTheta) return gemma4metadata.ConfigFile{TextConfig: gemma4metadata.TextConfig{HiddenSize: textHidden}, VisionConfig: v} } func validateGemma4VisionWeights(tensors map[string]*mlx.Array, cfg *VisionConfig, textHidden int, tq map[string]*model.TensorQuantInfo) (bool, error) { - if firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") == nil { - if firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight_packed", "model.vision_tower.patch_embedder.input_proj.weight_packed") != nil { + sentinel := firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") + packedSentinel := firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight_packed", "model.vision_tower.patch_embedder.input_proj.weight_packed") + if cfg != nil && cfg.unified() { + sentinel = firstNonNil(tensors, "model.vision_embedder.patch_ln1.weight") + packedSentinel = firstNonNil(tensors, "model.vision_embedder.patch_dense.weight_packed") + } + if sentinel == nil { + if packedSentinel != nil { return false, fmt.Errorf("runtime contains source-only Gemma4 vision packed sentinel") } return false, nil @@ -282,6 +329,43 @@ func validateGemma4VisionWeights(tensors map[string]*mlx.Array, cfg *VisionConfi return true, nil } +func loadUnifiedVisionEmbedder(tensors map[string]*mlx.Array, cfg *VisionConfig, groupSize, bits int, mode string, tq map[string]*model.TensorQuantInfo) (*UnifiedVisionEmbedder, error) { + const prefix = "model.vision_embedder." + patchDim, ok := checkedPositiveProduct(math.MaxInt32, int64(cfg.ModelPatchSize), int64(cfg.ModelPatchSize), 3) + if !ok { + return nil, errors.New("invalid Gemma4 unified patch dimension") + } + linears := model.NewLinearFactory(tensors, groupSize, bits, mode, tq) + patchDense := linears.Make(prefix + "patch_dense") + if patchDense == nil { + return nil, errors.New("missing Gemma4 unified patch projection") + } + makeLayerNorm := func(path string, width int) (*nn.LayerNorm, error) { + weight, bias := tensors[path+".weight"], tensors[path+".bias"] + if weight == nil || bias == nil || !slices.Equal(weight.Dims(), []int{width}) || !slices.Equal(bias.Dims(), []int{width}) { + return nil, fmt.Errorf("invalid Gemma4 unified layer norm %s", path) + } + return &nn.LayerNorm{Weight: weight, Bias: bias, Eps: 1e-5}, nil + } + patchLN1, err := makeLayerNorm(prefix+"patch_ln1", int(patchDim)) + if err != nil { + return nil, err + } + patchLN2, err := makeLayerNorm(prefix+"patch_ln2", int(cfg.MMEmbedDim)) + if err != nil { + return nil, err + } + posNorm, err := makeLayerNorm(prefix+"pos_norm", int(cfg.MMEmbedDim)) + if err != nil { + return nil, err + } + pos := tensors[prefix+"pos_embedding"] + if pos == nil || !slices.Equal(pos.Dims(), []int{int(cfg.MMPosembSize), 2, int(cfg.MMEmbedDim)}) { + return nil, errors.New("invalid Gemma4 unified position embedding") + } + return &UnifiedVisionEmbedder{PatchLN1: patchLN1, PatchDense: patchDense, PatchLN2: patchLN2, PosEmbedding: pos, PosNorm: posNorm, PatchDim: int32(patchDim)}, nil +} + func resolveVisionPrefix(tensors map[string]*mlx.Array) string { if tensors["vision_tower.patch_embedder.input_proj.weight"] != nil { return "" @@ -425,6 +509,7 @@ func (m *MultimodalEmbedder) Forward(x *mlx.Array) *mlx.Array { // in stream order and remains a separate cache-identity item. func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { prepared := &base.PreparedRequest{} + var layout gemma4MediaLayout for source, seg := range segments { if err := ctx.Err(); err != nil { return nil, err @@ -433,7 +518,7 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas prepared.Tokens = append(prepared.Tokens, seg.Tokens...) continue } - if m.VisionConfig == nil || m.Vision == nil || m.EmbedVision == nil { + if m.VisionConfig == nil || (m.Vision == nil && m.UnifiedVision == nil) || m.EmbedVision == nil { return nil, fmt.Errorf("this model does not support %s input", seg.Kind) } if seg.Kind != "image" { @@ -454,13 +539,20 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas prepared.Tokens = append(prepared.Tokens, m.EOITokenIDValue) geom := *img - pixels := geom.Pixels + mediaData := geom.Pixels + dims := []int{1, 3, geom.Height, geom.Width} geom.Pixels = nil + if m.UnifiedVision != nil { + mediaData = geom.Patches + dims = []int{1, geom.SoftTokens, int(m.UnifiedVision.PatchDim)} + geom.Patches = nil + layout.ImageSpans = append(layout.ImageSpans, [2]int{start + imageStart, start + imageEnd}) + } item := base.PreparedItem{ Range: [2]int{start, len(prepared.Tokens)}, Source: source, - MediaData: pixels, - Dims: []int{1, 3, geom.Height, geom.Width}, + MediaData: mediaData, + Dims: dims, Opaque: gemma4MediaPayload{ Image: geom, ImageStart: imageStart, @@ -472,6 +564,9 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas } prepared.Items = append(prepared.Items, item) } + if len(layout.ImageSpans) > 0 { + prepared.Layout = &layout + } if err := ctx.Err(); err != nil { return nil, err } @@ -482,8 +577,15 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas // MediaData, so pixels are always read from data rather than Opaque. func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { payload := item.Opaque.(gemma4MediaPayload) - pixels := mlx.Reshape(data, 1, 3, int32(payload.Image.Height), int32(payload.Image.Width)) - features := m.EmbedVision.Forward(m.Vision.Forward(pixels, &payload.Image)) + var encoded *mlx.Array + if m.UnifiedVision != nil { + patches := mlx.Reshape(data, 1, int32(payload.Image.SoftTokens), m.UnifiedVision.PatchDim) + encoded = m.UnifiedVision.Forward(patches, &payload.Image) + } else { + pixels := mlx.Reshape(data, 1, 3, int32(payload.Image.Height), int32(payload.Image.Width)) + encoded = m.Vision.Forward(pixels, &payload.Image) + } + features := m.EmbedVision.Forward(encoded) return mlx.Squeeze(features, 0) } @@ -584,32 +686,99 @@ func preprocessGemma4Image(ctx context.Context, data []byte, cfg *VisionConfig, } } - pixels, err := imageToCHWFloat32Context(ctx, resized) - if err != nil { - return nil, err - } patchW := targetW / patchSize patchH := targetH / patchSize patchCount, ok := checkedPositiveProduct(maxIntValue, int64(patchW), int64(patchH)) if !ok { return nil, errors.New("Gemma4 image patch count exceeds platform limits") } - if _, ok := checkedPositiveProduct(maxGemma4PositionValues, patchCount, int64(cfg.HeadDim)); !ok { - return nil, fmt.Errorf("Gemma4 image position allocation exceeds limit %d", maxGemma4PositionValues) + if !cfg.unified() { + if _, ok := checkedPositiveProduct(maxGemma4PositionValues, patchCount, int64(cfg.HeadDim)); !ok { + return nil, fmt.Errorf("Gemma4 image position allocation exceeds limit %d", maxGemma4PositionValues) + } } softTokens := int(patchCount / (int64(pooling) * int64(pooling))) if softTokens <= 0 || softTokens > maxSoftTokens { return nil, fmt.Errorf("Gemma4 image produced %d soft tokens, limit %d", softTokens, maxSoftTokens) } - return &gemma4ImageInput{ - Pixels: pixels, + input := &gemma4ImageInput{ Width: targetW, Height: targetH, PatchWidth: patchW, PatchHeight: patchH, SoftTokens: softTokens, - }, nil + } + if cfg.unified() { + input.Patches, input.Positions, err = imageToUnifiedPatchesContext(ctx, resized, int(cfg.ModelPatchSize)) + if err != nil { + return nil, err + } + expected, ok := checkedPositiveProduct(maxIntValue, int64(softTokens), int64(cfg.ModelPatchSize), int64(cfg.ModelPatchSize), 3) + if !ok || int64(len(input.Patches)) != expected || len(input.Positions) != softTokens*2 { + return nil, errors.New("Gemma4 unified patch count does not match soft token count") + } + } else { + input.Pixels, err = imageToCHWFloat32Context(ctx, resized) + if err != nil { + return nil, err + } + } + return input, nil +} + +func imageToUnifiedPatchesContext(ctx context.Context, img image.Image, patchSize int) ([]float32, []int32, error) { + if patchSize <= 0 { + return nil, nil, errors.New("invalid Gemma4 unified model patch size") + } + b := img.Bounds() + width, height := b.Dx(), b.Dy() + if width%patchSize != 0 || height%patchSize != 0 { + return nil, nil, fmt.Errorf("Gemma4 unified image dimensions %dx%d are not divisible by patch size %d", width, height, patchSize) + } + patchW, patchH := width/patchSize, height/patchSize + patchDim, ok := checkedPositiveProduct(maxIntValue, int64(patchSize), int64(patchSize), 3) + if !ok { + return nil, nil, errors.New("Gemma4 unified patch dimension exceeds platform limits") + } + patchCount, ok := checkedPositiveProduct(maxIntValue, int64(patchW), int64(patchH)) + if !ok || patchCount > math.MaxInt32 { + return nil, nil, errors.New("Gemma4 unified patch count exceeds platform limits") + } + values, ok := checkedPositiveProduct(maxIntValue, patchCount, patchDim) + if !ok { + return nil, nil, errors.New("Gemma4 unified patch allocation exceeds platform limits") + } + positionsCount, ok := checkedPositiveProduct(maxIntValue, patchCount, 2) + if !ok { + return nil, nil, errors.New("Gemma4 unified position allocation exceeds platform limits") + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + patches := make([]float32, int(values)) + if err := ctx.Err(); err != nil { + return nil, nil, err + } + positions := make([]int32, int(positionsCount)) + for py := range patchH { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + for px := range patchW { + patch := py*patchW + px + positions[2*patch], positions[2*patch+1] = int32(px), int32(py) + offset := patch * int(patchDim) + for y := range patchSize { + for x := range patchSize { + r, g, blue, _ := img.At(b.Min.X+px*patchSize+x, b.Min.Y+py*patchSize+y).RGBA() + i := offset + (y*patchSize+x)*3 + patches[i], patches[i+1], patches[i+2] = float32(r)/65535, float32(g)/65535, float32(blue)/65535 + } + } + } + } + return patches, positions, nil } type gemma4CancellationPanic struct{ err error } @@ -805,6 +974,17 @@ func (m *VisionModel) Forward(pixels *mlx.Array, img *gemma4ImageInput) *mlx.Arr return h } +func (m *UnifiedVisionEmbedder) Forward(patches *mlx.Array, img *gemma4ImageInput) *mlx.Array { + hidden := m.PatchLN2.Forward(m.PatchDense.Forward(m.PatchLN1.Forward(patches))) + positions := mlx.FromValues(img.Positions, 1, img.SoftTokens, 2) + x := mlx.Squeeze(mlx.SliceStartStop(positions, []int32{0, 0, 0}, []int32{1, int32(img.SoftTokens), 1}), 2) + y := mlx.Squeeze(mlx.SliceStartStop(positions, []int32{0, 0, 1}, []int32{1, int32(img.SoftTokens), 2}), 2) + tableX := mlx.Squeeze(mlx.SliceStartStop(m.PosEmbedding, []int32{0, 0, 0}, []int32{int32(m.PosEmbedding.Dim(0)), 1, int32(m.PosEmbedding.Dim(2))}), 1) + tableY := mlx.Squeeze(mlx.SliceStartStop(m.PosEmbedding, []int32{0, 1, 0}, []int32{int32(m.PosEmbedding.Dim(0)), 2, int32(m.PosEmbedding.Dim(2))}), 1) + pos := mlx.Add(tableX.TakeAxis(x, 0), tableY.TakeAxis(y, 0)).AsType(hidden.DType()) + return m.PosNorm.Forward(mlx.Add(hidden, pos)) +} + func (m *VisionModel) positionArrays(img *gemma4ImageInput) visionPositionArrays { L := img.PatchHeight * img.PatchWidth xs := make([]int32, L) diff --git a/x/models/gemma4/vision_test.go b/x/models/gemma4/vision_test.go index 6f8bf4bad79..c237c8745db 100644 --- a/x/models/gemma4/vision_test.go +++ b/x/models/gemma4/vision_test.go @@ -49,6 +49,28 @@ func TestParseVisionConfigDefaults(t *testing.T) { } } +func TestParseUnifiedVisionConfig(t *testing.T) { + cfg, err := parseVisionConfig([]byte(`{"text_config":{"hidden_size":5},"vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`)) + if err != nil { + t.Fatal(err) + } + if !cfg.unified() || cfg.MMEmbedDim != 3 || cfg.MMPosembSize != 4 || cfg.ModelPatchSize != 2 || cfg.DefaultOutputLength != 2 { + t.Fatalf("unexpected unified config: %+v", cfg) + } + topLevel, err := parseVisionConfig([]byte(`{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"hidden_size":5},"vision_config":{"mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`)) + if err != nil || !topLevel.unified() { + t.Fatalf("top-level unified dispatch = %+v, %v", topLevel, err) + } + for _, input := range []string{ + `{"text_config":{"hidden_size":5},"vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":3,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`, + `{"text_config":{"hidden_size":5},"vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":1,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`, + } { + if _, err := parseVisionConfig([]byte(input)); err == nil { + t.Fatalf("invalid unified config accepted: %s", input) + } + } +} + func TestParseVisionConfigRejectsUnsafeDimensions(t *testing.T) { tests := []struct { name string @@ -201,6 +223,40 @@ func TestImageToCHWFloat32UsesBoundsAndChannelOrder(t *testing.T) { } } +func TestImageToUnifiedPatchesUsesHWCModelPatchOrder(t *testing.T) { + img := image.NewRGBA(image.Rect(10, 20, 14, 22)) + for y := range 2 { + for x := range 4 { + img.SetRGBA(10+x, 20+y, color.RGBA{R: uint8(1 + x + 4*y), G: uint8(11 + x + 4*y), B: uint8(21 + x + 4*y), A: 255}) + } + } + patches, positions, err := imageToUnifiedPatchesContext(context.Background(), img, 2) + if err != nil { + t.Fatal(err) + } + if !slices.Equal(positions, []int32{0, 0, 1, 0}) { + t.Fatalf("positions = %v", positions) + } + want := []uint8{1, 11, 21, 2, 12, 22, 5, 15, 25, 6, 16, 26, 3, 13, 23, 4, 14, 24, 7, 17, 27, 8, 18, 28} + if len(patches) != len(want) { + t.Fatalf("patch values = %d, want %d", len(patches), len(want)) + } + for i, value := range want { + if math.Abs(float64(patches[i]-float32(value)/255)) > 1e-6 { + t.Fatalf("patch[%d] = %v, want %v", i, patches[i], float32(value)/255) + } + } + if _, _, err := imageToUnifiedPatchesContext(context.Background(), img, 3); err == nil { + t.Fatal("non-divisible patch geometry accepted") + } + canceled, cancel := context.WithCancel(context.Background()) + cancel() + patches, positions, err = imageToUnifiedPatchesContext(canceled, img, 2) + if !errors.Is(err, context.Canceled) || patches != nil || positions != nil { + t.Fatalf("canceled patch conversion = (%v, %v, %v), want (nil, nil, context.Canceled)", patches, positions, err) + } +} + func TestPreprocessGemma4ImageSoftTokenBudget(t *testing.T) { src := image.NewRGBA(image.Rect(0, 0, 8, 4)) src.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) @@ -291,6 +347,37 @@ func TestPrepareMediaPreservesOrderedImageItems(t *testing.T) { } } +func TestPrepareMediaUnifiedUsesPatchDataAndLayout(t *testing.T) { + img := image.NewRGBA(image.Rect(0, 0, 4, 2)) + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + t.Fatal(err) + } + cfg, err := parseVisionConfig([]byte(`{"text_config":{"hidden_size":5},"vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`)) + if err != nil { + t.Fatal(err) + } + m := &Model{ + TextConfig: &TextConfig{ImageTokenIDValue: 10, BOITokenIDValue: 11, EOITokenIDValue: 12, VisionSoftTokens: 2}, + VisionConfig: cfg, UnifiedVision: &UnifiedVisionEmbedder{PatchDim: 12}, EmbedVision: &MultimodalEmbedder{}, + } + got, err := m.PrepareMedia(context.Background(), []base.Segment{{Tokens: []int32{1}}, {Kind: "image", Data: buf.Bytes()}}) + if err != nil { + t.Fatal(err) + } + if len(got.Items) != 1 || !slices.Equal(got.Items[0].Dims, []int{1, 2, 12}) || len(got.Items[0].MediaData) != 24 { + t.Fatalf("unified item = %#v", got.Items) + } + payload := got.Items[0].Opaque.(gemma4MediaPayload) + if payload.Image.Patches != nil || !slices.Equal(payload.Image.Positions, []int32{0, 0, 1, 0}) { + t.Fatalf("unified payload = %#v", payload.Image) + } + layout, ok := got.Layout.(*gemma4MediaLayout) + if !ok || !slices.Equal(layout.ImageSpans, [][2]int{{2, 4}}) { + t.Fatalf("unified layout = %#v", got.Layout) + } +} + func TestPrepareMediaSequentialRequestsAreIsolated(t *testing.T) { pngData := func(c color.RGBA) []byte { t.Helper() @@ -379,6 +466,32 @@ func TestImageToCHWFloat32CancellationAllocationBoundary(t *testing.T) { } } +func TestImageToUnifiedPatchesCancellationAllocationBoundary(t *testing.T) { + img := image.NewRGBA(image.Rect(0, 0, 64, 64)) + measure := func(target int) float64 { + ctx := &nthCancelContext{target: target} + return testing.AllocsPerRun(100, func() { + ctx.calls = 0 + patches, positions, err := imageToUnifiedPatchesContext(ctx, img, 1) + if !errors.Is(err, context.Canceled) || patches != nil || positions != nil { + t.Fatalf("imageToUnifiedPatchesContext() = (%v, %v, %v), want (nil, nil, context.Canceled)", patches, positions, err) + } + }) + } + preAllocation := measure(1) + betweenAllocations := measure(2) + conversionLoop := measure(3) + if preAllocation != 0 { + t.Fatalf("pre-allocation cancellation allocated %.1f objects", preAllocation) + } + if betweenAllocations <= preAllocation { + t.Fatalf("between-allocation cancellation allocations = %.1f, want more than pre-allocation %.1f", betweenAllocations, preAllocation) + } + if conversionLoop <= betweenAllocations { + t.Fatalf("conversion-loop cancellation allocations = %.1f, want more than between-allocation %.1f", conversionLoop, betweenAllocations) + } +} + type blockingCancelContext struct { target int calls int From 8f3359b315792a79be5af33c99f5ce3eac48e197 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sat, 15 Aug 2026 23:32:47 +0000 Subject: [PATCH 46/58] gemma4: retain and gate MLX audio tensors Validate bounded released Gemma 4 audio configuration and complete source and installed tensor inventories before advertising audio. Inspect installed descriptors and payloads with the accepted media safeguards while keeping create, show, and list capability views consistent. Co-authored-by: Codex --- server/images.go | 68 +++++- server/images_test.go | 296 +++++++++++++++++++++++++ server/model_list_cache.go | 11 +- x/create/client/create.go | 16 +- x/create/client/create_test.go | 130 ++++++++++- x/models/gemma4/metadata/audio.go | 227 +++++++++++++++++++ x/models/gemma4/metadata/audio_test.go | 167 ++++++++++++++ x/models/gemma4/metadata/vision.go | 1 + 8 files changed, 905 insertions(+), 11 deletions(-) create mode 100644 x/models/gemma4/metadata/audio.go create mode 100644 x/models/gemma4/metadata/audio_test.go diff --git a/server/images.go b/server/images.go index 610c860687a..d968312e5da 100644 --- a/server/images.go +++ b/server/images.go @@ -87,6 +87,8 @@ type Model struct { TensorLayerNames []string Gemma4VisionConfig *gemma4metadata.ConfigFile `json:"-"` Gemma4VisionTensors map[string]gemma4metadata.TensorDescriptor `json:"-"` + Gemma4AudioConfig *gemma4metadata.ConfigFile `json:"-"` + Gemma4AudioTensors map[string]gemma4metadata.TensorDescriptor `json:"-"` System string License []string Digest string @@ -515,7 +517,11 @@ func suppressGemma4SafetensorsVisionCapability(m *Model) bool { func suppressAudioCapability(m *Model, arch string) bool { if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" { - return true + if !isLocalGemma4SafetensorsConfig(m.Config) { + return true + } + return m.Gemma4AudioConfig == nil || + gemma4metadata.ValidateAudioInstalledInventory(*m.Gemma4AudioConfig, m.Gemma4AudioTensors) != nil } if m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" { return true @@ -558,6 +564,59 @@ func hasGemma4VisionTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifes return err == nil && gemma4metadata.ValidateVisionInstalledInventory(cfg, tensors) == nil } +func hasGemma4AudioTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifest.Layer) bool { + tensors, err := gemma4AudioTensorDescriptors(layers) + return err == nil && gemma4metadata.ValidateAudioInstalledInventory(cfg, tensors) == nil +} + +func gemma4AudioTensorDescriptors(layers []manifest.Layer) (map[string]gemma4metadata.TensorDescriptor, error) { + tensors := make(map[string]gemma4metadata.TensorDescriptor) + descriptorWork := 0 + for _, layer := range layers { + if layer.MediaType != manifest.MediaTypeImageTensor { + continue + } + if !strings.HasPrefix(layer.Name, "model.audio_tower.") && !strings.HasPrefix(layer.Name, "model.embed_audio.") { + continue + } + filename, err := manifest.BlobsPath(layer.Digest) + if err != nil { + return nil, err + } + ext, err := openGemma4TensorLayer(filename) + if err != nil { + return nil, fmt.Errorf("open audio tensor layer %s: %w", layer.Name, err) + } + names := ext.ListTensors() + if len(names) > maxGemma4VisionDescriptors-len(tensors) { + ext.Close() + return nil, fmt.Errorf("Gemma4 audio tensor inventory exceeds %d descriptors", maxGemma4VisionDescriptors) + } + for _, name := range names { + tensor, err := ext.GetTensor(name) + if err != nil { + ext.Close() + return nil, err + } + if _, exists := tensors[name]; exists { + ext.Close() + return nil, fmt.Errorf("duplicate audio tensor %s", name) + } + work := len(name) + len(tensor.Dtype) + len(tensor.Shape)*4 + if work > maxGemma4VisionDescriptorWork-descriptorWork { + ext.Close() + return nil, fmt.Errorf("Gemma4 audio tensor inventory exceeds descriptor work limit %d", maxGemma4VisionDescriptorWork) + } + descriptorWork += work + tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + } + if err := ext.Close(); err != nil { + return nil, err + } + } + return tensors, nil +} + func gemma4VisionTensorDescriptors(layers []manifest.Layer) (map[string]gemma4metadata.TensorDescriptor, error) { tensors := make(map[string]gemma4metadata.TensorDescriptor) descriptorWork := 0 @@ -814,6 +873,9 @@ func (m *Model) CheckCapabilities(want ...model.Capability) error { if slices.Contains(errs, errCapabilityVision) && suppressGemma4SafetensorsVisionCapability(m) { return fmt.Errorf("%w. Recreate or pull the model so it includes Gemma 4 vision tensor layers", err) } + if slices.Contains(errs, errCapabilityAudio) && isLocalGemma4SafetensorsConfig(m.Config) && suppressAudioCapability(m, "") { + return fmt.Errorf("%w. Recreate or pull the model so it includes Gemma 4 audio tensor layers", err) + } return err } @@ -943,9 +1005,13 @@ func GetModel(name string) (*Model, error) { var cfg gemma4metadata.ConfigFile if err := mf.ReadConfigJSON("config.json", &cfg); err == nil { m.Gemma4VisionConfig = &cfg + m.Gemma4AudioConfig = &cfg if tensors, err := gemma4VisionTensorDescriptors(mf.Layers); err == nil { m.Gemma4VisionTensors = tensors } + if tensors, err := gemma4AudioTensorDescriptors(mf.Layers); err == nil { + m.Gemma4AudioTensors = tensors + } } } diff --git a/server/images_test.go b/server/images_test.go index 4a30fa38064..62bfe651024 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -598,6 +598,21 @@ func TestModelCapabilities(t *testing.T) { }, expectedCaps: []model.Capability{model.CapabilityVision}, }, + { + name: "gemma4 safetensors exposes complete audio", + model: Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererSmall, + Capabilities: []string{"audio"}, + }, + TensorLayerNames: gemma4AudioTensorNames(1), + Gemma4AudioConfig: gemma4AudioConfig(1), + Gemma4AudioTensors: testGemma4AudioTensorDescriptors(1), + Template: chatTemplate, + }, + expectedCaps: []model.Capability{model.CapabilityAudio}, + }, } // compare two slices of model.Capability regardless of order @@ -694,6 +709,247 @@ func TestGemma4SafetensorsVisionCapabilityRequiresTensorLayers(t *testing.T) { } } +func TestGemma4SafetensorsAudioCapabilityRequiresCompleteInventory(t *testing.T) { + complete := Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "audio"}, + }, + TensorLayerNames: gemma4AudioTensorNames(1), + Gemma4AudioConfig: gemma4AudioConfig(1), + Gemma4AudioTensors: testGemma4AudioTensorDescriptors(1), + } + if !slices.Contains(complete.Capabilities(), model.CapabilityAudio) { + t.Fatal("complete audio inventory did not expose audio") + } + + partial := complete + partial.TensorLayerNames = slices.Clone(complete.TensorLayerNames) + partial.Gemma4AudioTensors = maps.Clone(complete.Gemma4AudioTensors) + missing := "model.audio_tower.layers.0.self_attn.q_proj.input_max" + partial.TensorLayerNames = slices.DeleteFunc(partial.TensorLayerNames, func(name string) bool { return name == missing }) + delete(partial.Gemma4AudioTensors, missing) + if slices.Contains(partial.Capabilities(), model.CapabilityAudio) { + t.Fatal("partial audio inventory exposed audio") + } + if err := partial.CheckCapabilities(model.CapabilityAudio); err == nil || !strings.Contains(err.Error(), "includes Gemma 4 audio tensor layers") { + t.Fatalf("CheckCapabilities(audio) error = %v, want Gemma 4 audio hint", err) + } + + nearMatch := partial + nearMatch.TensorLayerNames = append(slices.Clone(partial.TensorLayerNames), missing+".extra") + nearMatch.Gemma4AudioTensors = maps.Clone(partial.Gemma4AudioTensors) + nearMatch.Gemma4AudioTensors[missing+".extra"] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{}} + if slices.Contains(nearMatch.Capabilities(), model.CapabilityAudio) { + t.Fatal("near-match audio tensor exposed audio") + } + + malformed := complete + malformed.Gemma4AudioConfig = gemma4AudioConfig(1) + malformed.Gemma4AudioConfig.AudioConfig.NumAttentionHeads = 3 + if slices.Contains(malformed.Capabilities(), model.CapabilityAudio) { + t.Fatal("malformed audio config exposed audio") + } + + remoteCfg := complete.Config + remoteCfg.RemoteHost = "https://example.invalid" + remoteCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityAudio}, remoteCfg) + if slices.Contains(remoteCaps, model.CapabilityAudio) { + t.Fatal("remote Gemma 4 list capability exposed unsupported local audio") + } + otherCfg := remoteCfg + otherCfg.Renderer = "other" + otherCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityAudio}, otherCfg) + if !slices.Contains(otherCaps, model.CapabilityAudio) { + t.Fatal("remote non-Gemma audio capability was suppressed") + } +} + +func TestGemma4InstalledAudioCapabilityDescriptorAndPayloadMatrix(t *testing.T) { + const target = "model.audio_tower.output_proj.weight" + tests := []struct { + name string + edit func(*testing.T, []manifest.Layer) []manifest.Layer + want bool + }{ + {name: "complete", want: true}, + {name: "partial", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == target }) + }}, + {name: "near-match internal name", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4AudioFixtureLayer(t, layers, target, target+".extra", gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{3, 4}}, nil) + }}, + {name: "wrong shape", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4AudioFixtureLayer(t, layers, target, target, gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{4, 3}}, nil) + }}, + {name: "wrong dtype", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4AudioFixtureLayer(t, layers, target, target, gemma4metadata.TensorDescriptor{Dtype: "U8", Shape: []int32{3, 4}}, nil) + }}, + {name: "duplicate descriptor", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return append(layers, gemma4AudioFixtureLayer(t, "model.audio_tower.duplicate.weight", target, gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{3, 4}}, nil)) + }}, + {name: "truncated payload", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4AudioFixtureLayer(t, layers, target, target, gemma4metadata.TensorDescriptor{}, []byte{1, 2, 3, 4}) + }}, + {name: "invalid payload range", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + header, err := json.Marshal(map[string]any{target: map[string]any{"dtype": "F32", "shape": []int{3, 4}, "data_offsets": []int{0, 4}}}) + if err != nil { + t.Fatal(err) + } + data := binary.LittleEndian.AppendUint64(nil, uint64(len(header))) + data = append(data, header...) + data = append(data, make([]byte, 4)...) + return replaceGemma4AudioFixtureLayer(t, layers, target, target, gemma4metadata.TensorDescriptor{}, data) + }}, + {name: "malformed header", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + data := binary.LittleEndian.AppendUint64(nil, 5) + data = append(data, []byte("nope!")...) + return replaceGemma4AudioFixtureLayer(t, layers, target, target, gemma4metadata.TensorDescriptor{}, data) + }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setTestHome(t, t.TempDir()) + layers := gemma4AudioManifestLayers(t) + if tt.edit != nil { + layers = tt.edit(t, layers) + } + cfg := model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererLarge, Capabilities: []string{"completion", "audio"}} + name := "gemma4-audio-" + strings.ReplaceAll(tt.name, " ", "-") + createSafetensorsTestModel(t, name, cfg, layers) + + m, err := GetModel(name) + if err != nil { + t.Fatal(err) + } + showAudio := slices.Contains(m.Capabilities(), model.CapabilityAudio) + mf, err := manifest.ParseNamedManifest(model.ParseName(name)) + if err != nil { + t.Fatal(err) + } + summary, err := buildModelListSummary(model.ParseName(name), mf) + if err != nil { + t.Fatal(err) + } + listAudio := slices.Contains(summary.Capabilities, model.CapabilityAudio) + if showAudio != tt.want || listAudio != tt.want { + t.Fatalf("show/list audio = %t/%t, want %t", showAudio, listAudio, tt.want) + } + }) + } +} + +func TestGemma4InstalledAudioCapabilityRejectsUnboundedConfig(t *testing.T) { + tests := []struct { + name string + edit func(*gemma4metadata.ConfigFile) + }{ + {name: "shape product overflow", edit: func(cfg *gemma4metadata.ConfigFile) { + cfg.AudioConfig.HiddenSize = 1 << 30 + cfg.AudioConfig.NumAttentionHeads = 1 + }}, + {name: "impractical layer count", edit: func(cfg *gemma4metadata.ConfigFile) { cfg.AudioConfig.NumHiddenLayers = 1 << 30 }}, + {name: "impractical convolution channels", edit: func(cfg *gemma4metadata.ConfigFile) { cfg.AudioConfig.SubsamplingConvChannels = []int{2, 1 << 30} }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + setTestHome(t, t.TempDir()) + cfg := gemma4AudioConfig(1) + tt.edit(cfg) + layers := replaceGemma4AudioConfigLayer(t, gemma4AudioManifestLayers(t), cfg) + modelCfg := model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererLarge, Capabilities: []string{"completion", "audio"}} + name := "gemma4-audio-config-" + strings.ReplaceAll(tt.name, " ", "-") + createSafetensorsTestModel(t, name, modelCfg, layers) + + m, err := GetModel(name) + if err != nil { + t.Fatal(err) + } + mf, err := manifest.ParseNamedManifest(model.ParseName(name)) + if err != nil { + t.Fatal(err) + } + summary, err := buildModelListSummary(model.ParseName(name), mf) + if err != nil { + t.Fatal(err) + } + if slices.Contains(m.Capabilities(), model.CapabilityAudio) || slices.Contains(summary.Capabilities, model.CapabilityAudio) { + t.Fatalf("unbounded config exposed show/list audio") + } + }) + } +} + +func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { + t.Helper() + descriptors := testGemma4AudioTensorDescriptors(1) + layers := make([]manifest.Layer, 0, len(descriptors)+1) + for name, descriptor := range descriptors { + layers = append(layers, gemma4AudioFixtureLayer(t, name, name, descriptor, nil)) + } + config, err := json.Marshal(gemma4AudioConfig(1)) + if err != nil { + t.Fatal(err) + } + digest := createTestBlob(t, config) + layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(config)), Name: "config.json"}) + return layers +} + +func replaceGemma4AudioConfigLayer(t *testing.T, layers []manifest.Layer, cfg *gemma4metadata.ConfigFile) []manifest.Layer { + t.Helper() + data, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + for i := range layers { + if layers[i].Name == "config.json" { + digest := createTestBlob(t, data) + layers[i] = manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(data)), Name: "config.json"} + return layers + } + } + t.Fatal("missing audio config fixture layer") + return nil +} + +func gemma4AudioFixtureLayer(t *testing.T, manifestName, internalName string, descriptor gemma4metadata.TensorDescriptor, raw []byte) manifest.Layer { + t.Helper() + data := raw + if data == nil { + shape := make([]int64, len(descriptor.Shape)) + for i, dim := range descriptor.Shape { + shape[i] = int64(dim) + } + size, err := gemma4SafetensorByteSize(descriptor.Dtype, shape) + if err != nil { + t.Fatal(err) + } + built, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{ + safetensors.NewTensorDataFromBytes(internalName, descriptor.Dtype, descriptor.Shape, make([]byte, int(size))), + })) + if err != nil { + t.Fatal(err) + } + data = built + } + digest := createTestBlob(t, data) + return manifest.Layer{MediaType: manifest.MediaTypeImageTensor, Digest: digest, Size: int64(len(data)), Name: manifestName} +} + +func replaceGemma4AudioFixtureLayer(t *testing.T, layers []manifest.Layer, manifestName, internalName string, descriptor gemma4metadata.TensorDescriptor, raw []byte) []manifest.Layer { + t.Helper() + for i := range layers { + if layers[i].Name == manifestName { + layers[i] = gemma4AudioFixtureLayer(t, manifestName, internalName, descriptor, raw) + return layers + } + } + t.Fatalf("missing audio fixture layer %s", manifestName) + return nil +} + func TestGemma4VisionTensorValidationRejectsIncompleteAndNearMatches(t *testing.T) { bareTensorNames := func() []string { names := gemma4VisionTensorNames(2) @@ -906,6 +1162,10 @@ func TestGemma4VisionTensorDescriptorsRejectsExcessiveInventory(t *testing.T) { if _, err := gemma4VisionTensorDescriptors(layers); err == nil || !strings.Contains(err.Error(), "descriptors") { t.Fatalf("excessive inventory error = %v", err) } + layers[0].Name = "model.audio_tower.synthetic.weight" + if _, err := gemma4AudioTensorDescriptors(layers); err == nil || !strings.Contains(err.Error(), "descriptors") { + t.Fatalf("excessive audio inventory error = %v", err) + } } func gemma4VisionConfig(layers int) *gemma4metadata.ConfigFile { @@ -915,6 +1175,42 @@ func gemma4VisionConfig(layers int) *gemma4metadata.ConfigFile { } } +func gemma4AudioConfig(layers int) *gemma4metadata.ConfigFile { + return &gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + AudioConfig: &gemma4metadata.AudioConfig{ + AttentionChunkSize: 2, AttentionContextLeft: 2, + ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, + NumHiddenLayers: layers, OutputProjDims: 3, + SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, + }, + } +} + +func gemma4AudioTensorNames(layers int) []string { + shapes, err := gemma4metadata.RequiredAudioTensorShapes(*gemma4AudioConfig(layers)) + if err != nil { + panic(err) + } + names := make([]string, 0, len(shapes)) + for name := range shapes { + names = append(names, name) + } + return names +} + +func testGemma4AudioTensorDescriptors(layers int) map[string]gemma4metadata.TensorDescriptor { + shapes, err := gemma4metadata.RequiredAudioTensorShapes(*gemma4AudioConfig(layers)) + if err != nil { + panic(err) + } + tensors := make(map[string]gemma4metadata.TensorDescriptor, len(shapes)) + for name, shape := range shapes { + tensors[name] = gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: shape} + } + return tensors +} + func testGemma4VisionTensorDescriptors(layers int) map[string]gemma4metadata.TensorDescriptor { return testGemma4VisionTensorDescriptorsForGeometry(layers, 4, 8, 6, 2, 16, 4) } diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 2c8d8655fdc..70c8d46b0ca 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -394,11 +394,17 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum if isLocalGemma4SafetensorsConfig(cfg) { var gemma4cfg gemma4metadata.ConfigFile - if err := mf.ReadConfigJSON("config.json", &gemma4cfg); err != nil || !hasGemma4VisionTensorLayers(gemma4cfg, mf.Layers) { + configErr := mf.ReadConfigJSON("config.json", &gemma4cfg) + if configErr != nil || !hasGemma4VisionTensorLayers(gemma4cfg, mf.Layers) { summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { return c == model.CapabilityVision }) } + if configErr != nil || !hasGemma4AudioTensorLayers(gemma4cfg, mf.Layers) { + summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { + return c == model.CapabilityAudio + }) + } } summary.Capabilities = filterUnsupportedModelListCapabilities(summary.Capabilities, cfg) @@ -406,7 +412,8 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum } func filterUnsupportedModelListCapabilities(capabilities []model.Capability, cfg model.ConfigV2) []model.Capability { - if cfg.ModelFormat == "safetensors" && isGemma4Renderer(cfg.Renderer) { + if cfg.ModelFormat == "safetensors" && isGemma4Renderer(cfg.Renderer) && + (cfg.RemoteHost != "" || cfg.RemoteModel != "") { capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { return c == model.CapabilityAudio }) diff --git a/x/create/client/create.go b/x/create/client/create.go index 67d07763e37..0cc93bacf00 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -562,13 +562,14 @@ func detectCapabilities(modelDir string) modelCapabilities { } vision := cfg.VisionConfig != nil || cfg.HasVision - if vision && isGemma4ModelConfig(cfg.Architectures, cfg.ModelType) { - vision = gemma4ModelDirHasVisionTensors(modelDir) + audio := cfg.AudioConfig != nil || cfg.SoundConfig != nil + if isGemma4ModelConfig(cfg.Architectures, cfg.ModelType) { + vision, audio = gemma4ModelDirMediaCapabilities(modelDir) } return modelCapabilities{ vision: vision, - audio: cfg.AudioConfig != nil || cfg.SoundConfig != nil, + audio: audio, thinking: chatTemplateHasThinkingSupport(readChatTemplate(modelDir)) || alwaysSupportsThinking(cfg.Architectures, cfg.ModelType), } @@ -626,14 +627,14 @@ func isGemma4ModelIdentifier(value string) bool { } } -func gemma4ModelDirHasVisionTensors(modelDir string) bool { +func gemma4ModelDirMediaCapabilities(modelDir string) (vision, audio bool) { inv, err := create.ReadInventory(modelDir) if err != nil { - return false + return false, false } var cfg gemma4metadata.ConfigFile if err := json.Unmarshal(inv.RawConfig, &cfg); err != nil { - return false + return false, false } tensors := make(map[string]gemma4metadata.TensorDescriptor, len(inv.Tensors)) for name, tensor := range inv.Tensors { @@ -641,7 +642,8 @@ func gemma4ModelDirHasVisionTensors(modelDir string) bool { // descriptor shapes rather than the released tower sentinel alone. tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} } - return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil + return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil, + gemma4metadata.ValidateAudioSourceInventory(cfg, tensors) == nil } // readChatTemplate returns the model's chat template, preferring the diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index 6b9f9fac2ac..bbcf9e164aa 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -435,7 +435,7 @@ func TestInferSafetensorsCapabilities(t *testing.T) { "vision_config": {"hidden_size": 1024}, "audio_config": {"num_mel_bins": 128} }`, - want: []string{"completion", "audio"}, + want: []string{"completion"}, }, { name: "model with audio but no vision", @@ -655,6 +655,131 @@ func TestInferSafetensorsCapabilitiesGemma4UnifiedVision(t *testing.T) { check(t, wrong, false) } +func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { + identities := []struct { + name string + architecture string + modelType string + }{ + {name: "released", architecture: "Gemma4ForConditionalGeneration", modelType: "gemma4"}, + {name: "unified", architecture: "Gemma4UnifiedForConditionalGeneration", modelType: "gemma4_unified"}, + } + for _, identity := range identities { + t.Run(identity.name, func(t *testing.T) { + cfg := gemma4metadata.ConfigFile{ + Architectures: []string{identity.architecture}, + ModelType: identity.modelType, + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + AudioConfig: &gemma4metadata.AudioConfig{ + AttentionChunkSize: 2, AttentionContextLeft: 2, + ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, + NumHiddenLayers: 1, OutputProjDims: 3, + SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, + }, + } + configJSON, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + shapes, err := gemma4metadata.RequiredAudioTensorShapes(cfg) + if err != nil { + t.Fatal(err) + } + valid := make(map[string]gemma4metadata.TensorDescriptor, len(shapes)) + for name, shape := range shapes { + valid[name] = gemma4metadata.TensorDescriptor{Dtype: "BF16", Shape: shape} + } + + check := func(t *testing.T, tensors map[string]gemma4metadata.TensorDescriptor, wantAudio bool) { + t.Helper() + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), configJSON, 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, tensors) + got := inferSafetensorsCapabilities(dir, "") + if slices.Contains(got, "audio") != wantAudio { + t.Fatalf("capabilities = %v, want audio %t", got, wantAudio) + } + } + + t.Run("complete", func(t *testing.T) { check(t, valid, true) }) + t.Run("partial", func(t *testing.T) { + partial := maps.Clone(valid) + delete(partial, "model.audio_tower.layers.0.self_attn.q_proj.input_max") + check(t, partial, false) + }) + t.Run("malformed", func(t *testing.T) { + malformed := maps.Clone(valid) + d := malformed["model.audio_tower.output_proj.weight"] + d.Shape = []int32{4, 3} + malformed["model.audio_tower.output_proj.weight"] = d + check(t, malformed, false) + }) + t.Run("near match", func(t *testing.T) { + near := maps.Clone(valid) + name := "model.embed_audio.embedding_projection.weight" + near[name+".extra"] = near[name] + delete(near, name) + check(t, near, false) + }) + }) + } +} + +func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testing.T) { + base := gemma4metadata.ConfigFile{ + Architectures: []string{"Gemma4ForConditionalGeneration"}, ModelType: "gemma4", + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + AudioConfig: &gemma4metadata.AudioConfig{ + AttentionChunkSize: 2, AttentionContextLeft: 2, + ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, + NumHiddenLayers: 1, OutputProjDims: 3, + SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, + }, + } + shapes, err := gemma4metadata.RequiredAudioTensorShapes(base) + if err != nil { + t.Fatal(err) + } + valid := make(map[string]gemma4metadata.TensorDescriptor, len(shapes)) + for name, shape := range shapes { + valid[name] = gemma4metadata.TensorDescriptor{Dtype: "BF16", Shape: shape} + } + tests := []struct { + name string + edit func(*gemma4metadata.ConfigFile) + }{ + {name: "shape product overflow", edit: func(cfg *gemma4metadata.ConfigFile) { + cfg.AudioConfig.HiddenSize = 1 << 30 + cfg.AudioConfig.NumAttentionHeads = 1 + }}, + {name: "impractical layer count", edit: func(cfg *gemma4metadata.ConfigFile) { cfg.AudioConfig.NumHiddenLayers = 1 << 30 }}, + {name: "impractical convolution channels", edit: func(cfg *gemma4metadata.ConfigFile) { cfg.AudioConfig.SubsamplingConvChannels = []int{2, 1 << 30} }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := base + audio := *base.AudioConfig + audio.SubsamplingConvChannels = slices.Clone(base.AudioConfig.SubsamplingConvChannels) + cfg.AudioConfig = &audio + tt.edit(&cfg) + configJSON, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), configJSON, 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, valid) + if got := inferSafetensorsCapabilities(dir, ""); slices.Contains(got, "audio") { + t.Fatalf("unbounded config capabilities = %v, did not expect audio", got) + } + }) + } +} + func TestInferSafetensorsCapabilitiesGemma4PackedSourceRequiresProducerContract(t *testing.T) { const configJSON = `{ "architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4", @@ -800,6 +925,9 @@ func TestGemma4ModelConfigRejectsNearMatches(t *testing.T) { if !isGemma4ModelConfig([]string{"Gemma4ForConditionalGeneration"}, "") { t.Fatal("released Gemma 4 architecture not classified") } + if !isGemma4ModelConfig([]string{"Gemma4UnifiedForConditionalGeneration"}, "") { + t.Fatal("unified Gemma 4 architecture not classified") + } } func gemma4ClientVisionTensorNames(layers int) []string { diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go new file mode 100644 index 00000000000..124d7558810 --- /dev/null +++ b/x/models/gemma4/metadata/audio.go @@ -0,0 +1,227 @@ +package metadata + +import ( + "fmt" + "math" + "slices" +) + +const gemma4AudioFeatureSize = 128 + +const ( + maxAudioHiddenSize = 8_192 + maxAudioLayers = 128 + maxAudioHeads = 256 + maxAudioOutputDims = 16_384 + maxAudioConvChannels = 4_096 + maxAudioKernelSize = 255 + maxAudioContextSize = 4_096 + maxTextHiddenSize = 65_536 +) + +type AudioConfig struct { + AttentionChunkSize int `json:"attention_chunk_size"` + AttentionContextLeft int `json:"attention_context_left"` + AttentionContextRight int `json:"attention_context_right"` + ConvKernelSize int `json:"conv_kernel_size"` + HiddenSize int `json:"hidden_size"` + NumAttentionHeads int `json:"num_attention_heads"` + NumHiddenLayers int `json:"num_hidden_layers"` + OutputProjDims int `json:"output_proj_dims"` + SubsamplingConvChannels []int `json:"subsampling_conv_channels"` + UseClippedLinears bool `json:"use_clipped_linears"` +} + +// ValidateAudioTensors verifies the normalized tensor names required by the +// released Gemma 4 MLX audio loader. +func ValidateAudioTensors(cfg ConfigFile, names []string) error { + if err := validateAudioConfig(cfg); err != nil { + return err + } + shapes, err := requiredAudioShapes(cfg) + if err != nil { + return err + } + present := make(map[string]struct{}, len(names)) + for _, name := range names { + present[name] = struct{}{} + } + for name := range shapes { + if _, ok := present[name]; !ok { + return fmt.Errorf("missing %s", name) + } + } + return nil +} + +// ValidateAudioSourceInventory additionally validates released tensor shapes +// and floating-point dtypes. +func ValidateAudioSourceInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + if err := validateAudioConfig(cfg); err != nil { + return err + } + shapes, err := requiredAudioShapes(cfg) + if err != nil { + return err + } + for name, shape := range shapes { + desc, ok := tensors[name] + if !ok { + return fmt.Errorf("missing %s", name) + } + if !slices.Equal(desc.Shape, shape) { + return fmt.Errorf("%s shape %v, want %v", name, desc.Shape, shape) + } + if !isFloat(desc.Dtype) { + return fmt.Errorf("%s dtype %s is not floating point", name, desc.Dtype) + } + } + return nil +} + +// ValidateAudioInstalledInventory validates the normalized descriptors stored +// in installed tensor layers. Installed audio remains source precision at this +// row, so its descriptor contract is identical to the released source form. +func ValidateAudioInstalledInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + return ValidateAudioSourceInventory(cfg, tensors) +} + +// RequiredAudioTensorShapes returns a copy of the normalized released audio +// tensor contract derived from config. +func RequiredAudioTensorShapes(cfg ConfigFile) (map[string][]int32, error) { + if err := validateAudioConfig(cfg); err != nil { + return nil, err + } + shapes, err := requiredAudioShapes(cfg) + if err != nil { + return nil, err + } + out := make(map[string][]int32, len(shapes)) + for name, shape := range shapes { + out[name] = slices.Clone(shape) + } + return out, nil +} + +func validateAudioConfig(cfg ConfigFile) error { + ac := cfg.AudioConfig + if ac == nil { + return fmt.Errorf("missing audio_config") + } + if ac.HiddenSize <= 0 || ac.HiddenSize > maxAudioHiddenSize || + ac.NumHiddenLayers <= 0 || ac.NumHiddenLayers > maxAudioLayers || + ac.NumAttentionHeads <= 0 || ac.NumAttentionHeads > maxAudioHeads || ac.HiddenSize%ac.NumAttentionHeads != 0 || + ac.OutputProjDims <= 0 || ac.OutputProjDims > maxAudioOutputDims || + ac.ConvKernelSize <= 0 || ac.ConvKernelSize > maxAudioKernelSize || ac.ConvKernelSize%2 == 0 || + ac.AttentionChunkSize <= 0 || ac.AttentionChunkSize > maxAudioContextSize || + ac.AttentionContextLeft <= 0 || ac.AttentionContextLeft > maxAudioContextSize || + ac.AttentionContextRight < 0 || ac.AttentionContextRight > maxAudioContextSize || + len(ac.SubsamplingConvChannels) != 2 || + ac.SubsamplingConvChannels[0] <= 0 || ac.SubsamplingConvChannels[0] > maxAudioConvChannels || + ac.SubsamplingConvChannels[1] <= 0 || ac.SubsamplingConvChannels[1] > maxAudioConvChannels || + cfg.TextConfig.HiddenSize <= 0 || cfg.TextConfig.HiddenSize > maxTextHiddenSize { + return fmt.Errorf("invalid Gemma 4 audio dimensions") + } + return nil +} + +func requiredAudioShapes(cfg ConfigFile) (map[string][]int32, error) { + ac := cfg.AudioConfig + hidden, err := checkedAudioShapeDim("hidden_size", int64(ac.HiddenSize)) + if err != nil { + return nil, err + } + output, err := checkedAudioShapeDim("output_proj_dims", int64(ac.OutputProjDims)) + if err != nil { + return nil, err + } + headDim, err := checkedAudioShapeDim("attention head dimension", int64(ac.HiddenSize)/int64(ac.NumAttentionHeads)) + if err != nil { + return nil, err + } + c0, err := checkedAudioShapeDim("subsampling channel 0", int64(ac.SubsamplingConvChannels[0])) + if err != nil { + return nil, err + } + c1, err := checkedAudioShapeDim("subsampling channel 1", int64(ac.SubsamplingConvChannels[1])) + if err != nil { + return nil, err + } + freq0 := int32((gemma4AudioFeatureSize + 1) / 2) + freq1 := (freq0 + 1) / 2 + inputWidth, err := checkedAudioShapeDim("subsampling projection width", int64(freq1), int64(c1)) + if err != nil { + return nil, err + } + ffWidth, err := checkedAudioShapeDim("feed-forward width", 4, int64(hidden)) + if err != nil { + return nil, err + } + convWidth, err := checkedAudioShapeDim("convolution input width", 2, int64(hidden)) + if err != nil { + return nil, err + } + textHidden, err := checkedAudioShapeDim("text hidden_size", int64(cfg.TextConfig.HiddenSize)) + if err != nil { + return nil, err + } + kernel, err := checkedAudioShapeDim("conv_kernel_size", int64(ac.ConvKernelSize)) + if err != nil { + return nil, err + } + + required := map[string][]int32{ + "model.audio_tower.subsample_conv_projection.layer0.conv.weight": {c0, 1, 3, 3}, + "model.audio_tower.subsample_conv_projection.layer0.norm.weight": {c0}, + "model.audio_tower.subsample_conv_projection.layer1.conv.weight": {c1, c0, 3, 3}, + "model.audio_tower.subsample_conv_projection.layer1.norm.weight": {c1}, + "model.audio_tower.subsample_conv_projection.input_proj_linear.weight": {hidden, inputWidth}, + "model.audio_tower.output_proj.weight": {output, hidden}, + "model.audio_tower.output_proj.bias": {output}, + "model.embed_audio.embedding_projection.weight": {textHidden, output}, + } + + addLinear := func(path string, shape []int32) { + required[path+".linear.weight"] = shape + if ac.UseClippedLinears { + for _, suffix := range []string{".input_min", ".input_max", ".output_min", ".output_max"} { + required[path+suffix] = []int32{} + } + } + } + + for i := range ac.NumHiddenLayers { + layer := fmt.Sprintf("model.audio_tower.layers.%d", i) + for _, ff := range []string{"feed_forward1", "feed_forward2"} { + required[layer+"."+ff+".pre_layer_norm.weight"] = []int32{hidden} + required[layer+"."+ff+".post_layer_norm.weight"] = []int32{hidden} + addLinear(layer+"."+ff+".ffw_layer_1", []int32{ffWidth, hidden}) + addLinear(layer+"."+ff+".ffw_layer_2", []int32{hidden, ffWidth}) + } + + required[layer+".norm_pre_attn.weight"] = []int32{hidden} + required[layer+".norm_post_attn.weight"] = []int32{hidden} + required[layer+".norm_out.weight"] = []int32{hidden} + for _, projection := range []string{"q_proj", "k_proj", "v_proj", "post"} { + addLinear(layer+".self_attn."+projection, []int32{hidden, hidden}) + } + required[layer+".self_attn.per_dim_scale"] = []int32{headDim} + required[layer+".self_attn.relative_k_proj.weight"] = []int32{hidden, hidden} + + required[layer+".lconv1d.pre_layer_norm.weight"] = []int32{hidden} + required[layer+".lconv1d.conv_norm.weight"] = []int32{hidden} + required[layer+".lconv1d.depthwise_conv1d.weight"] = []int32{hidden, 1, kernel} + addLinear(layer+".lconv1d.linear_start", []int32{convWidth, hidden}) + addLinear(layer+".lconv1d.linear_end", []int32{hidden, hidden}) + } + + return required, nil +} + +func checkedAudioShapeDim(name string, factors ...int64) (int32, error) { + value, ok := checkedProduct(math.MaxInt32, factors...) + if !ok { + return 0, fmt.Errorf("invalid Gemma 4 audio %s", name) + } + return int32(value), nil +} diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go new file mode 100644 index 00000000000..4bc79ba46ec --- /dev/null +++ b/x/models/gemma4/metadata/audio_test.go @@ -0,0 +1,167 @@ +package metadata + +import ( + "maps" + "slices" + "strings" + "testing" +) + +func releasedAudioConfig(layers int) ConfigFile { + return ConfigFile{ + TextConfig: TextConfig{HiddenSize: 2560}, + AudioConfig: &AudioConfig{ + AttentionChunkSize: 12, AttentionContextLeft: 13, + ConvKernelSize: 5, HiddenSize: 1024, NumAttentionHeads: 8, + NumHiddenLayers: layers, OutputProjDims: 1536, + SubsamplingConvChannels: []int{128, 32}, UseClippedLinears: true, + }, + } +} + +func completeAudioInventory(cfg ConfigFile) map[string]TensorDescriptor { + tensors := make(map[string]TensorDescriptor) + shapes, err := requiredAudioShapes(cfg) + if err != nil { + panic(err) + } + for name, shape := range shapes { + tensors[name] = TensorDescriptor{Dtype: "BF16", Shape: shape} + } + return tensors +} + +func TestValidateReleasedAudioInventory(t *testing.T) { + cfg := releasedAudioConfig(12) + tensors := completeAudioInventory(cfg) + if got := len(tensors); got != 752 { + t.Fatalf("audio tensor count = %d, want 752", got) + } + if err := ValidateAudioSourceInventory(cfg, tensors); err != nil { + t.Fatalf("ValidateAudioSourceInventory() error = %v", err) + } + + partial := maps.Clone(tensors) + delete(partial, "model.audio_tower.layers.11.self_attn.q_proj.input_max") + if err := ValidateAudioSourceInventory(cfg, partial); err == nil { + t.Fatal("partial clipping inventory: error = nil") + } + + nearMatch := maps.Clone(partial) + nearMatch["model.audio_tower.layers.11.self_attn.q_proj.input_max.extra"] = TensorDescriptor{Dtype: "BF16", Shape: []int32{}} + if err := ValidateAudioSourceInventory(cfg, nearMatch); err == nil { + t.Fatal("near-match clipping scalar substituted for exact name: error = nil") + } + + badShape := maps.Clone(tensors) + badShape["model.audio_tower.output_proj.weight"] = TensorDescriptor{Dtype: "BF16", Shape: []int32{1024, 1536}} + if err := ValidateAudioSourceInventory(cfg, badShape); err == nil { + t.Fatal("transposed output projection: error = nil") + } + + badDtype := maps.Clone(tensors) + badDtype["model.embed_audio.embedding_projection.weight"] = TensorDescriptor{Dtype: "U8", Shape: []int32{2560, 1536}} + if err := ValidateAudioSourceInventory(cfg, badDtype); err == nil { + t.Fatal("non-floating projector: error = nil") + } +} + +func TestValidateAudioTensorsExactNormalizedNames(t *testing.T) { + cfg := releasedAudioConfig(1) + shapes, err := RequiredAudioTensorShapes(cfg) + if err != nil { + t.Fatal(err) + } + names := make([]string, 0, len(shapes)) + for name := range shapes { + names = append(names, name) + } + if err := ValidateAudioTensors(cfg, names); err != nil { + t.Fatalf("complete normalized inventory: %v", err) + } + + missing := "model.audio_tower.subsample_conv_projection.layer0.conv.weight" + for i, name := range names { + if name == missing { + names[i] = "prefix." + name + break + } + } + if err := ValidateAudioTensors(cfg, names); err == nil || !strings.Contains(err.Error(), missing) { + t.Fatalf("near-match inventory error = %v, want exact missing name", err) + } +} + +func TestValidateAudioConfig(t *testing.T) { + tests := []struct { + name string + edit func(*ConfigFile) + }{ + {"missing", func(cfg *ConfigFile) { cfg.AudioConfig = nil }}, + {"heads", func(cfg *ConfigFile) { cfg.AudioConfig.NumAttentionHeads = 7 }}, + {"conv channels", func(cfg *ConfigFile) { cfg.AudioConfig.SubsamplingConvChannels = []int{128} }}, + {"even kernel", func(cfg *ConfigFile) { cfg.AudioConfig.ConvKernelSize = 4 }}, + {"text hidden", func(cfg *ConfigFile) { cfg.TextConfig.HiddenSize = 0 }}, + {"hidden overflow", func(cfg *ConfigFile) { cfg.AudioConfig.HiddenSize = 1 << 30; cfg.AudioConfig.NumAttentionHeads = 1 }}, + {"impractical layers", func(cfg *ConfigFile) { cfg.AudioConfig.NumHiddenLayers = maxAudioLayers + 1 }}, + {"heads bound", func(cfg *ConfigFile) { cfg.AudioConfig.NumAttentionHeads = maxAudioHeads + 1 }}, + {"output bound", func(cfg *ConfigFile) { cfg.AudioConfig.OutputProjDims = maxAudioOutputDims + 1 }}, + {"kernel bound", func(cfg *ConfigFile) { cfg.AudioConfig.ConvKernelSize = maxAudioKernelSize + 2 }}, + {"chunk bound", func(cfg *ConfigFile) { cfg.AudioConfig.AttentionChunkSize = maxAudioContextSize + 1 }}, + {"left context bound", func(cfg *ConfigFile) { cfg.AudioConfig.AttentionContextLeft = maxAudioContextSize + 1 }}, + {"right context bound", func(cfg *ConfigFile) { cfg.AudioConfig.AttentionContextRight = maxAudioContextSize + 1 }}, + {"conv channel bound", func(cfg *ConfigFile) { cfg.AudioConfig.SubsamplingConvChannels[1] = maxAudioConvChannels + 1 }}, + {"text hidden bound", func(cfg *ConfigFile) { cfg.TextConfig.HiddenSize = maxTextHiddenSize + 1 }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := releasedAudioConfig(12) + tt.edit(&cfg) + if err := ValidateAudioTensors(cfg, nil); err == nil { + t.Fatal("ValidateAudioTensors() error = nil") + } + }) + } +} + +func TestAudioConfigSupportedBoundaries(t *testing.T) { + cfg := releasedAudioConfig(maxAudioLayers) + cfg.TextConfig.HiddenSize = maxTextHiddenSize + cfg.AudioConfig.HiddenSize = maxAudioHiddenSize + cfg.AudioConfig.NumAttentionHeads = maxAudioHeads + cfg.AudioConfig.OutputProjDims = maxAudioOutputDims + cfg.AudioConfig.ConvKernelSize = maxAudioKernelSize + cfg.AudioConfig.AttentionChunkSize = maxAudioContextSize + cfg.AudioConfig.AttentionContextLeft = maxAudioContextSize + cfg.AudioConfig.AttentionContextRight = maxAudioContextSize + cfg.AudioConfig.SubsamplingConvChannels = []int{maxAudioConvChannels, maxAudioConvChannels} + shapes, err := RequiredAudioTensorShapes(cfg) + if err != nil { + t.Fatalf("supported boundary config: %v", err) + } + if got, want := len(shapes), 8+maxAudioLayers*62; got != want { + t.Fatalf("boundary inventory count = %d, want %d", got, want) + } + names := make([]string, 0, len(shapes)) + for name := range shapes { + names = append(names, name) + } + if err := ValidateAudioTensors(cfg, names); err != nil { + t.Fatalf("supported boundary inventory: %v", err) + } + if got, want := shapes["model.audio_tower.layers.127.feed_forward1.ffw_layer_1.linear.weight"], []int32{4 * maxAudioHiddenSize, maxAudioHiddenSize}; !slices.Equal(got, want) { + t.Fatalf("boundary feed-forward shape = %v, want %v", got, want) + } +} + +func TestAudioWithoutClippingScalars(t *testing.T) { + cfg := releasedAudioConfig(1) + cfg.AudioConfig.UseClippedLinears = false + tensors := completeAudioInventory(cfg) + if got := len(tensors); got != 30 { + t.Fatalf("audio tensor count = %d, want 30", got) + } + if err := ValidateAudioSourceInventory(cfg, tensors); err != nil { + t.Fatalf("ValidateAudioSourceInventory() error = %v", err) + } +} diff --git a/x/models/gemma4/metadata/vision.go b/x/models/gemma4/metadata/vision.go index 75028b9807a..0d26583b74c 100644 --- a/x/models/gemma4/metadata/vision.go +++ b/x/models/gemma4/metadata/vision.go @@ -28,6 +28,7 @@ type ConfigFile struct { ModelType string `json:"model_type"` TextConfig TextConfig `json:"text_config"` VisionConfig *VisionConfig `json:"vision_config"` + AudioConfig *AudioConfig `json:"audio_config"` Quantization Quantization `json:"quantization"` QuantizationConfig Quantization `json:"quantization_config"` } From 1784d0adb765fae8917079fc829e17a065af42eb Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 00:09:56 +0000 Subject: [PATCH 47/58] gemma4: add released MLX audio processor Add the bounded waveform-to-feature preprocessing used by released Gemma 4 audio checkpoints, with deterministic reference coverage. Co-authored-by: Codex --- x/models/gemma4/audio_processor.go | 477 ++++++++++++++++++++++++ x/models/gemma4/audio_processor_test.go | 467 +++++++++++++++++++++++ 2 files changed, 944 insertions(+) create mode 100644 x/models/gemma4/audio_processor.go create mode 100644 x/models/gemma4/audio_processor_test.go diff --git a/x/models/gemma4/audio_processor.go b/x/models/gemma4/audio_processor.go new file mode 100644 index 00000000000..17fe45276e5 --- /dev/null +++ b/x/models/gemma4/audio_processor.go @@ -0,0 +1,477 @@ +package gemma4 + +import ( + "context" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "math" + "math/cmplx" + + "gonum.org/v1/gonum/dsp/fourier" +) + +const ( + maxGemma4AudioBytes = 32 << 20 + maxGemma4AudioChannels = 8 + minGemma4AudioSampleRate = 8_000 + maxGemma4AudioSampleRate = 192_000 + maxGemma4AudioSamples = 480_000 + maxGemma4AudioFrames = 3_000 + maxGemma4AudioFeatures = maxGemma4AudioFrames * 128 +) + +type AudioProcessorConfig struct { + AudioSequenceLength int `json:"audio_seq_length"` + FeatureExtractor struct { + Dither float64 `json:"dither"` + FeatureSize int `json:"feature_size"` + FFTLength int `json:"fft_length"` + FFTOverdrive bool `json:"fft_overdrive"` + FrameLength int `json:"frame_length"` + HopLength int `json:"hop_length"` + InputScaleFactor float64 `json:"input_scale_factor"` + MaxFrequency float64 `json:"max_frequency"` + MelFloor float64 `json:"mel_floor"` + MinFrequency float64 `json:"min_frequency"` + PaddingSide string `json:"padding_side"` + PerBinMean []float64 `json:"per_bin_mean"` + PerBinStddev []float64 `json:"per_bin_stddev"` + Preemphasis float64 `json:"preemphasis"` + SamplingRate int `json:"sampling_rate"` + } `json:"feature_extractor"` +} + +type gemma4AudioInput struct { + Features []float32 + FeatureMask []bool + Frames int + SoftTokens int +} + +func defaultAudioProcessorConfig() AudioProcessorConfig { + var cfg AudioProcessorConfig + cfg.AudioSequenceLength = 750 + cfg.FeatureExtractor.FeatureSize = 128 + cfg.FeatureExtractor.FFTLength = 512 + cfg.FeatureExtractor.FrameLength = 320 + cfg.FeatureExtractor.HopLength = 160 + cfg.FeatureExtractor.InputScaleFactor = 1 + cfg.FeatureExtractor.MaxFrequency = 8000 + cfg.FeatureExtractor.MelFloor = 1e-3 + cfg.FeatureExtractor.PaddingSide = "right" + cfg.FeatureExtractor.SamplingRate = 16000 + return cfg +} + +func parseAudioProcessorConfig(data []byte) (*AudioProcessorConfig, error) { + cfg := defaultAudioProcessorConfig() + if len(data) > 0 { + if err := json.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("parse Gemma4 audio processor config: %w", err) + } + } + if err := validateReleasedAudioProcessorConfig(&cfg); err != nil { + return nil, err + } + return &cfg, nil +} + +func validateReleasedAudioProcessorConfig(cfg *AudioProcessorConfig) error { + if cfg == nil { + return errors.New("Gemma4 MLX model has no supported audio processor configuration") + } + f := cfg.FeatureExtractor + if cfg.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || + f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || + f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || + f.MelFloor != 1e-3 || f.Preemphasis != 0 || f.PaddingSide != "right" || + len(f.PerBinMean) != 0 || len(f.PerBinStddev) != 0 { + return errors.New("unsupported Gemma4 audio processor configuration") + } + return nil +} + +func preprocessGemma4Audio(ctx context.Context, data []byte, cfg *AudioProcessorConfig) (*gemma4AudioInput, error) { + if err := validateReleasedAudioProcessorConfig(cfg); err != nil { + return nil, err + } + samples, err := decodeGemma4WAV(ctx, data, cfg.FeatureExtractor.SamplingRate) + if err != nil { + return nil, err + } + features, featureMask, err := computeGemma4LogMel(ctx, samples, cfg) + if err != nil { + return nil, err + } + outputMask := downsampleGemma4AudioMask(featureMask, 2) + softTokens := 0 + seenPadding := false + for _, valid := range outputMask { + if valid { + if seenPadding { + return nil, errors.New("Gemma4 audio validity mask is not a prefix") + } + softTokens++ + } else { + seenPadding = true + } + } + if softTokens == 0 { + return nil, errors.New("Gemma4 audio is too short to encode") + } + if softTokens > cfg.AudioSequenceLength { + return nil, fmt.Errorf("Gemma4 audio token count %d exceeds limit %d", softTokens, cfg.AudioSequenceLength) + } + return &gemma4AudioInput{ + Features: features, FeatureMask: featureMask, Frames: len(featureMask), SoftTokens: softTokens, + }, nil +} + +func decodeGemma4WAV(ctx context.Context, data []byte, targetRate int) ([]float32, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if len(data) > maxGemma4AudioBytes { + return nil, fmt.Errorf("Gemma4 audio is %d bytes, limit %d", len(data), maxGemma4AudioBytes) + } + if targetRate != 16000 { + return nil, fmt.Errorf("unsupported Gemma4 audio target sample rate %d", targetRate) + } + if len(data) < 12 || string(data[:4]) != "RIFF" || string(data[8:12]) != "WAVE" { + return nil, errors.New("Gemma4 audio must be a RIFF/WAVE file") + } + + var format, channels, bits, blockAlign uint16 + var sampleRate uint32 + var pcm []byte + seenFormat := false + for offset := uint64(12); offset+8 <= uint64(len(data)); { + if err := ctx.Err(); err != nil { + return nil, err + } + start := int(offset) + size := uint64(binary.LittleEndian.Uint32(data[start+4 : start+8])) + chunkStart := offset + 8 + chunkEnd := chunkStart + size + if chunkEnd < chunkStart || chunkEnd > uint64(len(data)) { + return nil, fmt.Errorf("truncated WAV chunk %q", string(data[start:start+4])) + } + chunk := data[int(chunkStart):int(chunkEnd)] + switch string(data[start : start+4]) { + case "fmt ": + if seenFormat { + return nil, errors.New("WAV file contains duplicate fmt chunks") + } + seenFormat = true + if len(chunk) < 16 { + return nil, errors.New("WAV fmt chunk is too short") + } + format = binary.LittleEndian.Uint16(chunk[0:2]) + channels = binary.LittleEndian.Uint16(chunk[2:4]) + sampleRate = binary.LittleEndian.Uint32(chunk[4:8]) + blockAlign = binary.LittleEndian.Uint16(chunk[12:14]) + bits = binary.LittleEndian.Uint16(chunk[14:16]) + if format == 0xfffe { + if len(chunk) < 40 { + return nil, errors.New("WAV extensible fmt chunk is too short") + } + cbSize := binary.LittleEndian.Uint16(chunk[16:18]) + if cbSize != 22 || len(chunk) != int(cbSize)+18 { + return nil, errors.New("invalid WAV extensible fmt size") + } + validBits := binary.LittleEndian.Uint16(chunk[18:20]) + if validBits != bits { + return nil, fmt.Errorf("unsupported WAV valid bits %d for %d-bit samples", validBits, bits) + } + pcmGUID := [16]byte{1, 0, 0, 0, 0, 0, 0x10, 0, 0x80, 0, 0, 0xaa, 0, 0x38, 0x9b, 0x71} + floatGUID := pcmGUID + floatGUID[0] = 3 + var subformat [16]byte + copy(subformat[:], chunk[24:40]) + switch subformat { + case pcmGUID: + format = 1 + case floatGUID: + format = 3 + default: + return nil, errors.New("unsupported WAV extensible subformat") + } + } + case "data": + if pcm == nil { + pcm = chunk + } + } + offset = chunkEnd + size%2 + } + + if format == 0 || pcm == nil { + return nil, errors.New("WAV file is missing fmt or data chunk") + } + if channels == 0 || channels > maxGemma4AudioChannels { + return nil, fmt.Errorf("unsupported WAV channel count %d", channels) + } + if sampleRate < minGemma4AudioSampleRate || sampleRate > maxGemma4AudioSampleRate { + return nil, fmt.Errorf("unsupported WAV sample rate %d", sampleRate) + } + validEncoding := format == 1 && (bits == 8 || bits == 16 || bits == 24 || bits == 32) || format == 3 && bits == 32 + if !validEncoding { + return nil, fmt.Errorf("unsupported WAV encoding format=%d bits=%d", format, bits) + } + bytesPerSample := int(bits / 8) + wantBlockAlign := int(channels) * bytesPerSample + if int(blockAlign) != wantBlockAlign || len(pcm)%wantBlockAlign != 0 { + return nil, errors.New("invalid WAV block alignment") + } + + frames := len(pcm) / wantBlockAlign + maxSourceFrames := int64(maxGemma4AudioSamples) * int64(sampleRate) / int64(targetRate) + if int64(frames) > maxSourceFrames { + frames = int(maxSourceFrames) + } + if err := ctx.Err(); err != nil { + return nil, err + } + samples := make([]float32, frames) + for i := range frames { + if i&4095 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + var sum float64 + for channel := range int(channels) { + offset := (i*int(channels) + channel) * bytesPerSample + switch { + case format == 1 && bits == 8: + sum += (float64(pcm[offset]) - 128) / 128 + case format == 1 && bits == 16: + sum += float64(int16(binary.LittleEndian.Uint16(pcm[offset:offset+2]))) / 32768 + case format == 1 && bits == 24: + value := int32(pcm[offset]) | int32(pcm[offset+1])<<8 | int32(pcm[offset+2])<<16 + if value&0x800000 != 0 { + value |= ^int32(0xffffff) + } + sum += float64(value) / 8388608 + case format == 1 && bits == 32: + sum += float64(int32(binary.LittleEndian.Uint32(pcm[offset:offset+4]))) / 2147483648 + case format == 3 && bits == 32: + value := math.Float32frombits(binary.LittleEndian.Uint32(pcm[offset : offset+4])) + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return nil, errors.New("WAV contains a non-finite float sample") + } + sum += float64(value) + } + } + samples[i] = float32(sum / float64(channels)) + } + if int(sampleRate) != targetRate { + resampled, err := resampleGemma4Audio(ctx, samples, int(sampleRate), targetRate) + if err != nil { + return nil, err + } + samples = resampled + } + if len(samples) > maxGemma4AudioSamples { + samples = samples[:maxGemma4AudioSamples] + } + return samples, nil +} + +func resampleGemma4Audio(ctx context.Context, samples []float32, sourceRate, targetRate int) ([]float32, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if sourceRate < minGemma4AudioSampleRate || sourceRate > maxGemma4AudioSampleRate || targetRate != 16000 { + return nil, fmt.Errorf("unsupported Gemma4 audio resampling rate %d to %d", sourceRate, targetRate) + } + if sourceRate == targetRate { + if len(samples) > maxGemma4AudioSamples { + return samples[:maxGemma4AudioSamples], nil + } + return samples, nil + } + if len(samples) == 0 { + return nil, nil + } + maxSourceFrames := int64(maxGemma4AudioSamples) * int64(sourceRate) / int64(targetRate) + if int64(len(samples)) > maxSourceFrames { + samples = samples[:maxSourceFrames] + } + length64 := int64(len(samples)) * int64(targetRate) / int64(sourceRate) + if length64 > maxGemma4AudioSamples { + length64 = maxGemma4AudioSamples + } + if length64 <= 0 { + return nil, nil + } + if err := ctx.Err(); err != nil { + return nil, err + } + out := make([]float32, int(length64)) + for i := range out { + if i&4095 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } + position := float64(i) * float64(sourceRate) / float64(targetRate) + left := int(position) + if left >= len(samples)-1 { + out[i] = samples[len(samples)-1] + continue + } + fraction := float32(position - float64(left)) + out[i] = samples[left]*(1-fraction) + samples[left+1]*fraction + } + return out, nil +} + +func computeGemma4LogMel(ctx context.Context, samples []float32, cfg *AudioProcessorConfig) ([]float32, []bool, error) { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + if err := validateReleasedAudioProcessorConfig(cfg); err != nil { + return nil, nil, err + } + if len(samples) > maxGemma4AudioSamples { + return nil, nil, fmt.Errorf("Gemma4 audio has %d samples, limit %d", len(samples), maxGemma4AudioSamples) + } + f := cfg.FeatureExtractor + validSamples := len(samples) + paddedLength := validSamples + if remainder := paddedLength % 128; remainder != 0 { + var ok bool + paddedLength, ok = checkedGemma4AudioAdd(paddedLength, 128-remainder) + if !ok { + return nil, nil, errors.New("Gemma4 audio padded sample count overflows") + } + } + leftPadding := f.FrameLength / 2 + unfoldSize, ok := checkedGemma4AudioAdd(f.FrameLength, 1) + if !ok { + return nil, nil, errors.New("Gemma4 audio frame size overflows") + } + paddedWithLeft, ok := checkedGemma4AudioAdd(paddedLength, leftPadding) + if !ok { + return nil, nil, errors.New("Gemma4 audio padded frame count overflows") + } + if paddedWithLeft < unfoldSize { + return nil, nil, errors.New("Gemma4 audio is too short to frame") + } + frames := (paddedWithLeft-unfoldSize)/f.HopLength + 1 + if frames <= 0 || frames > maxGemma4AudioFrames { + return nil, nil, fmt.Errorf("Gemma4 audio frame count %d exceeds limit %d", frames, maxGemma4AudioFrames) + } + featureValues, ok := checkedGemma4AudioMul(frames, f.FeatureSize) + if !ok || featureValues > maxGemma4AudioFeatures { + return nil, nil, errors.New("Gemma4 audio feature dimensions exceed limit") + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + features := make([]float32, featureValues) + mask := make([]bool, frames) + if err := ctx.Err(); err != nil { + return nil, nil, err + } + window := make([]float32, f.FrameLength) + for i := range window { + window[i] = float32(0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(f.FrameLength))) + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + melFilters := gemma4MelFilterBank(f.FFTLength/2+1, f.FeatureSize, f.MinFrequency, f.MaxFrequency, f.SamplingRate) + if err := ctx.Err(); err != nil { + return nil, nil, err + } + fft := fourier.NewFFT(f.FFTLength) + sequence := make([]float64, f.FFTLength) + coefficients := make([]complex128, f.FFTLength/2+1) + magnitudes := make([]float64, len(coefficients)) + + for frame := range frames { + if frame&31 == 0 { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + } + clear(sequence) + start := frame*f.HopLength - leftPadding + for i := range f.FrameLength { + index := start + i + if index >= 0 && index < validSamples { + sequence[i] = float64(samples[index] * window[i]) + } + } + coefficients = fft.Coefficients(coefficients, sequence) + for i, value := range coefficients { + magnitudes[i] = cmplx.Abs(value) + } + valid := frame*f.HopLength+f.FrameLength < leftPadding+validSamples + mask[frame] = valid + if !valid { + continue + } + for mel := range f.FeatureSize { + value := 0.0 + for bin, magnitude := range magnitudes { + value += magnitude * melFilters[bin*f.FeatureSize+mel] + } + features[frame*f.FeatureSize+mel] = float32(math.Log(value + f.MelFloor)) + } + } + return features, mask, nil +} + +func gemma4MelFilterBank(frequencyBins, melBins int, minFrequency, maxFrequency float64, sampleRate int) []float64 { + hzToMel := func(value float64) float64 { return 2595 * math.Log10(1+value/700) } + melToHz := func(value float64) float64 { return 700 * (math.Pow(10, value/2595) - 1) } + melMin, melMax := hzToMel(minFrequency), hzToMel(maxFrequency) + centers := make([]float64, melBins+2) + for i := range centers { + centers[i] = melToHz(melMin + float64(i)*(melMax-melMin)/float64(melBins+1)) + } + filters := make([]float64, frequencyBins*melBins) + for bin := range frequencyBins { + frequency := float64(bin) * float64(sampleRate/2) / float64(frequencyBins-1) + for mel := range melBins { + down := (frequency - centers[mel]) / (centers[mel+1] - centers[mel]) + up := (centers[mel+2] - frequency) / (centers[mel+2] - centers[mel+1]) + filters[bin*melBins+mel] = math.Max(0, math.Min(down, up)) + } + } + return filters +} + +func checkedGemma4AudioAdd(a, b int) (int, bool) { + if a < 0 || b < 0 || a > int(^uint(0)>>1)-b { + return 0, false + } + return a + b, true +} + +func checkedGemma4AudioMul(a, b int) (int, bool) { + if a < 0 || b < 0 || a != 0 && b > int(^uint(0)>>1)/a { + return 0, false + } + return a * b, true +} + +func downsampleGemma4AudioMask(mask []bool, layers int) []bool { + if len(mask) == 0 { + return nil + } + for range layers { + outputLength := (len(mask)-1)/2 + 1 + output := make([]bool, outputLength) + for i := range output { + output[i] = mask[i*2] + } + mask = output + } + return mask +} diff --git a/x/models/gemma4/audio_processor_test.go b/x/models/gemma4/audio_processor_test.go new file mode 100644 index 00000000000..2b7aca358db --- /dev/null +++ b/x/models/gemma4/audio_processor_test.go @@ -0,0 +1,467 @@ +package gemma4 + +import ( + "bytes" + "context" + "encoding/binary" + "fmt" + "math" + "strings" + "testing" + "time" +) + +func TestParseReleasedAudioProcessorConfig(t *testing.T) { + data := []byte(`{ + "audio_seq_length":750, + "feature_extractor":{ + "dither":0.0,"feature_size":128,"fft_length":512,"fft_overdrive":false, + "frame_length":320,"hop_length":160,"input_scale_factor":1.0, + "max_frequency":8000.0,"mel_floor":0.001,"min_frequency":0.0, + "padding_side":"right","per_bin_mean":null,"per_bin_stddev":null, + "preemphasis":0.0,"sampling_rate":16000 + } + }`) + cfg, err := parseAudioProcessorConfig(data) + if err != nil { + t.Fatal(err) + } + if cfg.FeatureExtractor.FFTLength != 512 || cfg.AudioSequenceLength != 750 { + t.Fatalf("processor config = %#v", cfg) + } + + bad := bytes.Replace(data, []byte(`"fft_length":512`), []byte(`"fft_length":1024`), 1) + if _, err := parseAudioProcessorConfig(bad); err == nil { + t.Fatal("1024-point FFT processor: error = nil") + } +} + +func TestReleasedAudioProcessorRejectsMalformedDimensions(t *testing.T) { + base := defaultAudioProcessorConfig() + tests := []struct { + name string + mutate func(*AudioProcessorConfig) + }{ + {"sample rate", func(cfg *AudioProcessorConfig) { cfg.FeatureExtractor.SamplingRate = 0 }}, + {"frame length", func(cfg *AudioProcessorConfig) { cfg.FeatureExtractor.FrameLength = -1 }}, + {"hop length", func(cfg *AudioProcessorConfig) { cfg.FeatureExtractor.HopLength = 0 }}, + {"feature size", func(cfg *AudioProcessorConfig) { cfg.FeatureExtractor.FeatureSize = int(^uint(0) >> 1) }}, + {"FFT length", func(cfg *AudioProcessorConfig) { cfg.FeatureExtractor.FFTLength = 0 }}, + {"sequence length", func(cfg *AudioProcessorConfig) { cfg.AudioSequenceLength = int(^uint(0) >> 1) }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := base + tt.mutate(&cfg) + if _, _, err := computeGemma4LogMel(context.Background(), make([]float32, 400), &cfg); err == nil || !strings.Contains(err.Error(), "configuration") { + t.Fatalf("malformed configuration error = %v", err) + } + }) + } + + if _, _, err := computeGemma4LogMel(context.Background(), make([]float32, maxGemma4AudioSamples+1), &base); err == nil || !strings.Contains(err.Error(), "samples") { + t.Fatalf("oversized sample input error = %v", err) + } + maxInt := int(^uint(0) >> 1) + if _, ok := checkedGemma4AudioAdd(maxInt, 1); ok { + t.Fatal("checked add accepted overflow") + } + if _, ok := checkedGemma4AudioMul(maxInt, 2); ok { + t.Fatal("checked multiply accepted overflow") + } +} + +func TestGemma4LogMelReference(t *testing.T) { + cfg := defaultAudioProcessorConfig() + samples := make([]float32, 4000) + for i := range samples { + samples[i] = float32(0.25 * math.Sin(2*math.Pi*440*float64(i)/16000)) + } + features, mask, err := computeGemma4LogMel(context.Background(), samples, &cfg) + if err != nil { + t.Fatal(err) + } + if len(mask) != 25 || len(features) != 25*128 { + t.Fatalf("feature shape = (%d, %d), want (25, 128)", len(mask), len(features)/len(mask)) + } + valid := 0 + for _, value := range mask { + if value { + valid++ + } + } + if valid != 24 { + t.Fatalf("valid frames = %d, want 24", valid) + } + selected := map[int]float64{ + 0*128 + 0: -6.907755374908447, + 0*128 + 1: 0.11585413664579391, + 0*128 + 10: -0.6425663232803345, + 0*128 + 64: -1.9309545755386353, + 0*128 + 127: -2.837610960006714, + 1*128 + 10: -4.558506965637207, + 12*128 + 10: -4.529771327972412, + 23*128 + 10: -4.529770374298096, + 24*128 + 10: 0, + } + for index, want := range selected { + if got := float64(features[index]); math.Abs(got-want) > 1e-5 { + t.Errorf("features[%d] = %.9f, want %.9f", index, got, want) + } + } + sum := 0.0 + for _, value := range features { + sum += float64(value) + } + if math.Abs(sum-(-15998.611463362351)) > 0.05 { + t.Errorf("feature sum = %.9f, want -15998.611463362351", sum) + } + if got := len(downsampleGemma4AudioMask(mask, 2)); got != 7 { + t.Fatalf("downsampled mask length = %d, want 7", got) + } + softTokens := 0 + for _, value := range downsampleGemma4AudioMask(mask, 2) { + if value { + softTokens++ + } + } + if softTokens != 6 { + t.Fatalf("soft tokens = %d, want 6", softTokens) + } +} + +func TestDecodeGemma4WAVEncodings(t *testing.T) { + for _, tt := range []struct { + name string + format uint16 + bits uint16 + tol float64 + }{ + {"pcm8", 1, 8, 1.0 / 128}, + {"pcm16", 1, 16, 1.0 / 32768}, + {"pcm24", 1, 24, 1.0 / 8388608}, + {"pcm32", 1, 32, 1.0 / 2147483648}, + {"float32", 3, 32, 1e-6}, + } { + t.Run(tt.name, func(t *testing.T) { + data := makeTestWAV(t, tt.format, tt.bits, 16000, [][]float64{{-0.5}, {0}, {0.5}}) + samples, err := decodeGemma4WAV(context.Background(), data, 16000) + if err != nil { + t.Fatal(err) + } + for i, want := range []float64{-0.5, 0, 0.5} { + if math.Abs(float64(samples[i])-want) > tt.tol { + t.Errorf("sample %d = %v, want %v", i, samples[i], want) + } + } + }) + } +} + +func TestDecodeGemma4WAVExtensible(t *testing.T) { + pcmGUID := [16]byte{1, 0, 0, 0, 0, 0, 0x10, 0, 0x80, 0, 0, 0xaa, 0, 0x38, 0x9b, 0x71} + floatGUID := pcmGUID + floatGUID[0] = 3 + for _, tt := range []struct { + name string + format uint16 + bits uint16 + guid [16]byte + }{ + {"pcm", 1, 16, pcmGUID}, + {"float", 3, 32, floatGUID}, + } { + t.Run(tt.name, func(t *testing.T) { + data := makeTestExtensibleWAV(t, tt.format, tt.bits, tt.guid) + samples, err := decodeGemma4WAV(context.Background(), data, 16000) + if err != nil { + t.Fatal(err) + } + if len(samples) != 3 || math.Abs(float64(samples[0])+0.5) > 1e-5 || math.Abs(float64(samples[2])-0.5) > 1e-5 { + t.Fatalf("extensible samples = %v", samples) + } + }) + } + + valid := makeTestExtensibleWAV(t, 1, 16, pcmGUID) + tests := []struct { + name string + mutate func([]byte) []byte + want string + }{ + {"small extension", func(data []byte) []byte { + binary.LittleEndian.PutUint16(data[36:38], 21) + return data + }, "fmt size"}, + {"large extension", func(data []byte) []byte { + binary.LittleEndian.PutUint16(data[36:38], 23) + return data + }, "fmt size"}, + {"valid bits", func(data []byte) []byte { + binary.LittleEndian.PutUint16(data[38:40], 12) + return data + }, "valid bits"}, + {"GUID", func(data []byte) []byte { + data[59] ^= 1 + return data + }, "subformat"}, + {"duplicate fmt", func(data []byte) []byte { + duplicate := make([]byte, 0, len(data)+48) + duplicate = append(duplicate, data[:60]...) + duplicate = append(duplicate, data[12:60]...) + duplicate = append(duplicate, data[60:]...) + binary.LittleEndian.PutUint32(duplicate[4:8], uint32(len(duplicate)-8)) + return duplicate + }, "duplicate fmt"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data := tt.mutate(append([]byte(nil), valid...)) + if samples, err := decodeGemma4WAV(context.Background(), data, 16000); err == nil || !strings.Contains(err.Error(), tt.want) || samples != nil { + t.Fatalf("malformed extensible WAV = (%v, %v), want nil and %q error", samples, err, tt.want) + } + }) + } +} + +func TestDecodeGemma4WAVDownmixResampleAndTruncate(t *testing.T) { + frames := make([][]float64, 8000) + for i := range frames { + frames[i] = []float64{-0.5, 0.5} + } + data := makeTestWAV(t, 1, 16, 8000, frames) + samples, err := decodeGemma4WAV(context.Background(), data, 16000) + if err != nil { + t.Fatal(err) + } + if len(samples) != 16000 { + t.Fatalf("resampled length = %d, want 16000", len(samples)) + } + for _, index := range []int{0, 7999, 15999} { + if math.Abs(float64(samples[index])) > 1e-6 { + t.Errorf("downmixed sample %d = %v, want 0", index, samples[index]) + } + } + + limitFrames := make([][]float64, maxGemma4AudioSamples) + for i := range limitFrames { + limitFrames[i] = []float64{0} + } + limitData := makeTestWAV(t, 1, 16, 16000, limitFrames) + limited, err := decodeGemma4WAV(context.Background(), limitData, 16000) + if err != nil { + t.Fatal(err) + } + if len(limited) != maxGemma4AudioSamples { + t.Fatalf("exact-boundary length = %d, want %d", len(limited), maxGemma4AudioSamples) + } + overFrames := make([][]float64, len(limitFrames)+1) + copy(overFrames, limitFrames) + overFrames[len(limitFrames)] = []float64{0} + overData := makeTestWAV(t, 1, 16, 16000, overFrames) + truncated, err := decodeGemma4WAV(context.Background(), overData, 16000) + if err != nil { + t.Fatal(err) + } + if len(truncated) != maxGemma4AudioSamples { + t.Fatalf("boundary+1 truncated length = %d, want %d", len(truncated), maxGemma4AudioSamples) + } +} + +func TestPreprocessGemma4AudioSequentialIsolation(t *testing.T) { + cfg := defaultAudioProcessorConfig() + firstFrames := make([][]float64, 4000) + secondFrames := make([][]float64, 8000) + for i := range firstFrames { + firstFrames[i] = []float64{0.25 * math.Sin(2*math.Pi*440*float64(i)/16000)} + } + for i := range secondFrames { + secondFrames[i] = []float64{0.25 * math.Sin(2*math.Pi*880*float64(i)/16000)} + } + first, err := preprocessGemma4Audio(context.Background(), makeTestWAV(t, 1, 16, 16000, firstFrames), &cfg) + if err != nil { + t.Fatal(err) + } + firstSnapshot := append([]float32(nil), first.Features...) + second, err := preprocessGemma4Audio(context.Background(), makeTestWAV(t, 1, 16, 16000, secondFrames), &cfg) + if err != nil { + t.Fatal(err) + } + if first.Frames != 25 || second.Frames != 50 { + t.Fatalf("sequential frame counts = (%d, %d), want (25, 50)", first.Frames, second.Frames) + } + if len(first.Features) != len(firstSnapshot) || !equalFloat32s(first.Features, firstSnapshot) { + t.Fatal("second preprocess mutated or appended to first output") + } + if len(first.Features) > 0 && len(second.Features) > 0 && &first.Features[0] == &second.Features[0] { + t.Fatal("sequential preprocess calls share feature storage") + } +} + +func equalFloat32s(a, b []float32) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func TestGemma4AudioInputFailures(t *testing.T) { + cfg := defaultAudioProcessorConfig() + for _, tt := range []struct { + name string + data []byte + want string + }{ + {"mp3", []byte("ID3not-wav"), "RIFF/WAVE"}, + {"truncated chunk", append([]byte("RIFF\x00\x00\x00\x00WAVEdata\xff\xff\xff\x7f"), 0), "truncated"}, + {"too short", makeTestWAV(t, 1, 16, 16000, make([][]float64, 100)), "too short"}, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := preprocessGemma4Audio(context.Background(), tt.data, &cfg) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("error = %v, want containing %q", err, tt.want) + } + }) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := decodeGemma4WAV(ctx, []byte("anything"), 16000); err != context.Canceled { + t.Fatalf("canceled decode error = %v, want %v", err, context.Canceled) + } +} + +func TestGemma4AudioCancellationBeforeAllocations(t *testing.T) { + cfg := defaultAudioProcessorConfig() + frames := make([][]float64, 400) + for i := range frames { + frames[i] = []float64{0.25} + } + data := makeTestWAV(t, 1, 16, 16000, frames) + + decodeCtx := &cancelAfterErrChecks{remaining: 4} + if samples, err := decodeGemma4WAV(decodeCtx, data, 16000); err != context.Canceled || samples != nil { + t.Fatalf("decode cancellation = (%v, %v), want (nil, context.Canceled)", samples, err) + } + + resampleCtx := &cancelAfterErrChecks{remaining: 2} + if samples, err := resampleGemma4Audio(resampleCtx, make([]float32, 400), 8000, 16000); err != context.Canceled || samples != nil { + t.Fatalf("resample cancellation = (%v, %v), want (nil, context.Canceled)", samples, err) + } + + for _, checks := range []int{2, 3, 4, 5} { + t.Run(fmt.Sprintf("log-mel check %d", checks), func(t *testing.T) { + ctx := &cancelAfterErrChecks{remaining: checks} + features, mask, err := computeGemma4LogMel(ctx, make([]float32, 400), &cfg) + if err != context.Canceled || features != nil || mask != nil { + t.Fatalf("log-mel cancellation = (%v, %v, %v), want nil outputs and context.Canceled", features, mask, err) + } + }) + } + + preprocessCtx := &cancelAfterErrChecks{remaining: 6} + if input, err := preprocessGemma4Audio(preprocessCtx, data, &cfg); err != context.Canceled || input != nil { + t.Fatalf("preprocess cancellation = (%v, %v), want (nil, context.Canceled)", input, err) + } +} + +type cancelAfterErrChecks struct { + remaining int +} + +func (*cancelAfterErrChecks) Deadline() (time.Time, bool) { return time.Time{}, false } + +func (*cancelAfterErrChecks) Done() <-chan struct{} { return nil } + +func (c *cancelAfterErrChecks) Err() error { + c.remaining-- + if c.remaining <= 0 { + return context.Canceled + } + return nil +} + +func (*cancelAfterErrChecks) Value(any) any { return nil } + +func makeTestWAV(t *testing.T, format, bits uint16, sampleRate uint32, frames [][]float64) []byte { + t.Helper() + channels := 1 + if len(frames) > 0 && len(frames[0]) > 0 { + channels = len(frames[0]) + } + bytesPerSample := int(bits / 8) + var pcm bytes.Buffer + for _, frame := range frames { + if len(frame) == 0 { + frame = make([]float64, channels) + } + if len(frame) != channels { + t.Fatalf("inconsistent channel count") + } + for _, value := range frame { + switch { + case format == 1 && bits == 8: + pcm.WriteByte(byte(math.Round(value*128 + 128))) + case format == 1 && bits == 16: + _ = binary.Write(&pcm, binary.LittleEndian, int16(math.Round(value*32768))) + case format == 1 && bits == 24: + v := int32(math.Round(value * 8388608)) + pcm.Write([]byte{byte(v), byte(v >> 8), byte(v >> 16)}) + case format == 1 && bits == 32: + _ = binary.Write(&pcm, binary.LittleEndian, int32(math.Round(value*2147483648))) + case format == 3 && bits == 32: + _ = binary.Write(&pcm, binary.LittleEndian, float32(value)) + default: + t.Fatalf("unsupported test WAV encoding") + } + } + } + blockAlign := uint16(channels * bytesPerSample) + byteRate := sampleRate * uint32(blockAlign) + dataSize := pcm.Len() + var out bytes.Buffer + out.WriteString("RIFF") + _ = binary.Write(&out, binary.LittleEndian, uint32(36+dataSize)) + out.WriteString("WAVEfmt ") + _ = binary.Write(&out, binary.LittleEndian, uint32(16)) + _ = binary.Write(&out, binary.LittleEndian, format) + _ = binary.Write(&out, binary.LittleEndian, uint16(channels)) + _ = binary.Write(&out, binary.LittleEndian, sampleRate) + _ = binary.Write(&out, binary.LittleEndian, byteRate) + _ = binary.Write(&out, binary.LittleEndian, blockAlign) + _ = binary.Write(&out, binary.LittleEndian, bits) + out.WriteString("data") + _ = binary.Write(&out, binary.LittleEndian, uint32(dataSize)) + out.Write(pcm.Bytes()) + return out.Bytes() +} + +func makeTestExtensibleWAV(t *testing.T, format, bits uint16, subformat [16]byte) []byte { + t.Helper() + base := makeTestWAV(t, format, bits, 16000, [][]float64{{-0.5}, {0}, {0.5}}) + pcm := base[44:] + var out bytes.Buffer + out.WriteString("RIFF") + _ = binary.Write(&out, binary.LittleEndian, uint32(60+len(pcm))) + out.WriteString("WAVEfmt ") + _ = binary.Write(&out, binary.LittleEndian, uint32(40)) + _ = binary.Write(&out, binary.LittleEndian, uint16(0xfffe)) + _ = binary.Write(&out, binary.LittleEndian, uint16(1)) + _ = binary.Write(&out, binary.LittleEndian, uint32(16000)) + _ = binary.Write(&out, binary.LittleEndian, uint32(16000*int(bits/8))) + _ = binary.Write(&out, binary.LittleEndian, bits/8) + _ = binary.Write(&out, binary.LittleEndian, bits) + _ = binary.Write(&out, binary.LittleEndian, uint16(22)) + _ = binary.Write(&out, binary.LittleEndian, bits) + _ = binary.Write(&out, binary.LittleEndian, uint32(0)) + out.Write(subformat[:]) + out.WriteString("data") + _ = binary.Write(&out, binary.LittleEndian, uint32(len(pcm))) + out.Write(pcm) + return out.Bytes() +} From c59f098bc8cf79f1793d52fecbd468a4a4a02ff0 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 00:27:41 +0000 Subject: [PATCH 48/58] gemma4: execute native MLX audio embeddings Load the released conformer audio tower, prepare ordered WAV segments into cache-keyed feature rows, encode them lazily from runner-owned media data, and scatter projected audio features through the official media contract alongside images. Co-authored-by: Codex --- docs/third-party/mlx-vlm.md | 8 + x/models/gemma4/audio.go | 469 +++++++++++++++++++++++++++++++++ x/models/gemma4/audio_test.go | 244 +++++++++++++++++ x/models/gemma4/gemma4.go | 66 ++++- x/models/gemma4/vision.go | 161 +++++++---- x/models/gemma4/vision_test.go | 114 +++++++- 6 files changed, 1007 insertions(+), 55 deletions(-) create mode 100644 x/models/gemma4/audio.go create mode 100644 x/models/gemma4/audio_test.go diff --git a/docs/third-party/mlx-vlm.md b/docs/third-party/mlx-vlm.md index 0233657227c..965fa654590 100644 --- a/docs/third-party/mlx-vlm.md +++ b/docs/third-party/mlx-vlm.md @@ -9,6 +9,14 @@ implementation in MLX-VLM: `mlx_vlm/models/gemma4/vision.py`, and `mlx_vlm/models/gemma4_unified/gemma4_unified.py` +The native Gemma 4 audio encoder in `x/models/gemma4/audio.go` is adapted from +the same project's Gemma 4 Conformer implementation at revision +`84f43753380355c0455a2bafb291d4b7cbcf81d1` (MLX-VLM v0.6.5), source path +`mlx_vlm/models/gemma4/audio.py`. Its architecture and numerical behavior were +also cross-checked against Hugging Face Transformers revision +`dff4572dfa4bfa9f00cc8414e4b84877552fefe9`. No MLX-VLM or Transformers runtime +dependency is included. + The adapted implementation is distributed under the following license. ```text diff --git a/x/models/gemma4/audio.go b/x/models/gemma4/audio.go new file mode 100644 index 00000000000..db65f192e71 --- /dev/null +++ b/x/models/gemma4/audio.go @@ -0,0 +1,469 @@ +package gemma4 + +// The Gemma 4 audio encoder follows the Universal Speech Model Conformer used +// by the released checkpoint. The MLX layout and execution structure are +// adapted from MLX-VLM's MIT-licensed implementation pinned in +// docs/third-party/mlx-vlm.md and cross-checked against Transformers. + +import ( + "encoding/json" + "errors" + "fmt" + "math" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" + "github.com/ollama/ollama/x/models/nn" +) + +type AudioConfig struct { + AttentionChunkSize int32 `json:"attention_chunk_size"` + AttentionContextLeft int32 `json:"attention_context_left"` + AttentionContextRight int32 `json:"attention_context_right"` + AttentionInvalidLogit float32 `json:"attention_invalid_logits_value"` + AttentionLogitCap float32 `json:"attention_logit_cap"` + ConvKernelSize int32 `json:"conv_kernel_size"` + GradientClipping float32 `json:"gradient_clipping"` + HiddenSize int32 `json:"hidden_size"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + OutputProjDims int32 `json:"output_proj_dims"` + ResidualWeight float32 `json:"residual_weight"` + RMSNormEps float32 `json:"rms_norm_eps"` + SubsamplingConvChannels []int32 `json:"subsampling_conv_channels"` + UseClippedLinears bool `json:"use_clipped_linears"` +} + +func parseAudioConfig(configData []byte) (*AudioConfig, error) { + var wrapped struct { + AudioConfig *AudioConfig `json:"audio_config"` + } + if err := json.Unmarshal(configData, &wrapped); err != nil { + return nil, fmt.Errorf("parse Gemma4 audio config: %w", err) + } + if wrapped.AudioConfig == nil { + return nil, nil + } + cfg := wrapped.AudioConfig + if cfg.HiddenSize <= 0 || cfg.NumHiddenLayers <= 0 || cfg.NumAttentionHeads <= 0 || + cfg.HiddenSize%cfg.NumAttentionHeads != 0 || cfg.OutputProjDims <= 0 || + cfg.AttentionChunkSize <= 0 || cfg.AttentionContextLeft <= 0 || cfg.AttentionContextRight < 0 || + cfg.ConvKernelSize <= 0 || cfg.ConvKernelSize%2 == 0 || len(cfg.SubsamplingConvChannels) != 2 || + cfg.SubsamplingConvChannels[0] <= 0 || cfg.SubsamplingConvChannels[1] <= 0 || + cfg.RMSNormEps <= 0 || cfg.ResidualWeight <= 0 || cfg.GradientClipping <= 0 || + cfg.AttentionLogitCap <= 0 || cfg.AttentionInvalidLogit >= 0 { + return nil, errors.New("invalid Gemma4 audio configuration") + } + return cfg, nil +} + +type audioConvBlock struct { + Weight *mlx.Array + Norm *nn.LayerNorm +} + +type audioFeedForward struct { + PreNorm, PostNorm *nn.RMSNorm + Up, Down *ClippableLinear + Config *AudioConfig +} + +type audioAttention struct { + Q, K, V, Output *ClippableLinear + RelativeK nn.LinearLayer + RelativeKDType mlx.DType + PerDimScale *mlx.Array + Config *AudioConfig +} + +type audioLightConv struct { + PreNorm, ConvNorm *nn.RMSNorm + Start, End *ClippableLinear + DepthwiseWeight *mlx.Array + Config *AudioConfig +} + +type audioConformerBlock struct { + FeedForward1, FeedForward2 *audioFeedForward + Attention *audioAttention + LightConv *audioLightConv + PreAttentionNorm *nn.RMSNorm + PostAttentionNorm *nn.RMSNorm + OutputNorm *nn.RMSNorm + Config *AudioConfig +} + +type AudioModel struct { + Conv0, Conv1 *audioConvBlock + InputProj nn.LinearLayer + Layers []*audioConformerBlock + OutputProj nn.LinearLayer + Config *AudioConfig +} + +func hasCompleteGemma4AudioWeights(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32) bool { + if cfg == nil { + return false + } + names := make([]string, 0, len(tensors)) + for name, tensor := range tensors { + if tensor != nil { + names = append(names, name) + } + } + return gemma4metadata.ValidateAudioTensors(audioMetadataConfig(cfg, textHidden), names) == nil +} + +func audioMetadataConfig(cfg *AudioConfig, textHidden int32) gemma4metadata.ConfigFile { + channels := make([]int, len(cfg.SubsamplingConvChannels)) + for i, channel := range cfg.SubsamplingConvChannels { + channels[i] = int(channel) + } + return gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: int(textHidden)}, + AudioConfig: &gemma4metadata.AudioConfig{ + AttentionChunkSize: int(cfg.AttentionChunkSize), AttentionContextLeft: int(cfg.AttentionContextLeft), + AttentionContextRight: int(cfg.AttentionContextRight), ConvKernelSize: int(cfg.ConvKernelSize), + HiddenSize: int(cfg.HiddenSize), NumAttentionHeads: int(cfg.NumAttentionHeads), + NumHiddenLayers: int(cfg.NumHiddenLayers), OutputProjDims: int(cfg.OutputProjDims), + SubsamplingConvChannels: channels, UseClippedLinears: cfg.UseClippedLinears, + }, + } +} + +func validateGemma4AudioWeights(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32) error { + required, err := gemma4metadata.RequiredAudioTensorShapes(audioMetadataConfig(cfg, textHidden)) + if err != nil { + return err + } + for name, shape := range required { + tensor := tensors[name] + if tensor == nil { + return fmt.Errorf("missing Gemma4 audio tensor %s", name) + } + want := make([]int, len(shape)) + for i, dim := range shape { + want[i] = int(dim) + } + if !equalIntShape(tensor.Dims(), want) { + return fmt.Errorf("Gemma4 audio tensor %s shape %v, want %v", name, tensor.Dims(), want) + } + if !supportedGemma4AudioDType(tensor.DType()) { + return fmt.Errorf("Gemma4 audio tensor %s has unsupported dtype %s", name, tensor.DType()) + } + } + return nil +} + +func supportedGemma4AudioDType(dtype mlx.DType) bool { + switch dtype { + case mlx.DTypeBFloat16, mlx.DTypeFloat16, mlx.DTypeFloat32: + return true + default: + return false + } +} + +func equalIntShape(a, b []int) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func loadAudioModel(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32, groupSize, bits int, mode string, tq map[string]*model.TensorQuantInfo) (*AudioModel, error) { + if err := validateGemma4AudioWeights(tensors, cfg, textHidden); err != nil { + return nil, err + } + const prefix = "model.audio_tower." + linears := model.NewLinearFactory(tensors, groupSize, bits, mode, tq) + loadConv := func(path string) (*audioConvBlock, error) { + weight := tensors[path+".conv.weight"] + norm := tensors[path+".norm.weight"] + if weight == nil || norm == nil { + return nil, fmt.Errorf("missing Gemma4 audio convolution tensors at %s", path) + } + // Safetensors uses PyTorch OIHW; MLX Conv2d expects OHWI. + weight = mlx.Transpose(weight, 0, 2, 3, 1) + return &audioConvBlock{Weight: weight, Norm: &nn.LayerNorm{Weight: norm, Eps: cfg.RMSNormEps}}, nil + } + conv0, err := loadConv(prefix + "subsample_conv_projection.layer0") + if err != nil { + return nil, err + } + conv1, err := loadConv(prefix + "subsample_conv_projection.layer1") + if err != nil { + return nil, err + } + inputProj := linears.Make(prefix + "subsample_conv_projection.input_proj_linear") + outputProj := linears.Make(prefix + "output_proj") + if inputProj == nil || outputProj == nil { + return nil, errors.New("missing Gemma4 audio input or output projection") + } + + out := &AudioModel{Conv0: conv0, Conv1: conv1, InputProj: inputProj, OutputProj: outputProj, Config: cfg} + out.Layers = make([]*audioConformerBlock, cfg.NumHiddenLayers) + for i := range cfg.NumHiddenLayers { + path := fmt.Sprintf("%slayers.%d.", prefix, i) + makeNorm := func(name string) (*nn.RMSNorm, error) { + weight := tensors[path+name+".weight"] + if weight == nil { + return nil, fmt.Errorf("missing Gemma4 audio tensor %s.weight", path+name) + } + return nn.NewRMSNorm(weight, cfg.RMSNormEps), nil + } + makeFF := func(name string) (*audioFeedForward, error) { + pre, err := makeNorm(name + ".pre_layer_norm") + if err != nil { + return nil, err + } + post, err := makeNorm(name + ".post_layer_norm") + if err != nil { + return nil, err + } + up := makeClippableLinear(tensors, linears, path+name+".ffw_layer_1", cfg.UseClippedLinears) + down := makeClippableLinear(tensors, linears, path+name+".ffw_layer_2", cfg.UseClippedLinears) + if up == nil || down == nil { + return nil, fmt.Errorf("missing Gemma4 audio feed-forward tensors at %s%s", path, name) + } + return &audioFeedForward{PreNorm: pre, PostNorm: post, Up: up, Down: down, Config: cfg}, nil + } + ff1, err := makeFF("feed_forward1") + if err != nil { + return nil, err + } + ff2, err := makeFF("feed_forward2") + if err != nil { + return nil, err + } + preAttn, err := makeNorm("norm_pre_attn") + if err != nil { + return nil, err + } + postAttn, err := makeNorm("norm_post_attn") + if err != nil { + return nil, err + } + outNorm, err := makeNorm("norm_out") + if err != nil { + return nil, err + } + convPre, err := makeNorm("lconv1d.pre_layer_norm") + if err != nil { + return nil, err + } + convNorm, err := makeNorm("lconv1d.conv_norm") + if err != nil { + return nil, err + } + depthwise := tensors[path+"lconv1d.depthwise_conv1d.weight"] + if depthwise == nil { + return nil, fmt.Errorf("missing Gemma4 audio depthwise convolution at %s", path) + } + // PyTorch [channels, 1, kernel] to MLX [channels, kernel, 1]. + depthwise = mlx.Transpose(depthwise, 0, 2, 1) + convStart := makeClippableLinear(tensors, linears, path+"lconv1d.linear_start", cfg.UseClippedLinears) + convEnd := makeClippableLinear(tensors, linears, path+"lconv1d.linear_end", cfg.UseClippedLinears) + q := makeClippableLinear(tensors, linears, path+"self_attn.q_proj", cfg.UseClippedLinears) + k := makeClippableLinear(tensors, linears, path+"self_attn.k_proj", cfg.UseClippedLinears) + v := makeClippableLinear(tensors, linears, path+"self_attn.v_proj", cfg.UseClippedLinears) + attnOut := makeClippableLinear(tensors, linears, path+"self_attn.post", cfg.UseClippedLinears) + relativeK := linears.Make(path + "self_attn.relative_k_proj") + perDimScale := tensors[path+"self_attn.per_dim_scale"] + if convStart == nil || convEnd == nil || q == nil || k == nil || v == nil || attnOut == nil || relativeK == nil || perDimScale == nil { + return nil, fmt.Errorf("missing Gemma4 audio attention or convolution tensors at %s", path) + } + out.Layers[i] = &audioConformerBlock{ + FeedForward1: ff1, FeedForward2: ff2, + Attention: &audioAttention{ + Q: q, K: k, V: v, Output: attnOut, RelativeK: relativeK, + RelativeKDType: tensors[path+"self_attn.relative_k_proj.weight"].DType(), + PerDimScale: perDimScale, Config: cfg, + }, + LightConv: &audioLightConv{PreNorm: convPre, ConvNorm: convNorm, Start: convStart, End: convEnd, DepthwiseWeight: depthwise, Config: cfg}, + PreAttentionNorm: preAttn, PostAttentionNorm: postAttn, OutputNorm: outNorm, Config: cfg, + } + } + return out, nil +} + +func audioValidityArray(valid []bool, rank int) *mlx.Array { + shape := make([]int, rank) + shape[0], shape[1] = 1, len(valid) + for i := 2; i < rank; i++ { + shape[i] = 1 + } + values := make([]float32, len(valid)) + for i, ok := range valid { + if ok { + values[i] = 1 + } + } + return mlx.FromValues(values, shape...) +} + +func (b *audioConvBlock) Forward(x *mlx.Array, valid []bool) (*mlx.Array, []bool) { + x = mlx.Mul(x, audioValidityArray(valid, 4).AsType(x.DType())) + x = mlx.PadConstant(x, []int{1, 2}, []int{1, 1}, []int{1, 1}) + x = mlx.Conv2d(x, b.Weight, 2, 2, 0, 0, 1, 1, 1) + x = mlx.ReLU(b.Norm.Forward(x)) + return x, downsampleGemma4AudioMask(valid, 1) +} + +func (m *AudioModel) Forward(features *mlx.Array, input *gemma4AudioInput) *mlx.Array { + x := mlx.Reshape(features, 1, int32(input.Frames), 128) + valid := append([]bool(nil), input.FeatureMask...) + x = mlx.ExpandDims(x, -1) + x, valid = m.Conv0.Forward(x, valid) + x, valid = m.Conv1.Forward(x, valid) + x = mlx.Reshape(x, 1, int32(x.Dim(1)), int32(x.Dim(2)*x.Dim(3))) + x = m.InputProj.Forward(x) + for _, layer := range m.Layers { + x = layer.Forward(x, valid) + } + x = m.OutputProj.Forward(x) + x = mlx.Mul(x, audioValidityArray(valid, 3).AsType(x.DType())) + return mlx.SliceStartStop(x, []int32{0, 0, 0}, []int32{1, int32(input.SoftTokens), int32(x.Dim(2))}) +} + +func (f *audioFeedForward) Forward(x *mlx.Array) *mlx.Array { + residual := x + x = mlx.Clamp(x, -f.Config.GradientClipping, f.Config.GradientClipping) + x = f.PreNorm.Forward(x, 0) + x = f.Up.Forward(x) + x = mlx.Mul(x, mlx.Sigmoid(x)) // SiLU + x = f.Down.Forward(x) + x = mlx.Clamp(x, -f.Config.GradientClipping, f.Config.GradientClipping) + x = f.PostNorm.Forward(x, 0) + return mlx.Add(residual, mlx.MulScalar(x, f.Config.ResidualWeight)) +} + +func padAudioTime(x *mlx.Array, left, right int) *mlx.Array { + if left == 0 && right == 0 { + return x + } + return mlx.PadConstant(x, []int{1}, []int{left}, []int{right}) +} + +func (a *audioAttention) relativeLogits(q *mlx.Array, blocks, chunk, context int) *mlx.Array { + cfg := a.Config + headDim := cfg.HiddenSize / cfg.NumAttentionHeads + span := cfg.AttentionContextLeft + cfg.AttentionContextRight + half := cfg.HiddenSize / 2 + values := make([]float32, int(span*cfg.HiddenSize)) + logIncrement := math.Log(10000) / float64(max(half-1, 1)) + for p := range span { + position := float64(cfg.AttentionContextLeft - 1 - p) + for d := range half { + angle := position * math.Exp(-float64(d)*logIncrement) + values[int(p*cfg.HiddenSize+d)] = float32(math.Sin(angle)) + values[int(p*cfg.HiddenSize+half+d)] = float32(math.Cos(angle)) + } + } + position := mlx.FromValues(values, int(span), int(cfg.HiddenSize)).AsType(a.RelativeKDType) + position = a.RelativeK.Forward(position) + position = position.AsType(q.DType()) + position = mlx.Reshape(position, span, cfg.NumAttentionHeads, headDim) + position = mlx.Transpose(position, 1, 2, 0) + qFlat := mlx.Reshape(q, 1, cfg.NumAttentionHeads, int32(blocks*chunk), headDim) + term := mlx.Matmul(qFlat, position) + term = mlx.Reshape(term, 1, cfg.NumAttentionHeads, int32(blocks), int32(chunk), span) + pad := context + 1 - int(span) + term = mlx.PadConstant(term, []int{4}, []int{0}, []int{pad}) + term = mlx.Reshape(term, 1, cfg.NumAttentionHeads, int32(blocks), int32(chunk*(context+1))) + term = mlx.SliceStartStop(term, []int32{0, 0, 0, 0}, []int32{1, cfg.NumAttentionHeads, int32(blocks), int32(chunk * context)}) + return mlx.Reshape(term, 1, cfg.NumAttentionHeads, int32(blocks), int32(chunk), int32(context)) +} + +func audioAttentionMask(valid []bool, blocks, chunk, context, left, future int) *mlx.Array { + values := audioAttentionMaskValues(valid, blocks, chunk, context, left, future) + return mlx.FromValues(values, 1, 1, blocks, chunk, context) +} + +func audioAttentionMaskValues(valid []bool, blocks, chunk, context, left, future int) []bool { + values := make([]bool, blocks*chunk*context) + for u := range blocks { + for w := range chunk { + for c := range context { + actual := u*chunk + c - left + ok := c >= w && c <= w+left+future && actual >= 0 && actual < len(valid) && valid[actual] + values[(u*chunk+w)*context+c] = ok + } + } + } + return values +} + +func (a *audioAttention) Forward(x *mlx.Array, valid []bool) *mlx.Array { + cfg := a.Config + length := x.Dim(1) + heads, headDim := int(cfg.NumAttentionHeads), int(cfg.HiddenSize/cfg.NumAttentionHeads) + chunk := int(cfg.AttentionChunkSize) + left, future := int(cfg.AttentionContextLeft-1), int(cfg.AttentionContextRight) + context := chunk + left + future + blocks := (length + chunk - 1) / chunk + q := mlx.Reshape(a.Q.Forward(x).AsType(mlx.DTypeFloat32), 1, int32(length), cfg.NumAttentionHeads, int32(headDim)) + k := mlx.Reshape(a.K.Forward(x).AsType(mlx.DTypeFloat32), 1, int32(length), cfg.NumAttentionHeads, int32(headDim)) + v := mlx.Reshape(a.V.Forward(x).AsType(mlx.DTypeFloat32), 1, int32(length), cfg.NumAttentionHeads, int32(headDim)) + qScale := float32(math.Pow(float64(headDim), -0.5) / math.Log(2)) + q = mlx.Mul(q, mlx.MulScalar(mlx.Softplus(a.PerDimScale).AsType(q.DType()), qScale)) + k = mlx.MulScalar(k, float32(math.Log(1+math.E)/math.Log(2))) + q = padAudioTime(q, 0, blocks*chunk-length) + q = mlx.Reshape(q, 1, int32(blocks), int32(chunk), int32(heads), int32(headDim)) + q = mlx.Transpose(q, 0, 3, 1, 2, 4) + indices := make([]int32, blocks*context) + for u := range blocks { + for c := range context { + indices[u*context+c] = int32(u*chunk + c) + } + } + indicesArray := mlx.FromValues(indices, blocks, context) + k = mlx.Take(padAudioTime(k, left, future+chunk-1), indicesArray, 1) + v = mlx.Take(padAudioTime(v, left, future+chunk-1), indicesArray, 1) + k = mlx.Transpose(k, 0, 3, 1, 4, 2) + content := mlx.Matmul(q, k) + logits := mlx.Add(content, a.relativeLogits(q, blocks, chunk, context)) + logits = mlx.MulScalar(mlx.DivScalar(logits, cfg.AttentionLogitCap).Tanh(), cfg.AttentionLogitCap) + condition := audioAttentionMask(valid, blocks, chunk, context, left, future) + logits = mlx.Where(condition, logits, mlx.FromValue(cfg.AttentionInvalidLogit)) + probs := mlx.SoftmaxAxis(logits, -1, true) + v = mlx.Transpose(v, 0, 3, 1, 2, 4) + result := mlx.Matmul(probs, v) + result = mlx.Transpose(result, 0, 2, 3, 1, 4) + result = mlx.Reshape(result, 1, int32(blocks*chunk), cfg.HiddenSize) + result = mlx.SliceStartStop(result, []int32{0, 0, 0}, []int32{1, int32(length), cfg.HiddenSize}) + return a.Output.Forward(result) +} + +func (c *audioLightConv) Forward(x *mlx.Array) *mlx.Array { + residual := x + x = c.PreNorm.Forward(x, 0) + x = mlx.GLU(c.Start.Forward(x)) + x = padAudioTime(x, int(c.Config.ConvKernelSize-1), 0) + x = mlx.Conv1d(x, c.DepthwiseWeight, nil, 1, 0, 1, c.Config.HiddenSize) + x = mlx.Clamp(x, -c.Config.GradientClipping, c.Config.GradientClipping) + x = c.ConvNorm.Forward(x, 0) + x = mlx.Mul(x, mlx.Sigmoid(x)) + return mlx.Add(c.End.Forward(x), residual) +} + +func (b *audioConformerBlock) Forward(x *mlx.Array, valid []bool) *mlx.Array { + x = b.FeedForward1.Forward(x) + residual := x + x = mlx.Clamp(x, -b.Config.GradientClipping, b.Config.GradientClipping) + x = b.PreAttentionNorm.Forward(x, 0) + x = b.Attention.Forward(x, valid) + x = mlx.Clamp(x, -b.Config.GradientClipping, b.Config.GradientClipping) + x = mlx.Add(residual, b.PostAttentionNorm.Forward(x, 0)) + x = mlx.Mul(x, audioValidityArray(valid, 3).AsType(x.DType())) + x = b.LightConv.Forward(x) + x = b.FeedForward2.Forward(x) + x = mlx.Clamp(x, -b.Config.GradientClipping, b.Config.GradientClipping) + return b.OutputNorm.Forward(x, 0) +} diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go new file mode 100644 index 00000000000..b66f74df4ca --- /dev/null +++ b/x/models/gemma4/audio_test.go @@ -0,0 +1,244 @@ +package gemma4 + +import ( + "maps" + "strings" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" +) + +const releasedGemma4AudioConfig = `{ + "audio_config": { + "attention_chunk_size": 12, + "attention_context_left": 13, + "attention_context_right": 0, + "attention_invalid_logits_value": -1000000000.0, + "attention_logit_cap": 50.0, + "conv_kernel_size": 5, + "gradient_clipping": 10000000000.0, + "hidden_size": 1024, + "num_attention_heads": 8, + "num_hidden_layers": 12, + "output_proj_dims": 1536, + "residual_weight": 0.5, + "rms_norm_eps": 0.000001, + "subsampling_conv_channels": [128, 32], + "use_clipped_linears": true + } +}` + +func TestParseReleasedAudioConfig(t *testing.T) { + cfg, err := parseAudioConfig([]byte(releasedGemma4AudioConfig)) + if err != nil { + t.Fatal(err) + } + if cfg == nil || cfg.HiddenSize != 1024 || cfg.NumHiddenLayers != 12 || cfg.OutputProjDims != 1536 { + t.Fatalf("audio config = %+v", cfg) + } + if cfg.AttentionContextLeft != 13 || cfg.AttentionChunkSize != 12 || cfg.ConvKernelSize != 5 { + t.Fatalf("audio context config = %+v", cfg) + } + + if cfg, err := parseAudioConfig([]byte(`{"model_type":"gemma4"}`)); err != nil || cfg != nil { + t.Fatalf("missing audio config = %+v, %v; want nil, nil", cfg, err) + } + bad := strings.Replace(releasedGemma4AudioConfig, `"num_attention_heads": 8`, `"num_attention_heads": 7`, 1) + if _, err := parseAudioConfig([]byte(bad)); err == nil { + t.Fatal("non-divisible head count: error = nil") + } +} + +func TestParseTextConfigRejectsInvalidAudioMarkers(t *testing.T) { + valid := `"boi_token_id":1,"image_token_id":2,"eoi_token_id":3,"boa_token_id":4,"audio_token_id":5,"eoa_token_index":6,"text_config":{"vocab_size":7}` + if _, err := parseTextConfig([]byte(`{` + valid + `}`)); err != nil { + t.Fatalf("exact marker boundary: %v", err) + } + tests := []struct { + name string + json string + }{ + {"negative begin", `{"boa_token_id":-1}`}, + {"negative audio", `{"audio_token_id":-1}`}, + {"negative end", `{"eoa_token_index":-1}`}, + {"begin outside vocab", `{"boa_token_id":262144}`}, + {"audio outside vocab", `{"audio_token_id":262144}`}, + {"end outside vocab", `{"eoa_token_index":262144}`}, + {"duplicate begin and audio", `{"boa_token_id":258881,"audio_token_id":258881}`}, + {"duplicate audio and end", `{"audio_token_id":258883,"eoa_token_index":258883}`}, + {"duplicate begin and end", `{"boa_token_id":258883,"eoa_token_index":258883}`}, + {"audio begin duplicates image begin", `{"boa_token_id":255999}`}, + {"audio begin duplicates image", `{"boa_token_id":258880}`}, + {"audio begin duplicates image end", `{"boa_token_id":258882}`}, + {"audio token duplicates image begin", `{"audio_token_id":255999}`}, + {"audio token duplicates image", `{"audio_token_id":258880}`}, + {"audio token duplicates image end", `{"audio_token_id":258882}`}, + {"audio end duplicates image begin", `{"eoa_token_index":255999}`}, + {"audio end duplicates image", `{"eoa_token_index":258880}`}, + {"audio end duplicates image end", `{"eoa_token_index":258882}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if _, err := parseTextConfig([]byte(tt.json)); err == nil { + t.Fatal("parseTextConfig() error = nil") + } + }) + } + cfg, err := parseTextConfig([]byte(`{` + valid + `}`)) + if err != nil { + t.Fatal(err) + } + for _, tt := range []struct { + name string + mutate func(*TextConfig) + }{ + {"zero begin", func(cfg *TextConfig) { cfg.BOATokenIDValue = 0 }}, + {"zero audio", func(cfg *TextConfig) { cfg.AudioTokenIDValue = 0 }}, + {"zero end", func(cfg *TextConfig) { cfg.EOATokenIDValue = 0 }}, + } { + t.Run(tt.name, func(t *testing.T) { + candidate := cfg + tt.mutate(&candidate) + if err := validateGemma4MediaTokenConfig(&candidate); err == nil { + t.Fatal("marker validator error = nil") + } + }) + } +} + +func TestAudioAttentionMaskValues(t *testing.T) { + valid := []bool{true, true, true, false} + got := audioAttentionMaskValues(valid, 2, 2, 4, 2, 0) + allowed := func(block, query, context int) bool { + return got[(block*2+query)*4+context] + } + for _, tt := range []struct { + block, query, context int + want bool + }{ + {0, 0, 0, false}, + {0, 0, 2, true}, + {0, 0, 3, false}, + {0, 1, 1, false}, + {0, 1, 3, true}, + {1, 0, 0, true}, + {1, 0, 2, true}, + {1, 0, 3, false}, + {1, 1, 3, false}, + } { + if value := allowed(tt.block, tt.query, tt.context); value != tt.want { + t.Errorf("mask[%d,%d,%d] = %v, want %v", tt.block, tt.query, tt.context, value, tt.want) + } + } +} + +func TestParseGemma4AudioTokens(t *testing.T) { + tokens := parseGemma4MediaTokens([]byte(`{ + "boa_token":"audio-start","audio_token":"audio-soft","eoa_token":"audio-end" + }`), defaultGemma4MediaTokens()) + if tokens.BOA != "audio-start" || tokens.Audio != "audio-soft" || tokens.EOA != "audio-end" { + t.Fatalf("audio tokens = %+v", tokens) + } + if tokens.Image != defaultGemma4ImageToken { + t.Fatalf("image token = %q, want unchanged default", tokens.Image) + } +} + +func TestMediaTokenSpanAudio(t *testing.T) { + start, end, err := mediaTokenSpan([]int32{1, 258881, 258881, 2}, 258881) + if err != nil || start != 1 || end != 3 { + t.Fatalf("mediaTokenSpan() = %d, %d, %v; want 1, 3, nil", start, end, err) + } +} + +func TestSupportedGemma4AudioDType(t *testing.T) { + for _, dtype := range []mlx.DType{mlx.DTypeBFloat16, mlx.DTypeFloat16, mlx.DTypeFloat32} { + if !supportedGemma4AudioDType(dtype) { + t.Errorf("supportedGemma4AudioDType(%s) = false", dtype) + } + } + for _, dtype := range []mlx.DType{mlx.DTypeInt32, mlx.DTypeUint8, mlx.DTypeFloat64} { + if supportedGemma4AudioDType(dtype) { + t.Errorf("supportedGemma4AudioDType(%s) = true", dtype) + } + } +} + +func TestValidateAndLoadGemma4AudioWeights(t *testing.T) { + useMLXTestThread(t) + + cfg := tinyAudioConfig() + const textHidden = int32(6) + for _, dtype := range []mlx.DType{mlx.DTypeBFloat16, mlx.DTypeFloat16, mlx.DTypeFloat32} { + t.Run(dtype.String(), func(t *testing.T) { + tensors := actualAudioTensors(t, cfg, textHidden, dtype) + if err := validateGemma4AudioWeights(tensors, cfg, textHidden); err != nil { + t.Fatalf("validateGemma4AudioWeights(%s): %v", dtype, err) + } + loaded, err := loadAudioModel(tensors, cfg, textHidden, 0, 0, "", nil) + if err != nil || loaded == nil { + t.Fatalf("loadAudioModel(%s) = (%v, %v)", dtype, loaded, err) + } + embed, err := loadMultimodalEmbedder(tensors, "embed_audio", cfg.RMSNormEps, 0, 0, "", nil) + if err != nil || embed == nil { + t.Fatalf("loadMultimodalEmbedder(%s) = (%v, %v)", dtype, embed, err) + } + }) + } + + valid := actualAudioTensors(t, cfg, textHidden, mlx.DTypeFloat32) + const target = "model.audio_tower.output_proj.weight" + tests := []struct { + name string + mutate func(map[string]*mlx.Array) + }{ + {"missing", func(tensors map[string]*mlx.Array) { delete(tensors, target) }}, + {"wrong shape", func(tensors map[string]*mlx.Array) { tensors[target] = mlx.Zeros(mlx.DTypeFloat32, 1, 1) }}, + {"integer dtype", func(tensors map[string]*mlx.Array) { + tensors[target] = mlx.Zeros(mlx.DTypeInt32, int(cfg.OutputProjDims), int(cfg.HiddenSize)) + }}, + {"float64 dtype", func(tensors map[string]*mlx.Array) { + tensors[target] = mlx.Zeros(mlx.DTypeFloat64, int(cfg.OutputProjDims), int(cfg.HiddenSize)) + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tensors := maps.Clone(valid) + tt.mutate(tensors) + if err := validateGemma4AudioWeights(tensors, cfg, textHidden); err == nil { + t.Fatal("validateGemma4AudioWeights() error = nil") + } + if loaded, err := loadAudioModel(tensors, cfg, textHidden, 0, 0, "", nil); err == nil || loaded != nil { + t.Fatalf("loadAudioModel() = (%v, %v), want (nil, error)", loaded, err) + } + }) + } +} + +func tinyAudioConfig() *AudioConfig { + return &AudioConfig{ + AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionContextRight: 0, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, + GradientClipping: 1e4, HiddenSize: 4, NumAttentionHeads: 2, + NumHiddenLayers: 1, OutputProjDims: 4, ResidualWeight: 0.5, + RMSNormEps: 1e-6, SubsamplingConvChannels: []int32{2, 2}, + } +} + +func actualAudioTensors(t *testing.T, cfg *AudioConfig, textHidden int32, dtype mlx.DType) map[string]*mlx.Array { + t.Helper() + shapes, err := gemma4metadata.RequiredAudioTensorShapes(audioMetadataConfig(cfg, textHidden)) + if err != nil { + t.Fatal(err) + } + tensors := make(map[string]*mlx.Array, len(shapes)) + for name, shape := range shapes { + dims := make([]int, len(shape)) + for i, dim := range shape { + dims[i] = int(dim) + } + tensors[name] = mlx.Zeros(dtype, dims...) + } + return tensors +} diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index e20d8f77f72..0f19533fbba 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -70,6 +70,8 @@ type TextConfig struct { AudioTokenIDValue int32 `json:"-"` BOITokenIDValue int32 `json:"-"` EOITokenIDValue int32 `json:"-"` + BOATokenIDValue int32 `json:"-"` + EOATokenIDValue int32 `json:"-"` VisionSoftTokens int32 `json:"-"` // Quantization parameters. @@ -381,6 +383,8 @@ type Model struct { Vision *VisionModel UnifiedVision *UnifiedVisionEmbedder EmbedVision *MultimodalEmbedder + Audio *AudioModel + EmbedAudio *MultimodalEmbedder // PLE model-level components (nil if no PLE). EmbedTokensPerLayer nn.EmbeddingLayer @@ -393,7 +397,9 @@ type Model struct { tok *tokenizer.Tokenizer *TextConfig - VisionConfig *VisionConfig + VisionConfig *VisionConfig + AudioConfig *AudioConfig + AudioProcessorConfig *AudioProcessorConfig SuppressLogitBias *mlx.Array weightPrefix string @@ -421,6 +427,8 @@ func parseTextConfig(configData []byte) (TextConfig, error) { AudioTokenID int32 `json:"audio_token_id"` BOITokenID int32 `json:"boi_token_id"` EOITokenID int32 `json:"eoi_token_id"` + BOATokenID int32 `json:"boa_token_id"` + EOATokenID int32 `json:"eoa_token_index"` VisionSoftTokensPerImage int32 `json:"vision_soft_tokens_per_image"` } if err := json.Unmarshal(configData, &top); err != nil { @@ -432,6 +440,8 @@ func parseTextConfig(configData []byte) (TextConfig, error) { cfg.AudioTokenIDValue = top.AudioTokenID cfg.BOITokenIDValue = top.BOITokenID cfg.EOITokenIDValue = top.EOITokenID + cfg.BOATokenIDValue = top.BOATokenID + cfg.EOATokenIDValue = top.EOATokenID cfg.VisionSoftTokens = top.VisionSoftTokensPerImage // Apply defaults. @@ -471,10 +481,16 @@ func parseTextConfig(configData []byte) (TextConfig, error) { if cfg.EOITokenIDValue == 0 { cfg.EOITokenIDValue = 258882 } + if cfg.BOATokenIDValue == 0 { + cfg.BOATokenIDValue = 256000 + } + if cfg.EOATokenIDValue == 0 { + cfg.EOATokenIDValue = 258883 + } if cfg.VisionSoftTokens == 0 { cfg.VisionSoftTokens = 280 } - if err := validateGemma4ImageTokenConfig(&cfg); err != nil { + if err := validateGemma4MediaTokenConfig(&cfg); err != nil { return TextConfig{}, err } @@ -557,7 +573,7 @@ func parseTextConfig(configData []byte) (TextConfig, error) { return cfg, nil } -func validateGemma4ImageTokenConfig(cfg *TextConfig) error { +func validateGemma4MediaTokenConfig(cfg *TextConfig) error { if cfg.VisionSoftTokens <= 0 || cfg.VisionSoftTokens > maxGemma4VisionSoftTokens { return fmt.Errorf("invalid Gemma4 vision soft-token count %d", cfg.VisionSoftTokens) } @@ -568,10 +584,13 @@ func validateGemma4ImageTokenConfig(cfg *TextConfig) error { {"boi_token_id", cfg.BOITokenIDValue}, {"image_token_id", cfg.ImageTokenIDValue}, {"eoi_token_id", cfg.EOITokenIDValue}, + {"boa_token_id", cfg.BOATokenIDValue}, + {"audio_token_id", cfg.AudioTokenIDValue}, + {"eoa_token_index", cfg.EOATokenIDValue}, } seen := make(map[int32]string, len(tokens)) for _, token := range tokens { - if token.id < 0 || token.id >= cfg.VocabSize { + if token.id <= 0 || token.id >= cfg.VocabSize { return fmt.Errorf("invalid Gemma4 %s %d for vocab size %d", token.name, token.id, cfg.VocabSize) } if other, ok := seen[token.id]; ok { @@ -683,6 +702,19 @@ func newModel(root *model.Root) (base.Model, error) { if err != nil { return nil, err } + audioConfig, err := parseAudioConfig(configData) + if err != nil { + return nil, err + } + var audioProcessorConfig *AudioProcessorConfig + if audioConfig != nil { + if processorData, readErr := root.Manifest.ReadConfig("processor_config.json"); readErr == nil { + audioProcessorConfig, err = parseAudioProcessorConfig(processorData) + if err != nil { + return nil, err + } + } + } if qt := root.QuantType(); qt != "" { cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt) @@ -717,12 +749,14 @@ func newModel(root *model.Root) (base.Model, error) { } m := &Model{ - Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), - TextConfig: &cfg, - VisionConfig: visionConfig, - tok: tok, - mediaTokens: mediaTokens, - SuppressLogitBias: makeSuppressLogitBias(suppressTokens, cfg.VocabSize), + Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), + TextConfig: &cfg, + VisionConfig: visionConfig, + AudioConfig: audioConfig, + AudioProcessorConfig: audioProcessorConfig, + tok: tok, + mediaTokens: mediaTokens, + SuppressLogitBias: makeSuppressLogitBias(suppressTokens, cfg.VocabSize), } for i := range m.Layers { @@ -796,6 +830,18 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { } m.EmbedVision = embedVision } + if m.AudioConfig != nil && m.AudioProcessorConfig != nil && hasCompleteGemma4AudioWeights(tensors, m.AudioConfig, m.HiddenSize) { + audio, err := loadAudioModel(tensors, m.AudioConfig, m.HiddenSize, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + embedAudio, err := loadMultimodalEmbedder(tensors, "embed_audio", m.AudioConfig.RMSNormEps, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + m.Audio = audio + m.EmbedAudio = embedAudio + } // PLE model-level weights. if m.HiddenSizePerLayer > 0 { diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go index f8699537c13..d020fc68d4f 100644 --- a/x/models/gemma4/vision.go +++ b/x/models/gemma4/vision.go @@ -33,6 +33,9 @@ const ( defaultGemma4BOIToken = "<|image>" defaultGemma4ImageToken = "<|image|>" defaultGemma4EOIToken = "" + defaultGemma4BOAToken = "<|audio>" + defaultGemma4AudioToken = "<|audio|>" + defaultGemma4EOAToken = "" maxGemma4ImageBytes = 32 << 20 maxGemma4ImageDimension = 16_384 @@ -86,6 +89,9 @@ type gemma4MediaTokens struct { BOI string Image string EOI string + BOA string + Audio string + EOA string } type gemma4ImageInput struct { @@ -104,9 +110,12 @@ type gemma4MediaLayout struct { } type gemma4MediaPayload struct { - Image gemma4ImageInput + Image *gemma4ImageInput + Audio *gemma4AudioInput ImageStart int ImageEnd int + AudioStart int + AudioEnd int } type ClippableLinear struct { @@ -250,6 +259,9 @@ func defaultGemma4MediaTokens() gemma4MediaTokens { BOI: defaultGemma4BOIToken, Image: defaultGemma4ImageToken, EOI: defaultGemma4EOIToken, + BOA: defaultGemma4BOAToken, + Audio: defaultGemma4AudioToken, + EOA: defaultGemma4EOAToken, } } @@ -258,6 +270,9 @@ func parseGemma4MediaTokens(data []byte, fallback gemma4MediaTokens) gemma4Media BOIToken string `json:"boi_token"` ImageToken string `json:"image_token"` EOIToken string `json:"eoi_token"` + BOAToken string `json:"boa_token"` + AudioToken string `json:"audio_token"` + EOAToken string `json:"eoa_token"` } if err := json.Unmarshal(data, &cfg); err != nil { return fallback @@ -271,6 +286,15 @@ func parseGemma4MediaTokens(data []byte, fallback gemma4MediaTokens) gemma4Media if cfg.EOIToken != "" { fallback.EOI = cfg.EOIToken } + if cfg.BOAToken != "" { + fallback.BOA = cfg.BOAToken + } + if cfg.AudioToken != "" { + fallback.Audio = cfg.AudioToken + } + if cfg.EOAToken != "" { + fallback.EOA = cfg.EOAToken + } return fallback } @@ -505,8 +529,8 @@ func (m *MultimodalEmbedder) Forward(x *mlx.Array) *mlx.Array { return m.Projection.Forward(mlx.RMSNormFn(x, nil, m.Eps)) } -// PrepareMedia implements the runner's media contract. Each image is expanded -// in stream order and remains a separate cache-identity item. +// PrepareMedia implements the runner's ordered media contract. Each image or +// audio segment remains a separate cache-identity item. func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*base.PreparedRequest, error) { prepared := &base.PreparedRequest{} var layout gemma4MediaLayout @@ -518,46 +542,66 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas prepared.Tokens = append(prepared.Tokens, seg.Tokens...) continue } - if m.VisionConfig == nil || (m.Vision == nil && m.UnifiedVision == nil) || m.EmbedVision == nil { - return nil, fmt.Errorf("this model does not support %s input", seg.Kind) - } - if seg.Kind != "image" { - return nil, fmt.Errorf("gemma4 does not support %s input", seg.Kind) - } - - img, err := preprocessGemma4Image(ctx, seg.Data, m.VisionConfig, int(m.VisionSoftTokens)) - if err != nil { - return nil, err - } start := len(prepared.Tokens) - prepared.Tokens = append(prepared.Tokens, m.BOITokenIDValue) - imageStart := len(prepared.Tokens) - start - for range img.SoftTokens { - prepared.Tokens = append(prepared.Tokens, m.ImageTokenIDValue) - } - imageEnd := len(prepared.Tokens) - start - prepared.Tokens = append(prepared.Tokens, m.EOITokenIDValue) - - geom := *img - mediaData := geom.Pixels - dims := []int{1, 3, geom.Height, geom.Width} - geom.Pixels = nil - if m.UnifiedVision != nil { - mediaData = geom.Patches - dims = []int{1, geom.SoftTokens, int(m.UnifiedVision.PatchDim)} - geom.Patches = nil - layout.ImageSpans = append(layout.ImageSpans, [2]int{start + imageStart, start + imageEnd}) + var payload gemma4MediaPayload + var mediaData []float32 + var dims []int + switch seg.Kind { + case "image": + if m.VisionConfig == nil || (m.Vision == nil && m.UnifiedVision == nil) || m.EmbedVision == nil { + return nil, fmt.Errorf("this model does not support image input") + } + img, err := preprocessGemma4Image(ctx, seg.Data, m.VisionConfig, int(m.VisionSoftTokens)) + if err != nil { + return nil, err + } + prepared.Tokens = append(prepared.Tokens, m.BOITokenIDValue) + payload.ImageStart = len(prepared.Tokens) - start + for range img.SoftTokens { + prepared.Tokens = append(prepared.Tokens, m.ImageTokenIDValue) + } + payload.ImageEnd = len(prepared.Tokens) - start + prepared.Tokens = append(prepared.Tokens, m.EOITokenIDValue) + geom := *img + mediaData = geom.Pixels + dims = []int{1, 3, geom.Height, geom.Width} + geom.Pixels = nil + if m.UnifiedVision != nil { + mediaData = geom.Patches + dims = []int{1, geom.SoftTokens, int(m.UnifiedVision.PatchDim)} + geom.Patches = nil + layout.ImageSpans = append(layout.ImageSpans, [2]int{start + payload.ImageStart, start + payload.ImageEnd}) + } + payload.Image = &geom + case "audio": + if m.AudioConfig == nil || m.AudioProcessorConfig == nil || m.Audio == nil || m.EmbedAudio == nil { + return nil, fmt.Errorf("this model does not support audio input") + } + audio, err := preprocessGemma4Audio(ctx, seg.Data, m.AudioProcessorConfig) + if err != nil { + return nil, err + } + prepared.Tokens = append(prepared.Tokens, m.BOATokenIDValue) + payload.AudioStart = len(prepared.Tokens) - start + for range audio.SoftTokens { + prepared.Tokens = append(prepared.Tokens, m.AudioTokenIDValue) + } + payload.AudioEnd = len(prepared.Tokens) - start + prepared.Tokens = append(prepared.Tokens, m.EOATokenIDValue) + geom := *audio + mediaData = geom.Features + dims = []int{1, geom.Frames, 128} + geom.Features = nil + payload.Audio = &geom + default: + return nil, fmt.Errorf("gemma4 does not support %s input", seg.Kind) } item := base.PreparedItem{ Range: [2]int{start, len(prepared.Tokens)}, Source: source, MediaData: mediaData, Dims: dims, - Opaque: gemma4MediaPayload{ - Image: geom, - ImageStart: imageStart, - ImageEnd: imageEnd, - }, + Opaque: payload, } if err := ctx.Err(); err != nil { return nil, err @@ -573,24 +617,31 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas return prepared, nil } -// EncodeMedia builds the lazy image feature graph. The runner owns and frees -// MediaData, so pixels are always read from data rather than Opaque. +// EncodeMedia builds one lazy image or audio feature graph from runner-owned +// MediaData. func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { payload := item.Opaque.(gemma4MediaPayload) + if payload.Audio != nil { + features := m.EmbedAudio.Forward(m.Audio.Forward(data, payload.Audio)) + return mlx.Squeeze(features, 0) + } var encoded *mlx.Array if m.UnifiedVision != nil { patches := mlx.Reshape(data, 1, int32(payload.Image.SoftTokens), m.UnifiedVision.PatchDim) - encoded = m.UnifiedVision.Forward(patches, &payload.Image) + encoded = m.UnifiedVision.Forward(patches, payload.Image) } else { pixels := mlx.Reshape(data, 1, 3, int32(payload.Image.Height), int32(payload.Image.Width)) - encoded = m.Vision.Forward(pixels, &payload.Image) + encoded = m.Vision.Forward(pixels, payload.Image) } features := m.EmbedVision.Forward(encoded) return mlx.Squeeze(features, 0) } -func gemma4ImageRun(item batch.MediaItem) (start, end int) { +func gemma4MediaRun(item batch.MediaItem) (start, end int) { payload := item.Opaque.(gemma4MediaPayload) + if payload.Audio != nil { + return item.Pos + payload.AudioStart, item.Pos + payload.AudioEnd + } return item.Pos + payload.ImageStart, item.Pos + payload.ImageEnd } @@ -599,7 +650,7 @@ func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch) *mlx.Array { if item.Features == nil { continue } - start, end := gemma4ImageRun(item) + start, end := gemma4MediaRun(item) base := int(b.SeqOffsets[item.Seq]) lo := max(start, base) hi := min(end, base+int(b.SeqQueryLens[item.Seq])) @@ -613,11 +664,11 @@ func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch) *mlx.Array { return h } -// pleTokens masks feature-bearing image tokens from PLE exactly where the +// pleTokens masks feature-bearing media tokens from PLE exactly where the // same prepared items are scattered into the token embeddings. func gemma4PLETokens(tokens *mlx.Array, b *batch.Batch) *mlx.Array { for _, item := range b.Media { - start, end := gemma4ImageRun(item) + start, end := gemma4MediaRun(item) base := int(b.SeqOffsets[item.Seq]) lo := max(start, base) hi := min(end, base+int(b.SeqQueryLens[item.Seq])) @@ -961,6 +1012,28 @@ func imageToCHWFloat32Context(ctx context.Context, img image.Image) ([]float32, return out, nil } +func mediaTokenSpan(tokens []int32, mediaTokenID int32) (int, int, error) { + start, end := -1, -1 + for i, tok := range tokens { + if tok != mediaTokenID { + continue + } + if start == -1 { + start = i + } + end = i + 1 + } + if start == -1 { + return 0, 0, fmt.Errorf("Gemma4 prompt contains no media token id %d", mediaTokenID) + } + for i := start; i < end; i++ { + if tokens[i] != mediaTokenID { + return 0, 0, errors.New("Gemma4 media tokens are not contiguous") + } + } + return start, end, nil +} + func (m *VisionModel) Forward(pixels *mlx.Array, img *gemma4ImageInput) *mlx.Array { positions := m.positionArrays(img) h := m.PatchEmbedder.Forward(pixels, positions) diff --git a/x/models/gemma4/vision_test.go b/x/models/gemma4/vision_test.go index c237c8745db..edd0026ff6c 100644 --- a/x/models/gemma4/vision_test.go +++ b/x/models/gemma4/vision_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/mlxrunner/model/base" ) @@ -340,7 +341,7 @@ func TestPrepareMediaPreservesOrderedImageItems(t *testing.T) { t.Fatalf("item %d payload = start %d end %d pixels %d", i, payload.ImageStart, payload.ImageEnd, len(payload.Image.Pixels)) } media := batch.MediaItem{Pos: item.Range[0], Opaque: item.Opaque} - start, end := gemma4ImageRun(media) + start, end := gemma4MediaRun(media) if start != item.Range[0]+1 || end != item.Range[1]-1 { t.Fatalf("item %d scatter/PLE run = [%d,%d), want current image span [%d,%d)", i, start, end, item.Range[0]+1, item.Range[1]-1) } @@ -589,3 +590,114 @@ func TestPrepareMediaCancellationDuringProductionStages(t *testing.T) { }) } } + +func TestPrepareMediaAudioUsesFeatureData(t *testing.T) { + frames := make([][]float64, 1600) + for i := range frames { + frames[i] = []float64{0.1 * math.Sin(2*math.Pi*440*float64(i)/16000)} + } + processor := defaultAudioProcessorConfig() + m := &Model{ + TextConfig: &TextConfig{ + AudioTokenIDValue: 20, BOATokenIDValue: 21, EOATokenIDValue: 22, + }, + AudioConfig: &AudioConfig{}, + AudioProcessorConfig: &processor, + Audio: &AudioModel{}, + EmbedAudio: &MultimodalEmbedder{}, + } + got, err := m.PrepareMedia(context.Background(), []base.Segment{ + {Tokens: []int32{1}}, + {Kind: "audio", Data: makeTestWAV(t, 1, 16, 16000, frames)}, + {Tokens: []int32{2}}, + }) + if err != nil { + t.Fatalf("PrepareMedia() error = %v", err) + } + if len(got.Items) != 1 { + t.Fatalf("audio item count = %d, want 1", len(got.Items)) + } + item := got.Items[0] + payload := item.Opaque.(gemma4MediaPayload) + if payload.Audio == nil || payload.Audio.Features != nil || item.Dims[2] != 128 { + t.Fatalf("audio payload/dims = %#v/%v", payload.Audio, item.Dims) + } + wantRange := [2]int{1, len(got.Tokens) - 1} + if item.Range != wantRange || got.Tokens[1] != 21 || got.Tokens[len(got.Tokens)-2] != 22 { + t.Fatalf("audio range/tokens = %v/%v", item.Range, got.Tokens) + } + if got.Tokens[payload.AudioStart+item.Range[0]] != 20 { + t.Fatalf("audio feature token missing: %v", got.Tokens) + } + firstTokens := slices.Clone(got.Tokens) + firstFeatures := slices.Clone(item.MediaData) + firstMask := slices.Clone(payload.Audio.FeatureMask) + + secondFrames := make([][]float64, 3200) + for i := range secondFrames { + secondFrames[i] = []float64{0.1 * math.Sin(2*math.Pi*880*float64(i)/16000)} + } + second, err := m.PrepareMedia(context.Background(), []base.Segment{ + {Tokens: []int32{3, 4}}, + {Kind: "audio", Data: makeTestWAV(t, 1, 16, 16000, secondFrames)}, + }) + if err != nil { + t.Fatal(err) + } + if len(second.Items) != 1 || len(second.Items[0].MediaData) == len(item.MediaData) { + t.Fatalf("second audio items/features = %d/%d, first features %d", len(second.Items), len(second.Items[0].MediaData), len(item.MediaData)) + } + secondPayload := second.Items[0].Opaque.(gemma4MediaPayload) + if secondPayload.Audio == nil || secondPayload.Audio.Features != nil { + t.Fatalf("second audio payload = %#v", secondPayload.Audio) + } + if &second.Items[0].MediaData[0] == &item.MediaData[0] || &secondPayload.Audio.FeatureMask[0] == &payload.Audio.FeatureMask[0] { + t.Fatal("sequential audio prepares reused feature or mask storage") + } + second.Tokens[0] = 99 + second.Items[0].MediaData[0] = 99 + secondPayload.Audio.FeatureMask[0] = !secondPayload.Audio.FeatureMask[0] + if !slices.Equal(got.Tokens, firstTokens) || !slices.Equal(item.MediaData, firstFeatures) || !slices.Equal(payload.Audio.FeatureMask, firstMask) { + t.Fatal("mutating the later audio request changed the earlier request") + } + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if prepared, err := m.PrepareMedia(canceled, []base.Segment{{Kind: "audio", Data: makeTestWAV(t, 1, 16, 16000, frames)}}); !errors.Is(err, context.Canceled) || prepared != nil { + t.Fatalf("canceled audio prepare = (%v, %v), want (nil, context.Canceled)", prepared, err) + } +} + +func TestScatterMediaAudioSpanAndDType(t *testing.T) { + useMLXTestThread(t) + + m := &Model{TextConfig: &TextConfig{HiddenSize: 2}} + payload := gemma4MediaPayload{Audio: &gemma4AudioInput{}, AudioStart: 1, AudioEnd: 3} + features := mlx.FromValues([]float32{1, 2, 3, 4}, 2, 2) + hidden := mlx.Zeros(mlx.DTypeBFloat16, 1, 5, 2) + got := m.scatterMedia(hidden, &batch.Batch{ + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{5}, + Media: []batch.MediaItem{{Pos: 1, Features: features, Opaque: payload}}, + }) + if got.DType() != mlx.DTypeBFloat16 || !slices.Equal(got.Dims(), []int{1, 5, 2}) { + t.Fatalf("scattered dtype/shape = %s/%v, want bfloat16/[1 5 2]", got.DType(), got.Dims()) + } + gotFloat := got.AsType(mlx.DTypeFloat32) + mlx.Eval(gotFloat) + if values := gotFloat.Floats(); !equalFloat32s(values, []float32{0, 0, 0, 0, 1, 2, 3, 4, 0, 0}) { + t.Fatalf("scattered values = %v", values) + } + + partialHidden := mlx.Zeros(mlx.DTypeFloat16, 1, 2, 2) + partial := m.scatterMedia(partialHidden, &batch.Batch{ + SeqOffsets: []int32{3}, + SeqQueryLens: []int32{2}, + Media: []batch.MediaItem{{Pos: 1, Features: features, Opaque: payload}}, + }) + partialFloat := partial.AsType(mlx.DTypeFloat32) + mlx.Eval(partialFloat) + if partial.DType() != mlx.DTypeFloat16 || !equalFloat32s(partialFloat.Floats(), []float32{3, 4, 0, 0}) { + t.Fatalf("partial scattered dtype/values = %s/%v", partial.DType(), partialFloat.Floats()) + } +} From b96bd2442934f8578c51aa116cb86def174389c2 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 04:03:47 +0000 Subject: [PATCH 49/58] gemma4: add MLX audio forward parity test Add a deterministic end-to-end reference for the released audio tower and projector. Co-authored-by: Codex --- .../gemma4/audio_forward_reference_test.go | 144 ++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 x/models/gemma4/audio_forward_reference_test.go diff --git a/x/models/gemma4/audio_forward_reference_test.go b/x/models/gemma4/audio_forward_reference_test.go new file mode 100644 index 00000000000..599fc092557 --- /dev/null +++ b/x/models/gemma4/audio_forward_reference_test.go @@ -0,0 +1,144 @@ +package gemma4 + +import ( + "context" + "math" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" +) + +func TestAudioForwardReference(t *testing.T) { + modelDir := os.Getenv("GEMMA4_AUDIO_MODEL_DIR") + refDir := os.Getenv("GEMMA4_AUDIO_REF_DIR") + wavPath := os.Getenv("GEMMA4_AUDIO_WAV") + if modelDir == "" || refDir == "" || wavPath == "" { + t.Skip("set GEMMA4_AUDIO_MODEL_DIR, GEMMA4_AUDIO_REF_DIR, and GEMMA4_AUDIO_WAV for audio parity") + } + skipIfNoMLX(t) + runtime.LockOSThread() + defer runtime.UnlockOSThread() + if mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + + configData, err := os.ReadFile(filepath.Join(modelDir, "config.json")) + if err != nil { + t.Fatal(err) + } + audioConfig, err := parseAudioConfig(configData) + if err != nil || audioConfig == nil { + t.Fatalf("parse audio config: %v", err) + } + textConfig, err := parseTextConfig(configData) + if err != nil { + t.Fatal(err) + } + processorData, err := os.ReadFile(filepath.Join(modelDir, "processor_config.json")) + if err != nil { + t.Fatal(err) + } + processorConfig, err := parseAudioProcessorConfig(processorData) + if err != nil { + t.Fatal(err) + } + wav, err := os.ReadFile(wavPath) + if err != nil { + t.Fatal(err) + } + input, err := preprocessGemma4Audio(context.Background(), wav, processorConfig) + if err != nil { + t.Fatal(err) + } + + source, err := mlx.LoadSafetensorsNative(filepath.Join(modelDir, "model.safetensors")) + if err != nil { + t.Fatal(err) + } + defer source.Free() + reference, err := mlx.LoadSafetensorsNative(filepath.Join(refDir, "audio-checkpoints.safetensors")) + if err != nil { + t.Fatal(err) + } + defer reference.Free() + if wantFrames := reference.Get("input_features").Dim(1); wantFrames > input.Frames { + padding := wantFrames - input.Frames + input.Features = append(input.Features, make([]float32, padding*128)...) + input.FeatureMask = append(input.FeatureMask, make([]bool, padding)...) + input.Frames = wantFrames + } + + required, err := gemma4metadata.RequiredAudioTensorShapes(audioMetadataConfig(audioConfig, textConfig.HiddenSize)) + if err != nil { + t.Fatal(err) + } + tensors := make(map[string]*mlx.Array, len(required)) + for name := range required { + tensors[name] = source.Get(name) + if tensors[name] == nil { + t.Fatalf("source is missing %s", name) + } + } + audioModel, err := loadAudioModel(tensors, audioConfig, textConfig.HiddenSize, 0, 0, "", nil) + if err != nil { + t.Fatal(err) + } + embedAudio, err := loadMultimodalEmbedder(tensors, "embed_audio", audioConfig.RMSNormEps, 0, 0, "", nil) + if err != nil { + t.Fatal(err) + } + + features := mlx.FromValues(input.Features, 1, input.Frames, 128) + compareAudioReference(t, "input_features", features, reference.Get("input_features"), 1e-5, 1e-5) + valid := append([]bool(nil), input.FeatureMask...) + x := mlx.ExpandDims(features, -1) + x, valid = audioModel.Conv0.Forward(x, valid) + x, valid = audioModel.Conv1.Forward(x, valid) + x = mlx.Reshape(x, 1, int32(x.Dim(1)), int32(x.Dim(2)*x.Dim(3))) + x = audioModel.InputProj.Forward(x) + compareAudioReference(t, "subsample", x, reference.Get("subsample"), 0.5, 0.05) + for index, layer := range audioModel.Layers { + x = layer.Forward(x, valid) + if index == 0 || index == len(audioModel.Layers)-1 { + name := "layer_0" + if index != 0 { + name = "layer_11" + } + compareAudioReference(t, name, x, reference.Get(name), 0.5, 0.05) + } + } + x = audioModel.OutputProj.Forward(x) + compareAudioReference(t, "output_projection", x, reference.Get("output_projection"), 0.5, 0.05) + x = embedAudio.Forward(x) + compareAudioReference(t, "multimodal_projection", x, reference.Get("multimodal_projection"), 0.5, 0.05) +} + +func compareAudioReference(t *testing.T, name string, got, want *mlx.Array, atol, rtol float64) { + t.Helper() + if got == nil || want == nil { + t.Fatalf("%s tensor is missing", name) + } + if !equalIntShape(got.Dims(), want.Dims()) { + t.Fatalf("%s shape = %v, want %v", name, got.Dims(), want.Dims()) + } + got = got.AsType(mlx.DTypeFloat32) + want = want.AsType(mlx.DTypeFloat32) + mlx.Eval(got, want) + gotValues, wantValues := got.Floats(), want.Floats() + var maxDiff float64 + for i := range wantValues { + diff := math.Abs(float64(gotValues[i] - wantValues[i])) + if diff > maxDiff { + maxDiff = diff + } + tolerance := atol + rtol*math.Abs(float64(wantValues[i])) + if diff > tolerance { + t.Fatalf("%s[%d] = %g, want %g (diff %g, tolerance %g)", name, i, gotValues[i], wantValues[i], diff, tolerance) + } + } + t.Logf("%s matched %d values; max absolute difference %g", name, len(wantValues), maxDiff) +} From e80351eb1a700fe6d189c7a747847aa1b836fc01 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 04:20:15 +0000 Subject: [PATCH 50/58] gemma4: harden MLX audio readiness and validation Tighten audio configuration, source inventory, manifest, preprocessing, and forward-reference validation while exercising lazy media encoding and the two-result model Forward contract. Co-authored-by: Codex --- server/images.go | 10 ++ server/images_test.go | 20 ++- server/model_list_cache.go | 2 +- server/model_list_cache_test.go | 31 ++++ x/create/client/create.go | 6 +- x/create/client/create_test.go | 20 ++- .../gemma4/audio_forward_reference_test.go | 134 +++++++++++++++++- x/models/gemma4/audio_processor_test.go | 23 +++ x/models/gemma4/audio_test.go | 42 ++++++ x/models/gemma4/gemma4.go | 23 ++- x/models/gemma4/metadata/audio.go | 83 +++++++++-- x/models/gemma4/metadata/audio_test.go | 33 ++++- x/models/gemma4/metadata/vision.go | 2 + 13 files changed, 404 insertions(+), 25 deletions(-) diff --git a/server/images.go b/server/images.go index d968312e5da..c7017bb000e 100644 --- a/server/images.go +++ b/server/images.go @@ -89,6 +89,7 @@ type Model struct { Gemma4VisionTensors map[string]gemma4metadata.TensorDescriptor `json:"-"` Gemma4AudioConfig *gemma4metadata.ConfigFile `json:"-"` Gemma4AudioTensors map[string]gemma4metadata.TensorDescriptor `json:"-"` + Gemma4AudioReady bool `json:"-"` System string License []string Digest string @@ -521,6 +522,7 @@ func suppressAudioCapability(m *Model, arch string) bool { return true } return m.Gemma4AudioConfig == nil || + !m.Gemma4AudioReady || gemma4metadata.ValidateAudioInstalledInventory(*m.Gemma4AudioConfig, m.Gemma4AudioTensors) != nil } if m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" { @@ -569,6 +571,13 @@ func hasGemma4AudioTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifest return err == nil && gemma4metadata.ValidateAudioInstalledInventory(cfg, tensors) == nil } +func hasGemma4AudioRuntimeMetadata(cfg gemma4metadata.ConfigFile, mf *manifest.Manifest) bool { + var processorData, tokenizerConfigData json.RawMessage + return mf.ReadConfigJSON("processor_config.json", &processorData) == nil && + mf.ReadConfigJSON("tokenizer_config.json", &tokenizerConfigData) == nil && + gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData) == nil +} + func gemma4AudioTensorDescriptors(layers []manifest.Layer) (map[string]gemma4metadata.TensorDescriptor, error) { tensors := make(map[string]gemma4metadata.TensorDescriptor) descriptorWork := 0 @@ -1006,6 +1015,7 @@ func GetModel(name string) (*Model, error) { if err := mf.ReadConfigJSON("config.json", &cfg); err == nil { m.Gemma4VisionConfig = &cfg m.Gemma4AudioConfig = &cfg + m.Gemma4AudioReady = hasGemma4AudioRuntimeMetadata(cfg, mf) if tensors, err := gemma4VisionTensorDescriptors(mf.Layers); err == nil { m.Gemma4VisionTensors = tensors } diff --git a/server/images_test.go b/server/images_test.go index 62bfe651024..fccf83749d6 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -609,6 +609,7 @@ func TestModelCapabilities(t *testing.T) { TensorLayerNames: gemma4AudioTensorNames(1), Gemma4AudioConfig: gemma4AudioConfig(1), Gemma4AudioTensors: testGemma4AudioTensorDescriptors(1), + Gemma4AudioReady: true, Template: chatTemplate, }, expectedCaps: []model.Capability{model.CapabilityAudio}, @@ -718,6 +719,7 @@ func TestGemma4SafetensorsAudioCapabilityRequiresCompleteInventory(t *testing.T) TensorLayerNames: gemma4AudioTensorNames(1), Gemma4AudioConfig: gemma4AudioConfig(1), Gemma4AudioTensors: testGemma4AudioTensorDescriptors(1), + Gemma4AudioReady: true, } if !slices.Contains(complete.Capabilities(), model.CapabilityAudio) { t.Fatal("complete audio inventory did not expose audio") @@ -776,6 +778,12 @@ func TestGemma4InstalledAudioCapabilityDescriptorAndPayloadMatrix(t *testing.T) {name: "partial", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == target }) }}, + {name: "missing processor metadata", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == "processor_config.json" }) + }}, + {name: "missing tokenizer metadata", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == "tokenizer_config.json" }) + }}, {name: "near-match internal name", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { return replaceGemma4AudioFixtureLayer(t, layers, target, target+".extra", gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{3, 4}}, nil) }}, @@ -884,7 +892,7 @@ func TestGemma4InstalledAudioCapabilityRejectsUnboundedConfig(t *testing.T) { func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { t.Helper() descriptors := testGemma4AudioTensorDescriptors(1) - layers := make([]manifest.Layer, 0, len(descriptors)+1) + layers := make([]manifest.Layer, 0, len(descriptors)+3) for name, descriptor := range descriptors { layers = append(layers, gemma4AudioFixtureLayer(t, name, name, descriptor, nil)) } @@ -894,6 +902,13 @@ func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { } digest := createTestBlob(t, config) layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(config)), Name: "config.json"}) + for name, data := range map[string][]byte{ + "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), + "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), + } { + digest := createTestBlob(t, data) + layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(data)), Name: name}) + } return layers } @@ -1177,7 +1192,8 @@ func gemma4VisionConfig(layers int) *gemma4metadata.ConfigFile { func gemma4AudioConfig(layers int) *gemma4metadata.ConfigFile { return &gemma4metadata.ConfigFile{ - TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, + AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 70c8d46b0ca..0dbaa691052 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -400,7 +400,7 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum return c == model.CapabilityVision }) } - if configErr != nil || !hasGemma4AudioTensorLayers(gemma4cfg, mf.Layers) { + if configErr != nil || !hasGemma4AudioRuntimeMetadata(gemma4cfg, mf) || !hasGemma4AudioTensorLayers(gemma4cfg, mf.Layers) { summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { return c == model.CapabilityAudio }) diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index 465d527dfdf..a2207377198 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -178,6 +178,37 @@ func TestModelListSummaryGemma4SafetensorsVisionRequiresTensorLayers(t *testing. } } +func TestModelListSummaryGemma4AudioRequiresRuntimeMetadataAndTensors(t *testing.T) { + setTestHome(t, t.TempDir()) + cfg := model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererSmall, Capabilities: []string{"completion", "audio"}} + for _, tt := range []struct { + name string + layers []manifest.Layer + wantAudio bool + }{ + {"complete", gemma4AudioManifestLayers(t), true}, + {"partial", slices.DeleteFunc(gemma4AudioManifestLayers(t), func(layer manifest.Layer) bool { + return layer.Name == "model.audio_tower.layers.0.self_attn.q_proj.input_max" + }), false}, + } { + t.Run(tt.name, func(t *testing.T) { + modelName := "list-gemma4-audio-" + tt.name + createSafetensorsTestModel(t, modelName, cfg, tt.layers) + mf, err := manifest.ParseNamedManifest(model.ParseName(modelName)) + if err != nil { + t.Fatal(err) + } + summary, err := buildModelListSummary(model.ParseName(modelName), mf) + if err != nil { + t.Fatal(err) + } + if got := slices.Contains(summary.Capabilities, model.CapabilityAudio); got != tt.wantAudio { + t.Fatalf("audio capability = %v, want %v (%v)", got, tt.wantAudio, summary.Capabilities) + } + }) + } +} + func TestModelListCacheMutationHooks(t *testing.T) { gin.SetMode(gin.TestMode) setTestHome(t, t.TempDir()) diff --git a/x/create/client/create.go b/x/create/client/create.go index 0cc93bacf00..7fc7bf92594 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -642,8 +642,12 @@ func gemma4ModelDirMediaCapabilities(modelDir string) (vision, audio bool) { // descriptor shapes rather than the released tower sentinel alone. tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} } - return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil, + processorData, processorErr := os.ReadFile(filepath.Join(modelDir, "processor_config.json")) + tokenizerConfigData, tokenizerErr := os.ReadFile(filepath.Join(modelDir, "tokenizer_config.json")) + audioReady := processorErr == nil && tokenizerErr == nil && + gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData) == nil && gemma4metadata.ValidateAudioSourceInventory(cfg, tensors) == nil + return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil, audioReady } // readChatTemplate returns the model's chat template, preferring the diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index bbcf9e164aa..d6bc444d98f 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -669,7 +669,8 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { cfg := gemma4metadata.ConfigFile{ Architectures: []string{identity.architecture}, ModelType: identity.modelType, - TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, + AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, @@ -697,6 +698,7 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { t.Fatal(err) } writeClientSafetensorDescriptors(t, dir, tensors) + writeGemma4AudioRuntimeConfigs(t, dir) got := inferSafetensorsCapabilities(dir, "") if slices.Contains(got, "audio") != wantAudio { t.Fatalf("capabilities = %v, want audio %t", got, wantAudio) @@ -730,7 +732,8 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testing.T) { base := gemma4metadata.ConfigFile{ Architectures: []string{"Gemma4ForConditionalGeneration"}, ModelType: "gemma4", - TextConfig: gemma4metadata.TextConfig{HiddenSize: 5}, + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, + AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, @@ -780,6 +783,19 @@ func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testin } } +func writeGemma4AudioRuntimeConfigs(t *testing.T, dir string) { + t.Helper() + files := map[string][]byte{ + "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), + "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), + } + for name, data := range files { + if err := os.WriteFile(filepath.Join(dir, name), data, 0o644); err != nil { + t.Fatal(err) + } + } +} + func TestInferSafetensorsCapabilitiesGemma4PackedSourceRequiresProducerContract(t *testing.T) { const configJSON = `{ "architectures":["Gemma4ForConditionalGeneration"],"model_type":"gemma4", diff --git a/x/models/gemma4/audio_forward_reference_test.go b/x/models/gemma4/audio_forward_reference_test.go index 599fc092557..40a59356bc5 100644 --- a/x/models/gemma4/audio_forward_reference_test.go +++ b/x/models/gemma4/audio_forward_reference_test.go @@ -6,9 +6,12 @@ import ( "os" "path/filepath" "runtime" + "strings" "testing" + "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/mlx" + runnermodel "github.com/ollama/ollama/x/mlxrunner/model" gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" ) @@ -93,11 +96,22 @@ func TestAudioForwardReference(t *testing.T) { } features := mlx.FromValues(input.Features, 1, input.Frames, 128) - compareAudioReference(t, "input_features", features, reference.Get("input_features"), 1e-5, 1e-5) + compareAudioReference(t, "input_features", features, reference.Get("input_features"), 2e-5, 1e-5) + compareAudioMaskReference(t, "input_features_mask", input.FeatureMask, reference.Get("input_features_mask"), false) valid := append([]bool(nil), input.FeatureMask...) x := mlx.ExpandDims(features, -1) x, valid = audioModel.Conv0.Forward(x, valid) x, valid = audioModel.Conv1.Forward(x, valid) + compareAudioMaskReference(t, "audio_forward_mask", valid, reference.Get("audio_forward_mask"), true) + validOutputs := 0 + for _, value := range valid { + if value { + validOutputs++ + } + } + if validOutputs != input.SoftTokens { + t.Fatalf("valid audio outputs = %d, want %d", validOutputs, input.SoftTokens) + } x = mlx.Reshape(x, 1, int32(x.Dim(1)), int32(x.Dim(2)*x.Dim(3))) x = audioModel.InputProj.Forward(x) compareAudioReference(t, "subsample", x, reference.Get("subsample"), 0.5, 0.05) @@ -115,6 +129,124 @@ func TestAudioForwardReference(t *testing.T) { compareAudioReference(t, "output_projection", x, reference.Get("output_projection"), 0.5, 0.05) x = embedAudio.Forward(x) compareAudioReference(t, "multimodal_projection", x, reference.Get("multimodal_projection"), 0.5, 0.05) + + featureInput := mlx.FromValues(input.Features, 1, input.Frames, 128) + forward := audioModel.Forward(featureInput, input) + compareAudioReference(t, "audio_forward_output", forward, reference.Get("audio_forward_output"), 0.5, 0.05) + projectedForward := embedAudio.Forward(forward) + inputIDs := reference.Get("input_ids") + embedWeight := source.Get("model.language_model.embed_tokens.weight") + if inputIDs == nil || embedWeight == nil { + t.Fatal("reference input IDs or source token embeddings are missing") + } + inputIDs = inputIDs.AsType(mlx.DTypeInt32) + tokenEmbeddings := mlx.MulScalar(mlx.Take(embedWeight, inputIDs, 0), textConfig.EmbedScale) + mlx.Eval(inputIDs) + start, end, err := mediaTokenSpan(inputIDs.Ints(), textConfig.AudioTokenIDValue) + if err != nil { + t.Fatal(err) + } + if end-start != input.SoftTokens { + t.Fatalf("audio token span = %d, want %d", end-start, input.SoftTokens) + } + parts := make([]*mlx.Array, 0, 3) + if start > 0 { + parts = append(parts, mlx.SliceStartStop(tokenEmbeddings, []int32{0, 0, 0}, []int32{1, int32(start), textConfig.HiddenSize})) + } + parts = append(parts, projectedForward) + if end < inputIDs.Dim(1) { + parts = append(parts, mlx.SliceStartStop(tokenEmbeddings, []int32{0, int32(end), 0}, []int32{1, int32(inputIDs.Dim(1)), textConfig.HiddenSize})) + } + finalEmbeddings := mlx.Concatenate(parts, 1) + compareAudioReference(t, "final_input_embeddings", finalEmbeddings, reference.Get("final_input_embeddings"), 0.5, 0.05) + + if importedModel := os.Getenv("GEMMA4_AUDIO_MODEL_NAME"); importedModel != "" { + fullModel := loadImportedGemma4ReferenceModel(t, importedModel) + payload := gemma4MediaPayload{Audio: input, AudioStart: 1, AudioEnd: 1 + input.SoftTokens} + modelBatch := &batch.Batch{ + InputIDs: inputIDs, + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{int32(inputIDs.Dim(1))}, + Media: []batch.MediaItem{{ + Seq: 0, Pos: start - 1, Features: mlx.Squeeze(projectedForward, 0), Opaque: payload, + }}, + } + hidden, _ := fullModel.Forward(modelBatch, nil) + logits := fullModel.Unembed(hidden) + last := mlx.SliceStartStop(logits, []int32{0, int32(inputIDs.Dim(1) - 1), 0}, []int32{1, int32(inputIDs.Dim(1)), textConfig.VocabSize}) + wantLogits := reference.Get("prefill_logits") + compareAudioReference(t, "prefill_logits", last, wantLogits, 0.5, 0.05) + gotID := last.Argmax(-1, false).AsType(mlx.DTypeInt32) + wantID := wantLogits.Argmax(-1, false).AsType(mlx.DTypeInt32) + mlx.Eval(gotID, wantID) + if got, want := gotID.Int(), wantID.Int(); got != want { + t.Fatalf("prefill argmax = %d, want %d", got, want) + } + } +} + +func loadImportedGemma4ReferenceModel(t *testing.T, name string) *Model { + t.Helper() + root, err := runnermodel.Open(name) + if err != nil { + t.Fatalf("open imported model %q: %v", name, err) + } + defer root.Close() + bm, err := newModel(root) + if err != nil { + t.Fatalf("construct imported model: %v", err) + } + tensors := make(map[string]*mlx.Array) + seen := make(map[string]bool) + for _, layer := range root.Manifest.GetTensorLayers("") { + if seen[layer.Digest] { + continue + } + seen[layer.Digest] = true + for tensorName, array := range mlx.Load(root.Manifest.BlobPath(layer.Digest)) { + tensors[tensorName] = array + } + } + for tensorName, array := range tensors { + if strings.HasSuffix(tensorName, ".scale") { + tensors[strings.TrimSuffix(tensorName, ".scale")+"_scale"] = array + } + } + if err := bm.LoadWeights(tensors); err != nil { + t.Fatalf("load imported model weights: %v", err) + } + collected := mlx.Collect(bm) + for _, array := range collected { + mlx.Pin(array) + } + mlx.Eval(collected...) + m, ok := bm.(*Model) + if !ok { + t.Fatalf("imported model type = %T, want *gemma4.Model", bm) + } + return m +} + +func compareAudioMaskReference(t *testing.T, name string, got []bool, want *mlx.Array, invertWant bool) { + t.Helper() + if want == nil { + t.Fatalf("%s tensor is missing", name) + } + want = want.AsType(mlx.DTypeInt32) + mlx.Eval(want) + values := want.Ints() + if len(got) != len(values) { + t.Fatalf("%s length = %d, want %d", name, len(got), len(values)) + } + for i, value := range values { + expected := value != 0 + if invertWant { + expected = !expected + } + if got[i] != expected { + t.Fatalf("%s[%d] = %v, want %v", name, i, got[i], expected) + } + } } func compareAudioReference(t *testing.T, name string, got, want *mlx.Array, atol, rtol float64) { diff --git a/x/models/gemma4/audio_processor_test.go b/x/models/gemma4/audio_processor_test.go index 2b7aca358db..e3c561d9808 100644 --- a/x/models/gemma4/audio_processor_test.go +++ b/x/models/gemma4/audio_processor_test.go @@ -336,6 +336,29 @@ func TestGemma4AudioInputFailures(t *testing.T) { } } +func TestGemma4WAVValidationLimits(t *testing.T) { + if _, err := decodeGemma4WAV(context.Background(), make([]byte, maxGemma4AudioBytes+1), 16000); err == nil || !strings.Contains(err.Error(), "limit") { + t.Fatalf("oversized WAV error = %v", err) + } + tooManyChannels := makeTestWAV(t, 1, 16, 16000, [][]float64{make([]float64, maxGemma4AudioChannels+1)}) + if _, err := decodeGemma4WAV(context.Background(), tooManyChannels, 16000); err == nil || !strings.Contains(err.Error(), "channel count") { + t.Fatalf("channel-limit error = %v", err) + } + badRate := makeTestWAV(t, 1, 16, minGemma4AudioSampleRate-1, [][]float64{{0}}) + if _, err := decodeGemma4WAV(context.Background(), badRate, 16000); err == nil || !strings.Contains(err.Error(), "sample rate") { + t.Fatalf("sample-rate error = %v", err) + } + badAlignment := makeTestWAV(t, 1, 16, 16000, [][]float64{{0}}) + binary.LittleEndian.PutUint16(badAlignment[32:34], 1) + if _, err := decodeGemma4WAV(context.Background(), badAlignment, 16000); err == nil || !strings.Contains(err.Error(), "block alignment") { + t.Fatalf("block-alignment error = %v", err) + } + nonFinite := makeTestWAV(t, 3, 32, 16000, [][]float64{{math.NaN()}}) + if _, err := decodeGemma4WAV(context.Background(), nonFinite, 16000); err == nil || !strings.Contains(err.Error(), "non-finite") { + t.Fatalf("non-finite error = %v", err) + } +} + func TestGemma4AudioCancellationBeforeAllocations(t *testing.T) { cfg := defaultAudioProcessorConfig() frames := make([][]float64, 400) diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index b66f74df4ca..e822bf2fc8b 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -1,11 +1,13 @@ package gemma4 import ( + "context" "maps" "strings" "testing" "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model/base" gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" ) @@ -165,6 +167,46 @@ func TestSupportedGemma4AudioDType(t *testing.T) { } } +func TestPrepareAudioMedia(t *testing.T) { + cfg := defaultAudioProcessorConfig() + frames := make([][]float64, 16000) + for i := range frames { + frames[i] = []float64{0.25} + } + m := &Model{ + TextConfig: &TextConfig{AudioTokenIDValue: 1, BOATokenIDValue: 2, EOATokenIDValue: 3}, + AudioConfig: &AudioConfig{}, AudioProcessorConfig: &cfg, + Audio: &AudioModel{}, EmbedAudio: &MultimodalEmbedder{}, + } + wav := makeTestWAV(t, 1, 16, 16000, frames) + prepared, err := m.PrepareMedia(context.Background(), []base.Segment{{Tokens: []int32{9}}, {Kind: "audio", Data: wav}}) + if err != nil { + t.Fatal(err) + } + if len(prepared.Items) != 1 { + t.Fatalf("items = %d, want 1", len(prepared.Items)) + } + payload := prepared.Items[0].Opaque.(gemma4MediaPayload) + if got := payload.AudioEnd - payload.AudioStart; got != payload.Audio.SoftTokens { + t.Fatalf("audio span = %d, want %d", got, payload.Audio.SoftTokens) + } + if payload.Audio.Features != nil || len(prepared.Items[0].MediaData) == 0 { + t.Fatal("audio features must be owned by PreparedItem.MediaData") + } + ordered, err := m.PrepareMedia(context.Background(), []base.Segment{{Kind: "audio", Data: wav}, {Tokens: []int32{8}}, {Kind: "audio", Data: wav}}) + if err != nil || len(ordered.Items) != 2 || ordered.Items[0].Range[1] >= ordered.Items[1].Range[0] { + t.Fatalf("ordered audio = items %d error %v", len(ordered.Items), err) + } +} + +func TestPrepareAudioMediaRejectsMissingWeights(t *testing.T) { + cfg := defaultAudioProcessorConfig() + m := &Model{TextConfig: &TextConfig{}, AudioConfig: &AudioConfig{}, AudioProcessorConfig: &cfg} + if _, err := m.PrepareMedia(context.Background(), []base.Segment{{Kind: "audio", Data: []byte("wav")}}); err == nil || !strings.Contains(err.Error(), "does not support audio") { + t.Fatalf("PrepareMedia() error = %v", err) + } +} + func TestValidateAndLoadGemma4AudioWeights(t *testing.T) { useMLXTestThread(t) diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index 0f19533fbba..bfc3b4b3ac2 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -702,17 +702,14 @@ func newModel(root *model.Root) (base.Model, error) { if err != nil { return nil, err } - audioConfig, err := parseAudioConfig(configData) - if err != nil { - return nil, err + audioConfig, audioConfigErr := parseAudioConfig(configData) + if audioConfigErr != nil { + audioConfig = nil } var audioProcessorConfig *AudioProcessorConfig if audioConfig != nil { if processorData, readErr := root.Manifest.ReadConfig("processor_config.json"); readErr == nil { - audioProcessorConfig, err = parseAudioProcessorConfig(processorData) - if err != nil { - return nil, err - } + audioProcessorConfig, _ = parseAudioProcessorConfig(processorData) } } @@ -747,6 +744,11 @@ func newModel(root *model.Root) (base.Model, error) { if err != nil { return nil, fmt.Errorf("parse tokenizer: %w", err) } + if audioConfig != nil && audioProcessorConfig != nil && + !validGemma4AudioTokenizer(tok, mediaTokens, cfg.AudioTokenIDValue) { + audioConfig = nil + audioProcessorConfig = nil + } m := &Model{ Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), @@ -775,6 +777,13 @@ func newModel(root *model.Root) (base.Model, error) { return m, nil } +func validGemma4AudioTokenizer(tok *tokenizer.Tokenizer, tokens gemma4MediaTokens, audioTokenID int32) bool { + boa := tok.Encode(tokens.BOA, false) + audio := tok.Encode(tokens.Audio, false) + eoa := tok.Encode(tokens.EOA, false) + return len(boa) == 1 && len(audio) == 1 && len(eoa) == 1 && audio[0] == audioTokenID +} + // LoadWeights receives all tensors loaded from the manifest and assigns them // to model fields. func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index 124d7558810..012c25bbb64 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -1,22 +1,22 @@ package metadata import ( + "encoding/json" "fmt" "math" "slices" ) -const gemma4AudioFeatureSize = 128 - const ( - maxAudioHiddenSize = 8_192 - maxAudioLayers = 128 - maxAudioHeads = 256 - maxAudioOutputDims = 16_384 - maxAudioConvChannels = 4_096 - maxAudioKernelSize = 255 - maxAudioContextSize = 4_096 - maxTextHiddenSize = 65_536 + gemma4AudioFeatureSize = 128 + maxAudioHiddenSize = 8192 + maxAudioLayers = 128 + maxAudioHeads = 256 + maxAudioOutputDims = 16384 + maxAudioConvChannels = 4096 + maxAudioKernelSize = 255 + maxAudioContextSize = 4096 + maxTextHiddenSize = 65536 ) type AudioConfig struct { @@ -32,6 +32,69 @@ type AudioConfig struct { UseClippedLinears bool `json:"use_clipped_linears"` } +type audioProcessorConfig struct { + AudioSequenceLength int `json:"audio_seq_length"` + FeatureExtractor struct { + Dither float64 `json:"dither"` + FeatureSize int `json:"feature_size"` + FFTLength int `json:"fft_length"` + FFTOverdrive bool `json:"fft_overdrive"` + FrameLength int `json:"frame_length"` + HopLength int `json:"hop_length"` + InputScaleFactor float64 `json:"input_scale_factor"` + MaxFrequency float64 `json:"max_frequency"` + MelFloor float64 `json:"mel_floor"` + MinFrequency float64 `json:"min_frequency"` + PaddingSide string `json:"padding_side"` + PerBinMean []float64 `json:"per_bin_mean"` + PerBinStddev []float64 `json:"per_bin_stddev"` + Preemphasis float64 `json:"preemphasis"` + SamplingRate int `json:"sampling_rate"` + } `json:"feature_extractor"` +} + +// ValidateAudioRuntimeMetadata verifies the processor and tokenizer metadata +// required by the native Gemma 4 audio path. It deliberately accepts only the +// released processor contract implemented by the runner. +func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfigData []byte) error { + if err := validateAudioConfig(cfg); err != nil { + return err + } + var processor audioProcessorConfig + if len(processorData) == 0 { + return fmt.Errorf("missing processor_config.json") + } + if err := json.Unmarshal(processorData, &processor); err != nil { + return fmt.Errorf("parse processor_config.json: %w", err) + } + f := processor.FeatureExtractor + if processor.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || + f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || + f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || + f.MelFloor != 1e-3 || f.Preemphasis != 0 || f.PaddingSide != "right" || + len(f.PerBinMean) != 0 || len(f.PerBinStddev) != 0 { + return fmt.Errorf("unsupported Gemma 4 audio processor configuration") + } + if cfg.AudioTokenID <= 0 || cfg.TextConfig.VocabSize <= cfg.AudioTokenID { + return fmt.Errorf("invalid Gemma 4 audio token id %d", cfg.AudioTokenID) + } + if len(tokenizerConfigData) == 0 { + return fmt.Errorf("missing tokenizer_config.json") + } + var tokens struct { + BOA string `json:"boa_token"` + Audio string `json:"audio_token"` + EOA string `json:"eoa_token"` + } + if err := json.Unmarshal(tokenizerConfigData, &tokens); err != nil { + return fmt.Errorf("parse tokenizer_config.json: %w", err) + } + if tokens.BOA == "" || tokens.Audio == "" || tokens.EOA == "" { + return fmt.Errorf("missing Gemma 4 audio tokenizer tokens") + } + return nil +} + // ValidateAudioTensors verifies the normalized tensor names required by the // released Gemma 4 MLX audio loader. func ValidateAudioTensors(cfg ConfigFile, names []string) error { diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index 4bc79ba46ec..f995a4a5103 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -9,7 +9,8 @@ import ( func releasedAudioConfig(layers int) ConfigFile { return ConfigFile{ - TextConfig: TextConfig{HiddenSize: 2560}, + TextConfig: TextConfig{HiddenSize: 2560, VocabSize: 262144}, + AudioTokenID: 258881, AudioConfig: &AudioConfig{ AttentionChunkSize: 12, AttentionContextLeft: 13, ConvKernelSize: 5, HiddenSize: 1024, NumAttentionHeads: 8, @@ -19,6 +20,36 @@ func releasedAudioConfig(layers int) ConfigFile { } } +func TestValidateAudioRuntimeMetadata(t *testing.T) { + processor := []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`) + tokens := []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`) + cfg := releasedAudioConfig(12) + if err := ValidateAudioRuntimeMetadata(cfg, processor, tokens); err != nil { + t.Fatalf("ValidateAudioRuntimeMetadata() error = %v", err) + } + for _, tt := range []struct { + name string + processor, tokens []byte + edit func(*ConfigFile) + }{ + {"missing processor", nil, tokens, nil}, + {"unsupported processor", []byte(`{"audio_seq_length":749}`), tokens, nil}, + {"missing tokens", processor, nil, nil}, + {"incomplete tokens", processor, []byte(`{"audio_token":"<|audio|>"}`), nil}, + {"invalid token id", processor, tokens, func(cfg *ConfigFile) { cfg.AudioTokenID = cfg.TextConfig.VocabSize }}, + } { + t.Run(tt.name, func(t *testing.T) { + candidate := releasedAudioConfig(12) + if tt.edit != nil { + tt.edit(&candidate) + } + if err := ValidateAudioRuntimeMetadata(candidate, tt.processor, tt.tokens); err == nil { + t.Fatal("ValidateAudioRuntimeMetadata() error = nil") + } + }) + } +} + func completeAudioInventory(cfg ConfigFile) map[string]TensorDescriptor { tensors := make(map[string]TensorDescriptor) shapes, err := requiredAudioShapes(cfg) diff --git a/x/models/gemma4/metadata/vision.go b/x/models/gemma4/metadata/vision.go index 0d26583b74c..204e937570f 100644 --- a/x/models/gemma4/metadata/vision.go +++ b/x/models/gemma4/metadata/vision.go @@ -29,12 +29,14 @@ type ConfigFile struct { TextConfig TextConfig `json:"text_config"` VisionConfig *VisionConfig `json:"vision_config"` AudioConfig *AudioConfig `json:"audio_config"` + AudioTokenID int `json:"audio_token_id"` Quantization Quantization `json:"quantization"` QuantizationConfig Quantization `json:"quantization_config"` } type TextConfig struct { HiddenSize int `json:"hidden_size"` + VocabSize int `json:"vocab_size"` } type Quantization struct { From 19c8875b9495ddc023de0a528fcebead52653462 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 04:30:51 +0000 Subject: [PATCH 51/58] gemma4: align MLX audio capability with runtime Require the tokenizer and processor metadata consumed by the audio runtime, validate numerical configuration parity, and align capability reporting with the actual prepared audio path. Co-authored-by: Codex --- server/images.go | 5 +- server/images_test.go | 44 ++++++++++++++++- x/create/client/create.go | 5 +- x/create/client/create_test.go | 40 +++++++++++++++ x/models/gemma4/audio.go | 7 ++- .../gemma4/audio_forward_reference_test.go | 2 +- x/models/gemma4/metadata/audio.go | 49 ++++++++++++++----- x/models/gemma4/metadata/audio_test.go | 36 ++++++++++---- 8 files changed, 159 insertions(+), 29 deletions(-) diff --git a/server/images.go b/server/images.go index c7017bb000e..7c842c33132 100644 --- a/server/images.go +++ b/server/images.go @@ -572,10 +572,11 @@ func hasGemma4AudioTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifest } func hasGemma4AudioRuntimeMetadata(cfg gemma4metadata.ConfigFile, mf *manifest.Manifest) bool { - var processorData, tokenizerConfigData json.RawMessage + var processorData, tokenizerConfigData, tokenizerData json.RawMessage return mf.ReadConfigJSON("processor_config.json", &processorData) == nil && mf.ReadConfigJSON("tokenizer_config.json", &tokenizerConfigData) == nil && - gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData) == nil + mf.ReadConfigJSON("tokenizer.json", &tokenizerData) == nil && + gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData, tokenizerData) == nil } func gemma4AudioTensorDescriptors(layers []manifest.Layer) (map[string]gemma4metadata.TensorDescriptor, error) { diff --git a/server/images_test.go b/server/images_test.go index fccf83749d6..058b678e748 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -784,6 +784,31 @@ func TestGemma4InstalledAudioCapabilityDescriptorAndPayloadMatrix(t *testing.T) {name: "missing tokenizer metadata", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == "tokenizer_config.json" }) }}, + {name: "missing tokenizer vocabulary", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return slices.DeleteFunc(layers, func(layer manifest.Layer) bool { return layer.Name == "tokenizer.json" }) + }}, + {name: "invalid runtime scalar", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + cfg := gemma4AudioConfig(1) + cfg.AudioConfig.ResidualWeight = 0 + return replaceGemma4AudioConfigLayer(t, layers, cfg) + }}, + {name: "float32 overflow runtime scalar", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + valid, err := json.Marshal(gemma4AudioConfig(1)) + if err != nil { + t.Fatal(err) + } + overflow := []byte(strings.Replace(string(valid), `"gradient_clipping":10000000000`, `"gradient_clipping":1e39`, 1)) + if string(overflow) == string(valid) { + t.Fatal("failed to construct float32-overflow audio config") + } + return replaceGemma4ManifestJSON(t, layers, "config.json", overflow) + }}, + {name: "wrong audio token id", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4ManifestJSON(t, layers, "tokenizer.json", []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":8,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`)) + }}, + {name: "vocab-only markers are not singleton encodings", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { + return replaceGemma4ManifestJSON(t, layers, "tokenizer.json", []byte(`{"model":{"type":"BPE","vocab":{"<|audio>":5,"<|audio|>":7,"":6},"merges":[]},"added_tokens":[]}`)) + }}, {name: "near-match internal name", edit: func(t *testing.T, layers []manifest.Layer) []manifest.Layer { return replaceGemma4AudioFixtureLayer(t, layers, target, target+".extra", gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{3, 4}}, nil) }}, @@ -892,7 +917,7 @@ func TestGemma4InstalledAudioCapabilityRejectsUnboundedConfig(t *testing.T) { func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { t.Helper() descriptors := testGemma4AudioTensorDescriptors(1) - layers := make([]manifest.Layer, 0, len(descriptors)+3) + layers := make([]manifest.Layer, 0, len(descriptors)+4) for name, descriptor := range descriptors { layers = append(layers, gemma4AudioFixtureLayer(t, name, name, descriptor, nil)) } @@ -905,6 +930,7 @@ func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { for name, data := range map[string][]byte{ "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), + "tokenizer.json": []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), } { digest := createTestBlob(t, data) layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(data)), Name: name}) @@ -929,6 +955,20 @@ func replaceGemma4AudioConfigLayer(t *testing.T, layers []manifest.Layer, cfg *g return nil } +func replaceGemma4ManifestJSON(t *testing.T, layers []manifest.Layer, name string, contents []byte) []manifest.Layer { + t.Helper() + digest := createTestBlob(t, contents) + for i := range layers { + if layers[i].Name == name { + layers[i].Digest = digest + layers[i].Size = int64(len(contents)) + return layers + } + } + t.Fatalf("manifest config %q not found", name) + return nil +} + func gemma4AudioFixtureLayer(t *testing.T, manifestName, internalName string, descriptor gemma4metadata.TensorDescriptor, raw []byte) manifest.Layer { t.Helper() data := raw @@ -1196,8 +1236,10 @@ func gemma4AudioConfig(layers int) *gemma4metadata.ConfigFile { AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, NumHiddenLayers: layers, OutputProjDims: 3, + GradientClipping: 1e10, ResidualWeight: 0.5, RMSNormEps: 1e-6, SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, }, } diff --git a/x/create/client/create.go b/x/create/client/create.go index 7fc7bf92594..ce6426a0fa8 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -644,8 +644,9 @@ func gemma4ModelDirMediaCapabilities(modelDir string) (vision, audio bool) { } processorData, processorErr := os.ReadFile(filepath.Join(modelDir, "processor_config.json")) tokenizerConfigData, tokenizerErr := os.ReadFile(filepath.Join(modelDir, "tokenizer_config.json")) - audioReady := processorErr == nil && tokenizerErr == nil && - gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData) == nil && + tokenizerData, tokenizerDataErr := os.ReadFile(filepath.Join(modelDir, "tokenizer.json")) + audioReady := processorErr == nil && tokenizerErr == nil && tokenizerDataErr == nil && + gemma4metadata.ValidateAudioRuntimeMetadata(cfg, processorData, tokenizerConfigData, tokenizerData) == nil && gemma4metadata.ValidateAudioSourceInventory(cfg, tensors) == nil return gemma4metadata.ValidateVisionSourceInventory(cfg, tensors) == nil, audioReady } diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index d6bc444d98f..c6dd33023d5 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -673,8 +673,10 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, NumHiddenLayers: 1, OutputProjDims: 3, + GradientClipping: 1e10, ResidualWeight: 0.5, RMSNormEps: 1e-6, SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, }, } @@ -725,6 +727,21 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { delete(near, name) check(t, near, false) }) + t.Run("vocab-only markers are not singleton encodings", func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), configJSON, 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, valid) + writeGemma4AudioRuntimeConfigs(t, dir) + vocabOnly := []byte(`{"model":{"type":"BPE","vocab":{"<|audio>":5,"<|audio|>":7,"":6},"merges":[]},"added_tokens":[]}`) + if err := os.WriteFile(filepath.Join(dir, "tokenizer.json"), vocabOnly, 0o644); err != nil { + t.Fatal(err) + } + if got := inferSafetensorsCapabilities(dir, ""); slices.Contains(got, "audio") { + t.Fatalf("vocab-only marker capabilities = %v, did not expect audio", got) + } + }) }) } } @@ -736,8 +753,10 @@ func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testin AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, NumHiddenLayers: 1, OutputProjDims: 3, + GradientClipping: 1e10, ResidualWeight: 0.5, RMSNormEps: 1e-6, SubsamplingConvChannels: []int{2, 2}, UseClippedLinears: true, }, } @@ -781,6 +800,26 @@ func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testin } }) } + + t.Run("float32 overflow", func(t *testing.T) { + validConfig, err := json.Marshal(base) + if err != nil { + t.Fatal(err) + } + overflow := []byte(strings.Replace(string(validConfig), `"gradient_clipping":10000000000`, `"gradient_clipping":1e39`, 1)) + if string(overflow) == string(validConfig) { + t.Fatal("failed to construct float32-overflow audio config") + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), overflow, 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, valid) + writeGemma4AudioRuntimeConfigs(t, dir) + if got := inferSafetensorsCapabilities(dir, ""); slices.Contains(got, "audio") { + t.Fatalf("float32-overflow capabilities = %v, did not expect audio", got) + } + }) } func writeGemma4AudioRuntimeConfigs(t *testing.T, dir string) { @@ -788,6 +827,7 @@ func writeGemma4AudioRuntimeConfigs(t *testing.T, dir string) { files := map[string][]byte{ "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), + "tokenizer.json": []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), } for name, data := range files { if err := os.WriteFile(filepath.Join(dir, name), data, 0o644); err != nil { diff --git a/x/models/gemma4/audio.go b/x/models/gemma4/audio.go index db65f192e71..d9de6fb8275 100644 --- a/x/models/gemma4/audio.go +++ b/x/models/gemma4/audio.go @@ -124,9 +124,12 @@ func audioMetadataConfig(cfg *AudioConfig, textHidden int32) gemma4metadata.Conf TextConfig: gemma4metadata.TextConfig{HiddenSize: int(textHidden)}, AudioConfig: &gemma4metadata.AudioConfig{ AttentionChunkSize: int(cfg.AttentionChunkSize), AttentionContextLeft: int(cfg.AttentionContextLeft), - AttentionContextRight: int(cfg.AttentionContextRight), ConvKernelSize: int(cfg.ConvKernelSize), - HiddenSize: int(cfg.HiddenSize), NumAttentionHeads: int(cfg.NumAttentionHeads), + AttentionContextRight: int(cfg.AttentionContextRight), AttentionInvalidLogit: cfg.AttentionInvalidLogit, + AttentionLogitCap: cfg.AttentionLogitCap, ConvKernelSize: int(cfg.ConvKernelSize), + GradientClipping: cfg.GradientClipping, + HiddenSize: int(cfg.HiddenSize), NumAttentionHeads: int(cfg.NumAttentionHeads), NumHiddenLayers: int(cfg.NumHiddenLayers), OutputProjDims: int(cfg.OutputProjDims), + ResidualWeight: cfg.ResidualWeight, RMSNormEps: cfg.RMSNormEps, SubsamplingConvChannels: channels, UseClippedLinears: cfg.UseClippedLinears, }, } diff --git a/x/models/gemma4/audio_forward_reference_test.go b/x/models/gemma4/audio_forward_reference_test.go index 40a59356bc5..0fe6589d943 100644 --- a/x/models/gemma4/audio_forward_reference_test.go +++ b/x/models/gemma4/audio_forward_reference_test.go @@ -132,7 +132,7 @@ func TestAudioForwardReference(t *testing.T) { featureInput := mlx.FromValues(input.Features, 1, input.Frames, 128) forward := audioModel.Forward(featureInput, input) - compareAudioReference(t, "audio_forward_output", forward, reference.Get("audio_forward_output"), 0.5, 0.05) + compareAudioReference(t, "audio_physical_output", forward, reference.Get("audio_physical_output"), 0.5, 0.05) projectedForward := embedAudio.Forward(forward) inputIDs := reference.Get("input_ids") embedWeight := source.Get("model.language_model.embed_tokens.weight") diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index 012c25bbb64..eef2feb07f0 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -5,6 +5,8 @@ import ( "fmt" "math" "slices" + + "github.com/ollama/ollama/x/tokenizer" ) const ( @@ -20,16 +22,21 @@ const ( ) type AudioConfig struct { - AttentionChunkSize int `json:"attention_chunk_size"` - AttentionContextLeft int `json:"attention_context_left"` - AttentionContextRight int `json:"attention_context_right"` - ConvKernelSize int `json:"conv_kernel_size"` - HiddenSize int `json:"hidden_size"` - NumAttentionHeads int `json:"num_attention_heads"` - NumHiddenLayers int `json:"num_hidden_layers"` - OutputProjDims int `json:"output_proj_dims"` - SubsamplingConvChannels []int `json:"subsampling_conv_channels"` - UseClippedLinears bool `json:"use_clipped_linears"` + AttentionChunkSize int `json:"attention_chunk_size"` + AttentionContextLeft int `json:"attention_context_left"` + AttentionContextRight int `json:"attention_context_right"` + AttentionInvalidLogit float32 `json:"attention_invalid_logits_value"` + AttentionLogitCap float32 `json:"attention_logit_cap"` + ConvKernelSize int `json:"conv_kernel_size"` + GradientClipping float32 `json:"gradient_clipping"` + HiddenSize int `json:"hidden_size"` + NumAttentionHeads int `json:"num_attention_heads"` + NumHiddenLayers int `json:"num_hidden_layers"` + OutputProjDims int `json:"output_proj_dims"` + ResidualWeight float32 `json:"residual_weight"` + RMSNormEps float32 `json:"rms_norm_eps"` + SubsamplingConvChannels []int `json:"subsampling_conv_channels"` + UseClippedLinears bool `json:"use_clipped_linears"` } type audioProcessorConfig struct { @@ -56,7 +63,7 @@ type audioProcessorConfig struct { // ValidateAudioRuntimeMetadata verifies the processor and tokenizer metadata // required by the native Gemma 4 audio path. It deliberately accepts only the // released processor contract implemented by the runner. -func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfigData []byte) error { +func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfigData, tokenizerData []byte) error { if err := validateAudioConfig(cfg); err != nil { return err } @@ -92,6 +99,24 @@ func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfig if tokens.BOA == "" || tokens.Audio == "" || tokens.EOA == "" { return fmt.Errorf("missing Gemma 4 audio tokenizer tokens") } + if len(tokenizerData) == 0 { + return fmt.Errorf("missing tokenizer.json") + } + tok, err := tokenizer.LoadFromBytesWithConfig(tokenizerData, &tokenizer.TokenizerConfig{ + TokenizerConfigJSON: tokenizerConfigData, + }) + if err != nil { + return fmt.Errorf("parse tokenizer.json: %w", err) + } + for _, token := range []string{tokens.BOA, tokens.Audio, tokens.EOA} { + ids := tok.Encode(token, false) + if len(ids) != 1 || ids[0] < 0 || int(ids[0]) >= cfg.TextConfig.VocabSize { + return fmt.Errorf("audio tokenizer token %q is not a valid singleton token", token) + } + if token == tokens.Audio && int(ids[0]) != cfg.AudioTokenID { + return fmt.Errorf("audio tokenizer token id %d, want %d", ids[0], cfg.AudioTokenID) + } + } return nil } @@ -179,6 +204,8 @@ func validateAudioConfig(cfg ConfigFile) error { ac.AttentionChunkSize <= 0 || ac.AttentionChunkSize > maxAudioContextSize || ac.AttentionContextLeft <= 0 || ac.AttentionContextLeft > maxAudioContextSize || ac.AttentionContextRight < 0 || ac.AttentionContextRight > maxAudioContextSize || + ac.AttentionInvalidLogit >= 0 || ac.AttentionLogitCap <= 0 || + ac.GradientClipping <= 0 || ac.ResidualWeight <= 0 || ac.RMSNormEps <= 0 || len(ac.SubsamplingConvChannels) != 2 || ac.SubsamplingConvChannels[0] <= 0 || ac.SubsamplingConvChannels[0] > maxAudioConvChannels || ac.SubsamplingConvChannels[1] <= 0 || ac.SubsamplingConvChannels[1] > maxAudioConvChannels || diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index f995a4a5103..ce305b3d8e8 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -1,6 +1,7 @@ package metadata import ( + "encoding/json" "maps" "slices" "strings" @@ -13,8 +14,10 @@ func releasedAudioConfig(layers int) ConfigFile { AudioTokenID: 258881, AudioConfig: &AudioConfig{ AttentionChunkSize: 12, AttentionContextLeft: 13, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 5, HiddenSize: 1024, NumAttentionHeads: 8, NumHiddenLayers: layers, OutputProjDims: 1536, + GradientClipping: 1e10, ResidualWeight: 0.5, RMSNormEps: 1e-6, SubsamplingConvChannels: []int{128, 32}, UseClippedLinears: true, }, } @@ -23,27 +26,33 @@ func releasedAudioConfig(layers int) ConfigFile { func TestValidateAudioRuntimeMetadata(t *testing.T) { processor := []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`) tokens := []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`) + tokenizerData := []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":258881,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`) cfg := releasedAudioConfig(12) - if err := ValidateAudioRuntimeMetadata(cfg, processor, tokens); err != nil { + if err := ValidateAudioRuntimeMetadata(cfg, processor, tokens, tokenizerData); err != nil { t.Fatalf("ValidateAudioRuntimeMetadata() error = %v", err) } for _, tt := range []struct { - name string - processor, tokens []byte - edit func(*ConfigFile) + name string + processor []byte + tokens []byte + tokenizer []byte + edit func(*ConfigFile) }{ - {"missing processor", nil, tokens, nil}, - {"unsupported processor", []byte(`{"audio_seq_length":749}`), tokens, nil}, - {"missing tokens", processor, nil, nil}, - {"incomplete tokens", processor, []byte(`{"audio_token":"<|audio|>"}`), nil}, - {"invalid token id", processor, tokens, func(cfg *ConfigFile) { cfg.AudioTokenID = cfg.TextConfig.VocabSize }}, + {"missing processor", nil, tokens, tokenizerData, nil}, + {"unsupported processor", []byte(`{"audio_seq_length":749}`), tokens, tokenizerData, nil}, + {"missing tokens", processor, nil, tokenizerData, nil}, + {"incomplete tokens", processor, []byte(`{"audio_token":"<|audio|>"}`), tokenizerData, nil}, + {"missing tokenizer", processor, tokens, nil, nil}, + {"wrong tokenizer id", processor, tokens, []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), nil}, + {"vocab membership is not singleton encoding", processor, tokens, []byte(`{"model":{"type":"BPE","vocab":{"<|audio>":0,"<|audio|>":1,"":2},"merges":[]},"added_tokens":[]}`), func(cfg *ConfigFile) { cfg.AudioTokenID = 1; cfg.TextConfig.VocabSize = 3 }}, + {"invalid token id", processor, tokens, tokenizerData, func(cfg *ConfigFile) { cfg.AudioTokenID = cfg.TextConfig.VocabSize }}, } { t.Run(tt.name, func(t *testing.T) { candidate := releasedAudioConfig(12) if tt.edit != nil { tt.edit(&candidate) } - if err := ValidateAudioRuntimeMetadata(candidate, tt.processor, tt.tokens); err == nil { + if err := ValidateAudioRuntimeMetadata(candidate, tt.processor, tt.tokens, tt.tokenizer); err == nil { t.Fatal("ValidateAudioRuntimeMetadata() error = nil") } }) @@ -132,6 +141,8 @@ func TestValidateAudioConfig(t *testing.T) { {"heads", func(cfg *ConfigFile) { cfg.AudioConfig.NumAttentionHeads = 7 }}, {"conv channels", func(cfg *ConfigFile) { cfg.AudioConfig.SubsamplingConvChannels = []int{128} }}, {"even kernel", func(cfg *ConfigFile) { cfg.AudioConfig.ConvKernelSize = 4 }}, + {"invalid logit cap", func(cfg *ConfigFile) { cfg.AudioConfig.AttentionLogitCap = 0 }}, + {"invalid residual", func(cfg *ConfigFile) { cfg.AudioConfig.ResidualWeight = 0 }}, {"text hidden", func(cfg *ConfigFile) { cfg.TextConfig.HiddenSize = 0 }}, {"hidden overflow", func(cfg *ConfigFile) { cfg.AudioConfig.HiddenSize = 1 << 30; cfg.AudioConfig.NumAttentionHeads = 1 }}, {"impractical layers", func(cfg *ConfigFile) { cfg.AudioConfig.NumHiddenLayers = maxAudioLayers + 1 }}, @@ -153,6 +164,11 @@ func TestValidateAudioConfig(t *testing.T) { } }) } + + var overflow ConfigFile + if err := json.Unmarshal([]byte(`{"audio_config":{"gradient_clipping":1e39}}`), &overflow); err == nil { + t.Fatal("float32 overflow: json.Unmarshal() error = nil") + } } func TestAudioConfigSupportedBoundaries(t *testing.T) { From 78c28c9297ff49b63673d6e57ae9ad59556ecf09 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 04:52:03 +0000 Subject: [PATCH 52/58] gemma4: test malformed audio degradation Co-authored-by: Codex --- x/models/gemma4/audio_test.go | 174 ++++++++++++++++++++++++++++++++++ 1 file changed, 174 insertions(+) diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index e822bf2fc8b..fc86633c718 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -2,11 +2,16 @@ package gemma4 import ( "context" + "fmt" "maps" + "os" + "path/filepath" "strings" "testing" + "github.com/ollama/ollama/x/imagegen/manifest" "github.com/ollama/ollama/x/mlxrunner/mlx" + mlxmodel "github.com/ollama/ollama/x/mlxrunner/model" "github.com/ollama/ollama/x/mlxrunner/model/base" gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" ) @@ -258,6 +263,175 @@ func TestValidateAndLoadGemma4AudioWeights(t *testing.T) { } } +func TestNewModelDisablesMalformedAudioMetadata(t *testing.T) { + config := `{ + "architectures":["Gemma4ForConditionalGeneration"], + "boi_token_id":1, + "image_token_id":2, + "eoi_token_id":3, + "boa_token_id":4, + "audio_token_id":5, + "eoa_token_index":6, + "text_config":{ + "hidden_size":8, + "num_hidden_layers":1, + "intermediate_size":16, + "num_attention_heads":1, + "num_key_value_heads":1, + "head_dim":8, + "global_head_dim":8, + "vocab_size":7, + "rms_norm_eps":0.000001 + }, + "audio_config":{ + "attention_chunk_size":12, + "attention_context_left":13, + "attention_context_right":0, + "attention_invalid_logits_value":-1000000000.0, + "attention_logit_cap":50.0, + "conv_kernel_size":5, + "gradient_clipping":10000000000.0, + "hidden_size":1024, + "num_attention_heads":8, + "num_hidden_layers":12, + "output_proj_dims":1536, + "residual_weight":0.5, + "rms_norm_eps":0.000001, + "subsampling_conv_channels":[128,32], + "use_clipped_linears":true + } + }` + tokenizerData := []byte(`{ + "model":{"type":"BPE","vocab":{"a":0,"b":1,"c":2,"d":3,"<|audio>":4,"<|audio|>":5,"":6},"merges":[]}, + "added_tokens":[ + {"id":4,"content":"<|audio>","special":true}, + {"id":5,"content":"<|audio|>","special":true}, + {"id":6,"content":"","special":true} + ] + }`) + processorData := []byte(`{ + "feature_size":128,"sampling_rate":16000,"padding_value":0, + "return_attention_mask":true,"num_mel_bins":128,"n_fft":512, + "hop_length":160,"win_length":400,"max_length_seconds":30 + }`) + tokenizerConfigData := []byte(`{ + "boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":"" + }`) + tests := []struct { + name string + config string + tokenizer []byte + extra map[string][]byte + }{ + { + name: "malformed audio config", + config: strings.Replace(config, `"residual_weight":0.5`, `"residual_weight":0`, 1), + tokenizer: tokenizerData, + extra: map[string][]byte{ + "processor_config.json": processorData, + "tokenizer_config.json": tokenizerConfigData, + }, + }, + { + name: "missing processor config", + config: config, + tokenizer: tokenizerData, + extra: map[string][]byte{ + "tokenizer_config.json": tokenizerConfigData, + }, + }, + { + name: "malformed processor config", + config: config, + tokenizer: tokenizerData, + extra: map[string][]byte{ + "processor_config.json": []byte(`{"feature_extractor":{"sampling_rate":0}}`), + "tokenizer_config.json": tokenizerConfigData, + }, + }, + { + name: "incomplete tokenizer markers", + config: config, + tokenizer: []byte(`{ + "model":{"type":"BPE","vocab":{"a":0,"b":1,"c":2,"d":3,"<|audio>":4,"<|audio|>":5,"":6},"merges":[]}, + "added_tokens":[] + }`), + extra: map[string][]byte{ + "processor_config.json": processorData, + "tokenizer_config.json": tokenizerConfigData, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + loaded, err := newModel(testGemma4Root(t, []byte(tt.config), tt.tokenizer, tt.extra)) + if err != nil { + t.Fatalf("newModel() error = %v", err) + } + m := loaded.(*Model) + if m.AudioConfig != nil && m.AudioProcessorConfig != nil || m.Audio != nil || m.EmbedAudio != nil { + t.Fatalf("audio runtime retained state = (%+v, %+v, %v, %v)", m.AudioConfig, m.AudioProcessorConfig, m.Audio, m.EmbedAudio) + } + if m.TextConfig == nil || m.Tokenizer() == nil { + t.Fatal("text runtime was not preserved") + } + + var wantError string + for attempt := range 2 { + prepared, err := m.PrepareMedia(context.Background(), []base.Segment{{Kind: "audio", Data: []byte("wav")}}) + if err == nil || !strings.Contains(err.Error(), "does not support audio") { + t.Fatalf("PrepareMedia() attempt %d = (%+v, %v), want deterministic unsupported-audio error", attempt, prepared, err) + } + if attempt == 0 { + wantError = err.Error() + } else if err.Error() != wantError { + t.Fatalf("PrepareMedia() error = %q, want %q", err, wantError) + } + if prepared != nil || m.AudioConfig != nil && m.AudioProcessorConfig != nil || m.Audio != nil || m.EmbedAudio != nil { + t.Fatalf("PrepareMedia() attempt %d retained state: prepared=%+v model=%+v", attempt, prepared, m) + } + } + }) + } +} + +func testGemma4Root(t *testing.T, configData, tokenizerData []byte, extra map[string][]byte) *mlxmodel.Root { + t.Helper() + + blobDir := filepath.Join(t.TempDir(), "blobs") + if err := os.MkdirAll(blobDir, 0o755); err != nil { + t.Fatal(err) + } + layers := make([]manifest.ManifestLayer, 0, len(extra)+2) + writeConfig := func(name string, data []byte) { + digest := fmt.Sprintf("sha256:config-%d", len(layers)) + path := filepath.Join(blobDir, strings.Replace(digest, ":", "-", 1)) + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + layers = append(layers, manifest.ManifestLayer{ + MediaType: "application/vnd.ollama.image.json", + Digest: digest, + Name: name, + }) + } + writeConfig("config.json", configData) + writeConfig("tokenizer.json", tokenizerData) + for name, data := range extra { + writeConfig(name, data) + } + + return &mlxmodel.Root{Manifest: &manifest.ModelManifest{ + Manifest: &manifest.Manifest{ + SchemaVersion: 2, + MediaType: "application/vnd.ollama.image.model", + Layers: layers, + }, + BlobDir: blobDir, + }} +} + func tinyAudioConfig() *AudioConfig { return &AudioConfig{ AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionContextRight: 0, From 5620ab839ea552f4b3f820ca69b39e8c3ca333fc Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 05:00:46 +0000 Subject: [PATCH 53/58] gemma4: unify audio readiness parsing Co-authored-by: Codex --- x/models/gemma4/audio.go | 34 +++++++++++++------------- x/models/gemma4/audio_test.go | 25 +++++++++++++++++++ x/models/gemma4/metadata/audio.go | 30 +++++++++++++++++++++-- x/models/gemma4/metadata/audio_test.go | 32 ++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 19 deletions(-) diff --git a/x/models/gemma4/audio.go b/x/models/gemma4/audio.go index d9de6fb8275..aae6f6a431b 100644 --- a/x/models/gemma4/audio.go +++ b/x/models/gemma4/audio.go @@ -6,7 +6,6 @@ package gemma4 // docs/third-party/mlx-vlm.md and cross-checked against Transformers. import ( - "encoding/json" "errors" "fmt" "math" @@ -36,26 +35,27 @@ type AudioConfig struct { } func parseAudioConfig(configData []byte) (*AudioConfig, error) { - var wrapped struct { - AudioConfig *AudioConfig `json:"audio_config"` - } - if err := json.Unmarshal(configData, &wrapped); err != nil { - return nil, fmt.Errorf("parse Gemma4 audio config: %w", err) + cfg, err := gemma4metadata.ParseAudioConfig(configData) + if err != nil { + return nil, err } - if wrapped.AudioConfig == nil { + if cfg == nil { return nil, nil } - cfg := wrapped.AudioConfig - if cfg.HiddenSize <= 0 || cfg.NumHiddenLayers <= 0 || cfg.NumAttentionHeads <= 0 || - cfg.HiddenSize%cfg.NumAttentionHeads != 0 || cfg.OutputProjDims <= 0 || - cfg.AttentionChunkSize <= 0 || cfg.AttentionContextLeft <= 0 || cfg.AttentionContextRight < 0 || - cfg.ConvKernelSize <= 0 || cfg.ConvKernelSize%2 == 0 || len(cfg.SubsamplingConvChannels) != 2 || - cfg.SubsamplingConvChannels[0] <= 0 || cfg.SubsamplingConvChannels[1] <= 0 || - cfg.RMSNormEps <= 0 || cfg.ResidualWeight <= 0 || cfg.GradientClipping <= 0 || - cfg.AttentionLogitCap <= 0 || cfg.AttentionInvalidLogit >= 0 { - return nil, errors.New("invalid Gemma4 audio configuration") + channels := make([]int32, len(cfg.SubsamplingConvChannels)) + for i, channel := range cfg.SubsamplingConvChannels { + channels[i] = int32(channel) } - return cfg, nil + return &AudioConfig{ + AttentionChunkSize: int32(cfg.AttentionChunkSize), AttentionContextLeft: int32(cfg.AttentionContextLeft), + AttentionContextRight: int32(cfg.AttentionContextRight), AttentionInvalidLogit: cfg.AttentionInvalidLogit, + AttentionLogitCap: cfg.AttentionLogitCap, ConvKernelSize: int32(cfg.ConvKernelSize), + GradientClipping: cfg.GradientClipping, HiddenSize: int32(cfg.HiddenSize), + NumAttentionHeads: int32(cfg.NumAttentionHeads), NumHiddenLayers: int32(cfg.NumHiddenLayers), + OutputProjDims: int32(cfg.OutputProjDims), ResidualWeight: cfg.ResidualWeight, + RMSNormEps: cfg.RMSNormEps, SubsamplingConvChannels: channels, + UseClippedLinears: cfg.UseClippedLinears, + }, nil } type audioConvBlock struct { diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index fc86633c718..f2a108f96ea 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -57,6 +57,31 @@ func TestParseReleasedAudioConfig(t *testing.T) { } } +func TestParseAudioConfigUsesMetadataPredicate(t *testing.T) { + for _, edit := range []struct { + name string + old string + new string + }{ + {"hidden bound", `"hidden_size": 1024`, `"hidden_size": 8193`}, + {"layer bound", `"num_hidden_layers": 12`, `"num_hidden_layers": 129`}, + {"head bound", `"num_attention_heads": 8`, `"num_attention_heads": 257`}, + {"output bound", `"output_proj_dims": 1536`, `"output_proj_dims": 16385`}, + {"context bound", `"attention_context_left": 13`, `"attention_context_left": 4097`}, + {"float32 overflow", `"gradient_clipping": 10000000000.0`, `"gradient_clipping": 1e39`}, + } { + t.Run(edit.name, func(t *testing.T) { + data := strings.Replace(releasedGemma4AudioConfig, edit.old, edit.new, 1) + if data == releasedGemma4AudioConfig { + t.Fatal("test edit did not apply") + } + if _, err := parseAudioConfig([]byte(data)); err == nil { + t.Fatal("parseAudioConfig() error = nil") + } + }) + } +} + func TestParseTextConfigRejectsInvalidAudioMarkers(t *testing.T) { valid := `"boi_token_id":1,"image_token_id":2,"eoi_token_id":3,"boa_token_id":4,"audio_token_id":5,"eoa_token_index":6,"text_config":{"vocab_size":7}` if _, err := parseTextConfig([]byte(`{` + valid + `}`)); err != nil { diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index eef2feb07f0..305a6b22374 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -39,6 +39,23 @@ type AudioConfig struct { UseClippedLinears bool `json:"use_clipped_linears"` } +// ParseAudioConfig parses and validates the released Gemma 4 audio +// configuration. A model without audio_config is not an error and returns nil, +// allowing text-only Gemma 4 models to use the same parser. +func ParseAudioConfig(configData []byte) (*AudioConfig, error) { + var cfg ConfigFile + if err := json.Unmarshal(configData, &cfg); err != nil { + return nil, fmt.Errorf("parse Gemma 4 audio config: %w", err) + } + if cfg.AudioConfig == nil { + return nil, nil + } + if err := validateAudioConfigFields(cfg.AudioConfig); err != nil { + return nil, err + } + return cfg.AudioConfig, nil +} + type audioProcessorConfig struct { AudioSequenceLength int `json:"audio_seq_length"` FeatureExtractor struct { @@ -196,6 +213,16 @@ func validateAudioConfig(cfg ConfigFile) error { if ac == nil { return fmt.Errorf("missing audio_config") } + if err := validateAudioConfigFields(ac); err != nil { + return err + } + if cfg.TextConfig.HiddenSize <= 0 || cfg.TextConfig.HiddenSize > maxTextHiddenSize { + return fmt.Errorf("invalid Gemma 4 audio dimensions") + } + return nil +} + +func validateAudioConfigFields(ac *AudioConfig) error { if ac.HiddenSize <= 0 || ac.HiddenSize > maxAudioHiddenSize || ac.NumHiddenLayers <= 0 || ac.NumHiddenLayers > maxAudioLayers || ac.NumAttentionHeads <= 0 || ac.NumAttentionHeads > maxAudioHeads || ac.HiddenSize%ac.NumAttentionHeads != 0 || @@ -208,8 +235,7 @@ func validateAudioConfig(cfg ConfigFile) error { ac.GradientClipping <= 0 || ac.ResidualWeight <= 0 || ac.RMSNormEps <= 0 || len(ac.SubsamplingConvChannels) != 2 || ac.SubsamplingConvChannels[0] <= 0 || ac.SubsamplingConvChannels[0] > maxAudioConvChannels || - ac.SubsamplingConvChannels[1] <= 0 || ac.SubsamplingConvChannels[1] > maxAudioConvChannels || - cfg.TextConfig.HiddenSize <= 0 || cfg.TextConfig.HiddenSize > maxTextHiddenSize { + ac.SubsamplingConvChannels[1] <= 0 || ac.SubsamplingConvChannels[1] > maxAudioConvChannels { return fmt.Errorf("invalid Gemma 4 audio dimensions") } return nil diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index ce305b3d8e8..5df990a4958 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -23,6 +23,38 @@ func releasedAudioConfig(layers int) ConfigFile { } } +func TestParseAudioConfigEquivalence(t *testing.T) { + valid, err := json.Marshal(releasedAudioConfig(12)) + if err != nil { + t.Fatal(err) + } + for _, tt := range []struct { + name string + data []byte + ok bool + }{ + {"released", valid, true}, + {"missing", []byte(`{"model_type":"gemma4"}`), true}, + {"malformed", []byte(`{"audio_config":`), false}, + {"partial", []byte(`{"audio_config":{"hidden_size":1024}}`), false}, + {"near match", []byte(`{"audio_config_extra":{"hidden_size":1024}}`), true}, + {"bounded width", []byte(strings.Replace(string(valid), `"hidden_size":1024`, `"hidden_size":1073741824`, 1)), false}, + {"float32 overflow", []byte(strings.Replace(string(valid), `"gradient_clipping":10000000000`, `"gradient_clipping":1e39`, 1)), false}, + } { + t.Run(tt.name, func(t *testing.T) { + cfg, err := ParseAudioConfig(tt.data) + if (err == nil) != tt.ok { + t.Fatalf("ParseAudioConfig() = (%+v, %v), ok = %v", cfg, err, tt.ok) + } + if tt.name == "missing" || tt.name == "near match" { + if cfg != nil { + t.Fatalf("ParseAudioConfig() = %+v, want nil", cfg) + } + } + }) + } +} + func TestValidateAudioRuntimeMetadata(t *testing.T) { processor := []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`) tokens := []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`) From 163a747308b8a82f271885e6f38c454fa699f6ce Mon Sep 17 00:00:00 2001 From: Philipp Date: Fri, 14 Aug 2026 18:31:53 +0000 Subject: [PATCH 54/58] gemma4: support unified MLX audio projection Support the unified 12B raw-waveform audio projection within the official lazy media contract while preserving the released conformer path and shared readiness validation. Co-authored-by: Codex --- server/images_test.go | 89 +++++++++++++++++++ server/model_list_cache_test.go | 2 + x/create/client/create_test.go | 47 ++++++++++ x/models/gemma4/audio.go | 19 ++-- .../gemma4/audio_forward_reference_test.go | 83 +++++++++-------- x/models/gemma4/audio_processor.go | 70 +++++++++++---- x/models/gemma4/audio_processor_test.go | 53 +++++++++++ x/models/gemma4/audio_test.go | 54 +++++++++++ x/models/gemma4/gemma4.go | 10 ++- x/models/gemma4/metadata/audio.go | 70 +++++++++++---- x/models/gemma4/metadata/audio_test.go | 68 ++++++++++++++ x/models/gemma4/vision.go | 12 ++- 12 files changed, 493 insertions(+), 84 deletions(-) diff --git a/server/images_test.go b/server/images_test.go index 058b678e748..93e7ad37c9c 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -938,6 +938,63 @@ func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { return layers } +func gemma4UnifiedAudioConfig() *gemma4metadata.ConfigFile { + return &gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: 3840, VocabSize: 262144}, + AudioTokenID: 258881, + AudioConfig: &gemma4metadata.AudioConfig{ + ModelType: "gemma4_unified_audio", AudioEmbedDim: 640, + AudioSamplesPerToken: 640, HiddenSize: 640, OutputProjDims: 640, RMSNormEps: 1e-6, + }, + } +} + +func gemma4UnifiedAudioManifestLayers(t *testing.T, complete bool) []manifest.Layer { + t.Helper() + layers := make([]manifest.Layer, 0, 5) + if complete { + layers = append(layers, gemma4AudioFixtureLayer(t, "model.embed_audio.embedding_projection.weight", + "model.embed_audio.embedding_projection.weight", gemma4metadata.TensorDescriptor{Dtype: "F32", Shape: []int32{3840, 640}}, nil)) + } + config, err := json.Marshal(gemma4UnifiedAudioConfig()) + if err != nil { + t.Fatal(err) + } + for name, data := range map[string][]byte{ + "config.json": config, + "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_extractor_type":"Gemma4UnifiedAudioFeatureExtractor","audio_samples_per_token":640,"feature_size":640,"padding_side":"right","sampling_rate":16000}}`), + "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), + "tokenizer.json": []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":258880,"content":"<|audio>","special":true},{"id":258881,"content":"<|audio|>","special":true},{"id":258883,"content":"","special":true}]}`), + } { + digest := createTestBlob(t, data) + layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(data)), Name: name}) + } + return layers +} + +func TestGemma4UnifiedAudioManifestHydratesCapabilityState(t *testing.T) { + setTestHome(t, t.TempDir()) + name := "gemma4-unified-audio-hydration" + createSafetensorsTestModel( + t, name, + model.ConfigV2{ + ModelFormat: "safetensors", Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "audio"}, + }, + gemma4UnifiedAudioManifestLayers(t, true), + ) + m, err := GetModel(name) + if err != nil { + t.Fatal(err) + } + if m.Gemma4AudioConfig == nil || !m.Gemma4AudioReady || len(m.Gemma4AudioTensors) != 1 { + t.Fatalf("unified hydration state = config:%v ready:%v tensors:%v", m.Gemma4AudioConfig != nil, m.Gemma4AudioReady, m.Gemma4AudioTensors) + } + if !slices.Contains(m.Capabilities(), model.CapabilityAudio) { + t.Fatalf("hydrated unified capabilities = %v, want audio", m.Capabilities()) + } +} + func replaceGemma4AudioConfigLayer(t *testing.T, layers []manifest.Layer, cfg *gemma4metadata.ConfigFile) []manifest.Layer { t.Helper() data, err := json.Marshal(cfg) @@ -1087,6 +1144,38 @@ func TestGemma4InstalledVisionCapabilityRejectsMissingOrZeroTextWidth(t *testing } } +func TestGemma4UnifiedAudioCapabilityRequiresProjection(t *testing.T) { + setTestHome(t, t.TempDir()) + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererSmall, + Capabilities: []string{"completion", "audio"}, + } + + complete := gemma4UnifiedAudioManifestLayers(t, true) + createSafetensorsTestModel(t, "gemma4-unified-audio-complete", cfg, complete) + m, err := GetModel("gemma4-unified-audio-complete") + if err != nil { + t.Fatal(err) + } + if !slices.Contains(m.Capabilities(), model.CapabilityAudio) || !m.Gemma4AudioReady { + t.Fatalf("complete unified audio capabilities = %v ready=%v, want audio", m.Capabilities(), m.Gemma4AudioReady) + } + + partial := gemma4UnifiedAudioManifestLayers(t, false) + createSafetensorsTestModel(t, "gemma4-unified-audio-partial", cfg, partial) + m, err = GetModel("gemma4-unified-audio-partial") + if err != nil { + t.Fatal(err) + } + if slices.Contains(m.Capabilities(), model.CapabilityAudio) { + t.Fatalf("partial unified audio capabilities = %v, did not expect audio", m.Capabilities()) + } + if !m.Gemma4AudioReady { + t.Fatal("valid unified audio runtime metadata was not marked ready") + } +} + func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { t.Helper() diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index a2207377198..b937cf9379c 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -190,6 +190,8 @@ func TestModelListSummaryGemma4AudioRequiresRuntimeMetadataAndTensors(t *testing {"partial", slices.DeleteFunc(gemma4AudioManifestLayers(t), func(layer manifest.Layer) bool { return layer.Name == "model.audio_tower.layers.0.self_attn.q_proj.input_max" }), false}, + {"unified-complete", gemma4UnifiedAudioManifestLayers(t, true), true}, + {"unified-partial", gemma4UnifiedAudioManifestLayers(t, false), false}, } { t.Run(tt.name, func(t *testing.T) { modelName := "list-gemma4-audio-" + tt.name diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index c6dd33023d5..eefdb653c1c 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -822,6 +822,44 @@ func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testin }) } +func TestInferSafetensorsCapabilitiesGemma4UnifiedAudio(t *testing.T) { + cfg := gemma4metadata.ConfigFile{ + Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}, + ModelType: "gemma4_unified", + TextConfig: gemma4metadata.TextConfig{HiddenSize: 3840, VocabSize: 262144}, + AudioTokenID: 258881, + AudioConfig: &gemma4metadata.AudioConfig{ + ModelType: "gemma4_unified_audio", AudioEmbedDim: 640, + AudioSamplesPerToken: 640, HiddenSize: 640, OutputProjDims: 640, RMSNormEps: 1e-6, + }, + } + dir := t.TempDir() + configJSON, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "config.json"), configJSON, 0o644); err != nil { + t.Fatal(err) + } + for name, data := range map[string]string{ + "processor_config.json": `{"audio_seq_length":750,"feature_extractor":{"audio_samples_per_token":640,"feature_extractor_type":"Gemma4UnifiedAudioFeatureExtractor","feature_size":640,"padding_side":"right","sampling_rate":16000}}`, + "tokenizer_config.json": `{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`, + "tokenizer.json": `{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":258881,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`, + } { + if err := os.WriteFile(filepath.Join(dir, name), []byte(data), 0o644); err != nil { + t.Fatal(err) + } + } + shapes, err := gemma4metadata.RequiredAudioTensorShapes(cfg) + if err != nil { + t.Fatal(err) + } + writeClientSafetensorsWithShapes(t, dir, shapes) + if got, want := inferSafetensorsCapabilities(dir, ""), []string{"completion", "audio"}; !slices.Equal(got, want) { + t.Fatalf("unified capabilities = %#v, want %#v", got, want) + } +} + func writeGemma4AudioRuntimeConfigs(t *testing.T, dir string) { t.Helper() files := map[string][]byte{ @@ -1030,6 +1068,15 @@ func writeClientSafetensors(t *testing.T, dir string, names ...string) { } } +func writeClientSafetensorsWithShapes(t *testing.T, dir string, shapes map[string][]int32) { + t.Helper() + descriptors := make(map[string]gemma4metadata.TensorDescriptor, len(shapes)) + for name, shape := range shapes { + descriptors[name] = gemma4metadata.TensorDescriptor{Dtype: "BF16", Shape: shape} + } + writeClientSafetensorDescriptors(t, dir, descriptors) +} + func writeClientSafetensorDescriptors(t *testing.T, dir string, descriptors map[string]gemma4metadata.TensorDescriptor) { t.Helper() tensors := make([]*safetensors.TensorData, 0, len(descriptors)) diff --git a/x/models/gemma4/audio.go b/x/models/gemma4/audio.go index aae6f6a431b..6d1d0762f9a 100644 --- a/x/models/gemma4/audio.go +++ b/x/models/gemma4/audio.go @@ -1,9 +1,8 @@ package gemma4 -// The Gemma 4 audio encoder follows the Universal Speech Model Conformer used -// by the released checkpoint. The MLX layout and execution structure are -// adapted from MLX-VLM's MIT-licensed implementation pinned in -// docs/third-party/mlx-vlm.md and cross-checked against Transformers. +// Gemma 4 E2B/E4B use a Universal Speech Model Conformer. The unified 12B +// architecture projects raw 640-sample waveform blocks directly into the +// language model. Both paths use the canonical metadata readiness contract. import ( "errors" @@ -17,6 +16,9 @@ import ( ) type AudioConfig struct { + ModelType string `json:"model_type"` + AudioEmbedDim int32 `json:"audio_embed_dim"` + AudioSamplesPerToken int32 `json:"audio_samples_per_token"` AttentionChunkSize int32 `json:"attention_chunk_size"` AttentionContextLeft int32 `json:"attention_context_left"` AttentionContextRight int32 `json:"attention_context_right"` @@ -47,6 +49,7 @@ func parseAudioConfig(configData []byte) (*AudioConfig, error) { channels[i] = int32(channel) } return &AudioConfig{ + ModelType: cfg.ModelType, AudioEmbedDim: int32(cfg.AudioEmbedDim), AudioSamplesPerToken: int32(cfg.AudioSamplesPerToken), AttentionChunkSize: int32(cfg.AttentionChunkSize), AttentionContextLeft: int32(cfg.AttentionContextLeft), AttentionContextRight: int32(cfg.AttentionContextRight), AttentionInvalidLogit: cfg.AttentionInvalidLogit, AttentionLogitCap: cfg.AttentionLogitCap, ConvKernelSize: int32(cfg.ConvKernelSize), @@ -58,6 +61,10 @@ func parseAudioConfig(configData []byte) (*AudioConfig, error) { }, nil } +func (c *AudioConfig) unified() bool { + return c != nil && c.ModelType == "gemma4_unified_audio" +} + type audioConvBlock struct { Weight *mlx.Array Norm *nn.LayerNorm @@ -123,7 +130,9 @@ func audioMetadataConfig(cfg *AudioConfig, textHidden int32) gemma4metadata.Conf return gemma4metadata.ConfigFile{ TextConfig: gemma4metadata.TextConfig{HiddenSize: int(textHidden)}, AudioConfig: &gemma4metadata.AudioConfig{ - AttentionChunkSize: int(cfg.AttentionChunkSize), AttentionContextLeft: int(cfg.AttentionContextLeft), + ModelType: cfg.ModelType, AudioEmbedDim: int(cfg.AudioEmbedDim), + AudioSamplesPerToken: int(cfg.AudioSamplesPerToken), + AttentionChunkSize: int(cfg.AttentionChunkSize), AttentionContextLeft: int(cfg.AttentionContextLeft), AttentionContextRight: int(cfg.AttentionContextRight), AttentionInvalidLogit: cfg.AttentionInvalidLogit, AttentionLogitCap: cfg.AttentionLogitCap, ConvKernelSize: int(cfg.ConvKernelSize), GradientClipping: cfg.GradientClipping, diff --git a/x/models/gemma4/audio_forward_reference_test.go b/x/models/gemma4/audio_forward_reference_test.go index 0fe6589d943..9fe14216e24 100644 --- a/x/models/gemma4/audio_forward_reference_test.go +++ b/x/models/gemma4/audio_forward_reference_test.go @@ -70,7 +70,7 @@ func TestAudioForwardReference(t *testing.T) { defer reference.Free() if wantFrames := reference.Get("input_features").Dim(1); wantFrames > input.Frames { padding := wantFrames - input.Frames - input.Features = append(input.Features, make([]float32, padding*128)...) + input.Features = append(input.Features, make([]float32, padding*input.FeatureSize)...) input.FeatureMask = append(input.FeatureMask, make([]bool, padding)...) input.Frames = wantFrames } @@ -86,54 +86,59 @@ func TestAudioForwardReference(t *testing.T) { t.Fatalf("source is missing %s", name) } } - audioModel, err := loadAudioModel(tensors, audioConfig, textConfig.HiddenSize, 0, 0, "", nil) - if err != nil { - t.Fatal(err) - } embedAudio, err := loadMultimodalEmbedder(tensors, "embed_audio", audioConfig.RMSNormEps, 0, 0, "", nil) if err != nil { t.Fatal(err) } - features := mlx.FromValues(input.Features, 1, input.Frames, 128) + features := mlx.FromValues(input.Features, 1, input.Frames, input.FeatureSize) compareAudioReference(t, "input_features", features, reference.Get("input_features"), 2e-5, 1e-5) compareAudioMaskReference(t, "input_features_mask", input.FeatureMask, reference.Get("input_features_mask"), false) - valid := append([]bool(nil), input.FeatureMask...) - x := mlx.ExpandDims(features, -1) - x, valid = audioModel.Conv0.Forward(x, valid) - x, valid = audioModel.Conv1.Forward(x, valid) - compareAudioMaskReference(t, "audio_forward_mask", valid, reference.Get("audio_forward_mask"), true) - validOutputs := 0 - for _, value := range valid { - if value { - validOutputs++ + var projectedForward *mlx.Array + if audioConfig.unified() { + projectedForward = embedAudio.Forward(features) + compareAudioReference(t, "multimodal_projection", projectedForward, reference.Get("multimodal_projection"), 0.5, 0.05) + } else { + audioModel, err := loadAudioModel(tensors, audioConfig, textConfig.HiddenSize, 0, 0, "", nil) + if err != nil { + t.Fatal(err) } - } - if validOutputs != input.SoftTokens { - t.Fatalf("valid audio outputs = %d, want %d", validOutputs, input.SoftTokens) - } - x = mlx.Reshape(x, 1, int32(x.Dim(1)), int32(x.Dim(2)*x.Dim(3))) - x = audioModel.InputProj.Forward(x) - compareAudioReference(t, "subsample", x, reference.Get("subsample"), 0.5, 0.05) - for index, layer := range audioModel.Layers { - x = layer.Forward(x, valid) - if index == 0 || index == len(audioModel.Layers)-1 { - name := "layer_0" - if index != 0 { - name = "layer_11" + valid := append([]bool(nil), input.FeatureMask...) + x := mlx.ExpandDims(features, -1) + x, valid = audioModel.Conv0.Forward(x, valid) + x, valid = audioModel.Conv1.Forward(x, valid) + compareAudioMaskReference(t, "audio_forward_mask", valid, reference.Get("audio_forward_mask"), true) + validOutputs := 0 + for _, value := range valid { + if value { + validOutputs++ } - compareAudioReference(t, name, x, reference.Get(name), 0.5, 0.05) } - } - x = audioModel.OutputProj.Forward(x) - compareAudioReference(t, "output_projection", x, reference.Get("output_projection"), 0.5, 0.05) - x = embedAudio.Forward(x) - compareAudioReference(t, "multimodal_projection", x, reference.Get("multimodal_projection"), 0.5, 0.05) + if validOutputs != input.SoftTokens { + t.Fatalf("valid audio outputs = %d, want %d", validOutputs, input.SoftTokens) + } + x = mlx.Reshape(x, 1, int32(x.Dim(1)), int32(x.Dim(2)*x.Dim(3))) + x = audioModel.InputProj.Forward(x) + compareAudioReference(t, "subsample", x, reference.Get("subsample"), 0.5, 0.05) + for index, layer := range audioModel.Layers { + x = layer.Forward(x, valid) + if index == 0 || index == len(audioModel.Layers)-1 { + name := "layer_0" + if index != 0 { + name = "layer_11" + } + compareAudioReference(t, name, x, reference.Get(name), 0.5, 0.05) + } + } + x = audioModel.OutputProj.Forward(x) + compareAudioReference(t, "output_projection", x, reference.Get("output_projection"), 0.5, 0.05) + x = embedAudio.Forward(x) + compareAudioReference(t, "multimodal_projection", x, reference.Get("multimodal_projection"), 0.5, 0.05) - featureInput := mlx.FromValues(input.Features, 1, input.Frames, 128) - forward := audioModel.Forward(featureInput, input) - compareAudioReference(t, "audio_physical_output", forward, reference.Get("audio_physical_output"), 0.5, 0.05) - projectedForward := embedAudio.Forward(forward) + forward := audioModel.Forward(features, input) + compareAudioReference(t, "audio_physical_output", forward, reference.Get("audio_physical_output"), 0.5, 0.05) + projectedForward = embedAudio.Forward(forward) + } inputIDs := reference.Get("input_ids") embedWeight := source.Get("model.language_model.embed_tokens.weight") if inputIDs == nil || embedWeight == nil { @@ -171,7 +176,7 @@ func TestAudioForwardReference(t *testing.T) { Seq: 0, Pos: start - 1, Features: mlx.Squeeze(projectedForward, 0), Opaque: payload, }}, } - hidden, _ := fullModel.Forward(modelBatch, nil) + hidden, _ := fullModel.Forward(modelBatch, fullModel.NewCaches()) logits := fullModel.Unembed(hidden) last := mlx.SliceStartStop(logits, []int32{0, int32(inputIDs.Dim(1) - 1), 0}, []int32{1, int32(inputIDs.Dim(1)), textConfig.VocabSize}) wantLogits := reference.Get("prefill_logits") diff --git a/x/models/gemma4/audio_processor.go b/x/models/gemma4/audio_processor.go index 17fe45276e5..1a751d2b836 100644 --- a/x/models/gemma4/audio_processor.go +++ b/x/models/gemma4/audio_processor.go @@ -25,27 +25,30 @@ const ( type AudioProcessorConfig struct { AudioSequenceLength int `json:"audio_seq_length"` FeatureExtractor struct { - Dither float64 `json:"dither"` - FeatureSize int `json:"feature_size"` - FFTLength int `json:"fft_length"` - FFTOverdrive bool `json:"fft_overdrive"` - FrameLength int `json:"frame_length"` - HopLength int `json:"hop_length"` - InputScaleFactor float64 `json:"input_scale_factor"` - MaxFrequency float64 `json:"max_frequency"` - MelFloor float64 `json:"mel_floor"` - MinFrequency float64 `json:"min_frequency"` - PaddingSide string `json:"padding_side"` - PerBinMean []float64 `json:"per_bin_mean"` - PerBinStddev []float64 `json:"per_bin_stddev"` - Preemphasis float64 `json:"preemphasis"` - SamplingRate int `json:"sampling_rate"` + Type string `json:"feature_extractor_type"` + AudioSamplesPerToken int `json:"audio_samples_per_token"` + Dither float64 `json:"dither"` + FeatureSize int `json:"feature_size"` + FFTLength int `json:"fft_length"` + FFTOverdrive bool `json:"fft_overdrive"` + FrameLength int `json:"frame_length"` + HopLength int `json:"hop_length"` + InputScaleFactor float64 `json:"input_scale_factor"` + MaxFrequency float64 `json:"max_frequency"` + MelFloor float64 `json:"mel_floor"` + MinFrequency float64 `json:"min_frequency"` + PaddingSide string `json:"padding_side"` + PerBinMean []float64 `json:"per_bin_mean"` + PerBinStddev []float64 `json:"per_bin_stddev"` + Preemphasis float64 `json:"preemphasis"` + SamplingRate int `json:"sampling_rate"` } `json:"feature_extractor"` } type gemma4AudioInput struct { Features []float32 FeatureMask []bool + FeatureSize int Frames int SoftTokens int } @@ -83,6 +86,13 @@ func validateReleasedAudioProcessorConfig(cfg *AudioProcessorConfig) error { return errors.New("Gemma4 MLX model has no supported audio processor configuration") } f := cfg.FeatureExtractor + if f.Type == "Gemma4UnifiedAudioFeatureExtractor" { + if cfg.AudioSequenceLength != 750 || f.FeatureSize != 640 || f.SamplingRate != 16000 || + f.AudioSamplesPerToken != 640 || f.PaddingSide != "right" { + return errors.New("unsupported Gemma4 unified audio processor configuration") + } + return nil + } if cfg.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || @@ -101,6 +111,9 @@ func preprocessGemma4Audio(ctx context.Context, data []byte, cfg *AudioProcessor if err != nil { return nil, err } + if cfg.FeatureExtractor.Type == "Gemma4UnifiedAudioFeatureExtractor" { + return preprocessGemma4UnifiedAudio(samples, cfg) + } features, featureMask, err := computeGemma4LogMel(ctx, samples, cfg) if err != nil { return nil, err @@ -125,7 +138,32 @@ func preprocessGemma4Audio(ctx context.Context, data []byte, cfg *AudioProcessor return nil, fmt.Errorf("Gemma4 audio token count %d exceeds limit %d", softTokens, cfg.AudioSequenceLength) } return &gemma4AudioInput{ - Features: features, FeatureMask: featureMask, Frames: len(featureMask), SoftTokens: softTokens, + Features: features, FeatureMask: featureMask, FeatureSize: cfg.FeatureExtractor.FeatureSize, + Frames: len(featureMask), SoftTokens: softTokens, + }, nil +} + +func preprocessGemma4UnifiedAudio(samples []float32, cfg *AudioProcessorConfig) (*gemma4AudioInput, error) { + frameSize := cfg.FeatureExtractor.AudioSamplesPerToken + if frameSize <= 0 { + return nil, errors.New("invalid Gemma4 unified audio frame size") + } + softTokens := (len(samples) + frameSize - 1) / frameSize + if softTokens == 0 { + return nil, errors.New("Gemma4 audio is too short to encode") + } + if softTokens > cfg.AudioSequenceLength { + return nil, fmt.Errorf("Gemma4 audio token count %d exceeds limit %d", softTokens, cfg.AudioSequenceLength) + } + features := make([]float32, softTokens*frameSize) + copy(features, samples) + mask := make([]bool, softTokens) + for i := range mask { + mask[i] = true + } + return &gemma4AudioInput{ + Features: features, FeatureMask: mask, FeatureSize: frameSize, + Frames: softTokens, SoftTokens: softTokens, }, nil } diff --git a/x/models/gemma4/audio_processor_test.go b/x/models/gemma4/audio_processor_test.go index e3c561d9808..14b37eb5628 100644 --- a/x/models/gemma4/audio_processor_test.go +++ b/x/models/gemma4/audio_processor_test.go @@ -71,6 +71,59 @@ func TestReleasedAudioProcessorRejectsMalformedDimensions(t *testing.T) { } } +func TestParseReleasedUnifiedAudioProcessorConfig(t *testing.T) { + data := []byte(`{ + "audio_seq_length":750, + "feature_extractor":{ + "audio_samples_per_token":640, + "feature_extractor_type":"Gemma4UnifiedAudioFeatureExtractor", + "feature_size":640,"padding_side":"right","padding_value":0, + "return_attention_mask":true,"sampling_rate":16000 + } + }`) + cfg, err := parseAudioProcessorConfig(data) + if err != nil { + t.Fatal(err) + } + if cfg.FeatureExtractor.AudioSamplesPerToken != 640 || cfg.FeatureExtractor.FeatureSize != 640 { + t.Fatalf("processor config = %#v", cfg) + } + + bad := bytes.Replace(data, []byte(`"audio_samples_per_token":640`), []byte(`"audio_samples_per_token":320`), 1) + if _, err := parseAudioProcessorConfig(bad); err == nil { + t.Fatal("invalid unified processor: error = nil") + } +} + +func TestPreprocessUnifiedAudioFrames(t *testing.T) { + data := []byte(`{ + "audio_seq_length":750, + "feature_extractor":{ + "audio_samples_per_token":640, + "feature_extractor_type":"Gemma4UnifiedAudioFeatureExtractor", + "feature_size":640,"padding_side":"right","sampling_rate":16000 + } + }`) + cfg, err := parseAudioProcessorConfig(data) + if err != nil { + t.Fatal(err) + } + frames := make([][]float64, 641) + for i := range frames { + frames[i] = []float64{0.25} + } + input, err := preprocessGemma4Audio(context.Background(), makeTestWAV(t, 1, 16, 16000, frames), cfg) + if err != nil { + t.Fatal(err) + } + if input.SoftTokens != 2 || input.Frames != 2 || input.FeatureSize != 640 || len(input.Features) != 1280 { + t.Fatalf("unified input = %+v, feature count %d", input, len(input.Features)) + } + if input.Features[640] != 0.25 || input.Features[641] != 0 { + t.Fatalf("second frame values = %v, %v; want 0.25, 0", input.Features[640], input.Features[641]) + } +} + func TestGemma4LogMelReference(t *testing.T) { cfg := defaultAudioProcessorConfig() samples := make([]float32, 4000) diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index f2a108f96ea..500bff794ca 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -2,6 +2,7 @@ package gemma4 import ( "context" + "encoding/json" "fmt" "maps" "os" @@ -14,9 +15,14 @@ import ( mlxmodel "github.com/ollama/ollama/x/mlxrunner/model" "github.com/ollama/ollama/x/mlxrunner/model/base" gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" + "github.com/ollama/ollama/x/models/nn" ) const releasedGemma4AudioConfig = `{ + "text_config": { + "hidden_size": 2560, + "vocab_size": 262144 + }, "audio_config": { "attention_chunk_size": 12, "attention_context_left": 13, @@ -139,6 +145,54 @@ func TestParseTextConfigRejectsInvalidAudioMarkers(t *testing.T) { } } +func TestParseReleasedUnifiedAudioConfig(t *testing.T) { + cfg, err := parseAudioConfig([]byte(`{ + "text_config":{"hidden_size":3840,"vocab_size":262144}, + "audio_config":{ + "model_type":"gemma4_unified_audio", + "audio_embed_dim":640,"audio_samples_per_token":640, + "hidden_size":640,"output_proj_dims":640,"rms_norm_eps":0.000001 + } + }`)) + if err != nil { + t.Fatal(err) + } + if !cfg.unified() || cfg.AudioSamplesPerToken != 640 || cfg.OutputProjDims != 640 { + t.Fatalf("unified audio config = %+v", cfg) + } + + bad := *cfg + bad.AudioSamplesPerToken = 320 + data, err := json.Marshal(map[string]any{ + "text_config": map[string]any{"hidden_size": 3840, "vocab_size": 262144}, + "audio_config": bad, + }) + if err != nil { + t.Fatal(err) + } + if _, err := parseAudioConfig(data); err == nil { + t.Fatal("invalid unified audio config: error = nil") + } +} + +func TestEncodeUnifiedAudioMedia(t *testing.T) { + skipIfNoMLX(t) + m := &Model{ + AudioConfig: &AudioConfig{ModelType: "gemma4_unified_audio"}, + EmbedAudio: &MultimodalEmbedder{ + Projection: nn.NewLinear(mlx.FromValues([]float32{1, 0, 0, 1}, 2, 2), nil), + Eps: 1e-6, + }, + } + item := &base.PreparedItem{ + Opaque: gemma4MediaPayload{Audio: &gemma4AudioInput{FeatureSize: 2, Frames: 1, SoftTokens: 1}}, + } + features := m.EncodeMedia(item, mlx.FromValues([]float32{3, 4}, 1, 1, 2)) + if features.NumDims() != 2 || features.Dim(0) != 1 || features.Dim(1) != 2 { + t.Fatalf("unified audio features = %v", features.Dims()) + } +} + func TestAudioAttentionMaskValues(t *testing.T) { valid := []bool{true, true, true, false} got := audioAttentionMaskValues(valid, 2, 2, 4, 2, 0) diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index bfc3b4b3ac2..3ebc9578c49 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -840,15 +840,17 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { m.EmbedVision = embedVision } if m.AudioConfig != nil && m.AudioProcessorConfig != nil && hasCompleteGemma4AudioWeights(tensors, m.AudioConfig, m.HiddenSize) { - audio, err := loadAudioModel(tensors, m.AudioConfig, m.HiddenSize, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) - if err != nil { - return err + if !m.AudioConfig.unified() { + audio, err := loadAudioModel(tensors, m.AudioConfig, m.HiddenSize, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if err != nil { + return err + } + m.Audio = audio } embedAudio, err := loadMultimodalEmbedder(tensors, "embed_audio", m.AudioConfig.RMSNormEps, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) if err != nil { return err } - m.Audio = audio m.EmbedAudio = embedAudio } diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index 305a6b22374..ee53748752d 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -22,6 +22,9 @@ const ( ) type AudioConfig struct { + ModelType string `json:"model_type"` + AudioEmbedDim int `json:"audio_embed_dim"` + AudioSamplesPerToken int `json:"audio_samples_per_token"` AttentionChunkSize int `json:"attention_chunk_size"` AttentionContextLeft int `json:"attention_context_left"` AttentionContextRight int `json:"attention_context_right"` @@ -50,7 +53,7 @@ func ParseAudioConfig(configData []byte) (*AudioConfig, error) { if cfg.AudioConfig == nil { return nil, nil } - if err := validateAudioConfigFields(cfg.AudioConfig); err != nil { + if err := validateAudioConfig(cfg); err != nil { return nil, err } return cfg.AudioConfig, nil @@ -59,21 +62,23 @@ func ParseAudioConfig(configData []byte) (*AudioConfig, error) { type audioProcessorConfig struct { AudioSequenceLength int `json:"audio_seq_length"` FeatureExtractor struct { - Dither float64 `json:"dither"` - FeatureSize int `json:"feature_size"` - FFTLength int `json:"fft_length"` - FFTOverdrive bool `json:"fft_overdrive"` - FrameLength int `json:"frame_length"` - HopLength int `json:"hop_length"` - InputScaleFactor float64 `json:"input_scale_factor"` - MaxFrequency float64 `json:"max_frequency"` - MelFloor float64 `json:"mel_floor"` - MinFrequency float64 `json:"min_frequency"` - PaddingSide string `json:"padding_side"` - PerBinMean []float64 `json:"per_bin_mean"` - PerBinStddev []float64 `json:"per_bin_stddev"` - Preemphasis float64 `json:"preemphasis"` - SamplingRate int `json:"sampling_rate"` + Type string `json:"feature_extractor_type"` + AudioSamplesPerToken int `json:"audio_samples_per_token"` + Dither float64 `json:"dither"` + FeatureSize int `json:"feature_size"` + FFTLength int `json:"fft_length"` + FFTOverdrive bool `json:"fft_overdrive"` + FrameLength int `json:"frame_length"` + HopLength int `json:"hop_length"` + InputScaleFactor float64 `json:"input_scale_factor"` + MaxFrequency float64 `json:"max_frequency"` + MelFloor float64 `json:"mel_floor"` + MinFrequency float64 `json:"min_frequency"` + PaddingSide string `json:"padding_side"` + PerBinMean []float64 `json:"per_bin_mean"` + PerBinStddev []float64 `json:"per_bin_stddev"` + Preemphasis float64 `json:"preemphasis"` + SamplingRate int `json:"sampling_rate"` } `json:"feature_extractor"` } @@ -92,7 +97,13 @@ func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfig return fmt.Errorf("parse processor_config.json: %w", err) } f := processor.FeatureExtractor - if processor.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || + if isUnifiedAudioConfig(cfg) { + if processor.AudioSequenceLength != 750 || f.Type != "Gemma4UnifiedAudioFeatureExtractor" || + f.FeatureSize != 640 || f.SamplingRate != 16000 || f.AudioSamplesPerToken != 640 || + f.PaddingSide != "right" { + return fmt.Errorf("unsupported Gemma 4 unified audio processor configuration") + } + } else if processor.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || f.MelFloor != 1e-3 || f.Preemphasis != 0 || f.PaddingSide != "right" || @@ -213,6 +224,14 @@ func validateAudioConfig(cfg ConfigFile) error { if ac == nil { return fmt.Errorf("missing audio_config") } + if isUnifiedAudioConfig(cfg) { + if ac.AudioEmbedDim != 640 || ac.AudioSamplesPerToken != 640 || ac.HiddenSize != 640 || + ac.OutputProjDims != 640 || ac.RMSNormEps <= 0 || + cfg.TextConfig.HiddenSize <= 0 || cfg.TextConfig.HiddenSize > maxTextHiddenSize { + return fmt.Errorf("invalid Gemma 4 unified audio dimensions") + } + return nil + } if err := validateAudioConfigFields(ac); err != nil { return err } @@ -243,6 +262,19 @@ func validateAudioConfigFields(ac *AudioConfig) error { func requiredAudioShapes(cfg ConfigFile) (map[string][]int32, error) { ac := cfg.AudioConfig + if isUnifiedAudioConfig(cfg) { + textHidden, err := checkedAudioShapeDim("text hidden_size", int64(cfg.TextConfig.HiddenSize)) + if err != nil { + return nil, err + } + output, err := checkedAudioShapeDim("output_proj_dims", int64(ac.OutputProjDims)) + if err != nil { + return nil, err + } + return map[string][]int32{ + "model.embed_audio.embedding_projection.weight": {textHidden, output}, + }, nil + } hidden, err := checkedAudioShapeDim("hidden_size", int64(ac.HiddenSize)) if err != nil { return nil, err @@ -334,6 +366,10 @@ func requiredAudioShapes(cfg ConfigFile) (map[string][]int32, error) { return required, nil } +func isUnifiedAudioConfig(cfg ConfigFile) bool { + return cfg.AudioConfig != nil && cfg.AudioConfig.ModelType == "gemma4_unified_audio" +} + func checkedAudioShapeDim(name string, factors ...int64) (int32, error) { value, ok := checkedProduct(math.MaxInt32, factors...) if !ok { diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index 5df990a4958..47166f31952 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -23,6 +23,19 @@ func releasedAudioConfig(layers int) ConfigFile { } } +func releasedUnifiedAudioConfig() ConfigFile { + return ConfigFile{ + Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}, + ModelType: "gemma4_unified", + TextConfig: TextConfig{HiddenSize: 3840, VocabSize: 262144}, + AudioTokenID: 258881, + AudioConfig: &AudioConfig{ + ModelType: "gemma4_unified_audio", AudioEmbedDim: 640, + AudioSamplesPerToken: 640, HiddenSize: 640, OutputProjDims: 640, RMSNormEps: 1e-6, + }, + } +} + func TestParseAudioConfigEquivalence(t *testing.T) { valid, err := json.Marshal(releasedAudioConfig(12)) if err != nil { @@ -164,6 +177,38 @@ func TestValidateAudioTensorsExactNormalizedNames(t *testing.T) { } } +func TestValidateReleasedE2BAudioInventory(t *testing.T) { + cfg := releasedAudioConfig(12) + cfg.TextConfig.HiddenSize = 1536 + tensors := completeAudioInventory(cfg) + if got := tensors["model.embed_audio.embedding_projection.weight"].Shape; len(got) != 2 || got[0] != 1536 || got[1] != 1536 { + t.Fatalf("E2B audio projection shape = %v, want [1536 1536]", got) + } + if err := ValidateAudioSourceInventory(cfg, tensors); err != nil { + t.Fatalf("ValidateAudioSourceInventory() error = %v", err) + } +} + +func TestValidateReleasedUnifiedAudioInventory(t *testing.T) { + cfg := releasedUnifiedAudioConfig() + tensors := completeAudioInventory(cfg) + if got := len(tensors); got != 1 { + t.Fatalf("unified audio tensor count = %d, want 1", got) + } + if got := tensors["model.embed_audio.embedding_projection.weight"].Shape; len(got) != 2 || got[0] != 3840 || got[1] != 640 { + t.Fatalf("unified audio projection shape = %v, want [3840 640]", got) + } + if err := ValidateAudioSourceInventory(cfg, tensors); err != nil { + t.Fatalf("ValidateAudioSourceInventory() error = %v", err) + } + + bad := maps.Clone(tensors) + bad["model.embed_audio.embedding_projection.weight"] = TensorDescriptor{Dtype: "BF16", Shape: []int32{3840, 639}} + if err := ValidateAudioSourceInventory(cfg, bad); err == nil { + t.Fatal("bad unified projection shape: error = nil") + } +} + func TestValidateAudioConfig(t *testing.T) { tests := []struct { name string @@ -233,6 +278,29 @@ func TestAudioConfigSupportedBoundaries(t *testing.T) { } } +func TestValidateUnifiedAudioRuntimeMetadata(t *testing.T) { + processor := []byte(`{ + "audio_seq_length":750, + "feature_extractor":{ + "audio_samples_per_token":640, + "feature_extractor_type":"Gemma4UnifiedAudioFeatureExtractor", + "feature_size":640,"padding_side":"right","sampling_rate":16000 + } + }`) + tokens := []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`) + tokenizerData := []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":258881,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`) + cfg := releasedUnifiedAudioConfig() + if err := ValidateAudioRuntimeMetadata(cfg, processor, tokens, tokenizerData); err != nil { + t.Fatalf("ValidateAudioRuntimeMetadata() error = %v", err) + } + + badProcessor := []byte(string(processor)) + badProcessor = []byte(strings.Replace(string(badProcessor), `"audio_samples_per_token":640`, `"audio_samples_per_token":320`, 1)) + if err := ValidateAudioRuntimeMetadata(cfg, badProcessor, tokens, tokenizerData); err == nil { + t.Fatal("bad unified frame size: error = nil") + } +} + func TestAudioWithoutClippingScalars(t *testing.T) { cfg := releasedAudioConfig(1) cfg.AudioConfig.UseClippedLinears = false diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go index d020fc68d4f..ce90c4e3e75 100644 --- a/x/models/gemma4/vision.go +++ b/x/models/gemma4/vision.go @@ -574,7 +574,7 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas } payload.Image = &geom case "audio": - if m.AudioConfig == nil || m.AudioProcessorConfig == nil || m.Audio == nil || m.EmbedAudio == nil { + if m.AudioConfig == nil || m.AudioProcessorConfig == nil || m.EmbedAudio == nil || (!m.AudioConfig.unified() && m.Audio == nil) { return nil, fmt.Errorf("this model does not support audio input") } audio, err := preprocessGemma4Audio(ctx, seg.Data, m.AudioProcessorConfig) @@ -590,7 +590,7 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas prepared.Tokens = append(prepared.Tokens, m.EOATokenIDValue) geom := *audio mediaData = geom.Features - dims = []int{1, geom.Frames, 128} + dims = []int{1, geom.Frames, geom.FeatureSize} geom.Features = nil payload.Audio = &geom default: @@ -622,7 +622,13 @@ func (m *Model) PrepareMedia(ctx context.Context, segments []base.Segment) (*bas func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { payload := item.Opaque.(gemma4MediaPayload) if payload.Audio != nil { - features := m.EmbedAudio.Forward(m.Audio.Forward(data, payload.Audio)) + var features *mlx.Array + if m.AudioConfig.unified() { + raw := mlx.Reshape(data, 1, int32(payload.Audio.SoftTokens), int32(payload.Audio.FeatureSize)) + features = m.EmbedAudio.Forward(raw) + } else { + features = m.EmbedAudio.Forward(m.Audio.Forward(data, payload.Audio)) + } return mlx.Squeeze(features, 0) } var encoded *mlx.Array From 48629643cd9bf5abdf2b45b6c09e52c1c7f1e487 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 07:52:46 +0000 Subject: [PATCH 55/58] gemma4: tighten MLX audio architecture validation Reject mixed or unsupported Gemma 4 audio architecture metadata before capability reporting or runtime construction. Co-authored-by: Codex --- server/images_test.go | 3 +- server/model_inference_cache.go | 37 ++++++++++++ server/model_inference_cache_test.go | 75 +++++++++++++++++++++++++ x/create/client/create_test.go | 4 +- x/models/gemma4/audio_processor.go | 4 +- x/models/gemma4/audio_processor_test.go | 10 +++- x/models/gemma4/audio_test.go | 5 ++ x/models/gemma4/metadata/audio.go | 6 +- x/models/gemma4/metadata/audio_test.go | 5 +- 9 files changed, 142 insertions(+), 7 deletions(-) diff --git a/server/images_test.go b/server/images_test.go index 93e7ad37c9c..5a2cde8cd0e 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -928,7 +928,7 @@ func gemma4AudioManifestLayers(t *testing.T) []manifest.Layer { digest := createTestBlob(t, config) layers = append(layers, manifest.Layer{MediaType: "application/vnd.ollama.image.json", Digest: digest, Size: int64(len(config)), Name: "config.json"}) for name, data := range map[string][]byte{ - "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), + "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_extractor_type":"Gemma4AudioFeatureExtractor","feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), "tokenizer.json": []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), } { @@ -1324,6 +1324,7 @@ func gemma4AudioConfig(layers int) *gemma4metadata.ConfigFile { TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, diff --git a/server/model_inference_cache.go b/server/model_inference_cache.go index 7636e98640c..d08bc045846 100644 --- a/server/model_inference_cache.go +++ b/server/model_inference_cache.go @@ -9,6 +9,7 @@ import ( "github.com/ollama/ollama/envconfig" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/types/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" "golang.org/x/sync/singleflight" ) @@ -106,6 +107,10 @@ func cloneInferenceModel(src *Model) *Model { dst.AdapterPaths = slices.Clone(src.AdapterPaths) dst.ProjectorPaths = slices.Clone(src.ProjectorPaths) dst.TensorLayerNames = slices.Clone(src.TensorLayerNames) + dst.Gemma4VisionConfig = cloneGemma4MetadataConfig(src.Gemma4VisionConfig) + dst.Gemma4VisionTensors = cloneGemma4TensorDescriptors(src.Gemma4VisionTensors) + dst.Gemma4AudioConfig = cloneGemma4MetadataConfig(src.Gemma4AudioConfig) + dst.Gemma4AudioTensors = cloneGemma4TensorDescriptors(src.Gemma4AudioTensors) dst.License = slices.Clone(src.License) dst.Options = maps.Clone(src.Options) dst.Messages = slices.Clone(src.Messages) @@ -114,6 +119,38 @@ func cloneInferenceModel(src *Model) *Model { return &dst } +func cloneGemma4MetadataConfig(src *gemma4metadata.ConfigFile) *gemma4metadata.ConfigFile { + if src == nil { + return nil + } + + dst := *src + dst.Architectures = slices.Clone(src.Architectures) + if src.VisionConfig != nil { + vision := *src.VisionConfig + dst.VisionConfig = &vision + } + if src.AudioConfig != nil { + audio := *src.AudioConfig + audio.SubsamplingConvChannels = slices.Clone(src.AudioConfig.SubsamplingConvChannels) + dst.AudioConfig = &audio + } + return &dst +} + +func cloneGemma4TensorDescriptors(src map[string]gemma4metadata.TensorDescriptor) map[string]gemma4metadata.TensorDescriptor { + if src == nil { + return nil + } + + dst := make(map[string]gemma4metadata.TensorDescriptor, len(src)) + for name, descriptor := range src { + descriptor.Shape = slices.Clone(descriptor.Shape) + dst[name] = descriptor + } + return dst +} + func (s *Server) getModel(name string) (*Model, error) { if s != nil && s.modelCaches != nil && s.modelCaches.inference != nil { return s.modelCaches.inference.Get(name) diff --git a/server/model_inference_cache_test.go b/server/model_inference_cache_test.go index f841ef062ca..106fe8d0ccb 100644 --- a/server/model_inference_cache_test.go +++ b/server/model_inference_cache_test.go @@ -175,3 +175,78 @@ func TestInferenceModelCacheGemma4VisionTensorCapabilities(t *testing.T) { t.Fatalf("refreshed capabilities = %v, did not expect vision", third.Capabilities()) } } + +func TestInferenceModelCacheGemma4AudioCapabilities(t *testing.T) { + setTestHome(t, t.TempDir()) + + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "audio"}, + } + createSafetensorsTestModel(t, "gemma4-audio-cache", cfg, gemma4AudioManifestLayers(t)) + + cache := newInferenceModelCache() + loadCount := 0 + cache.loadModel = func(name string) (*Model, error) { + loadCount++ + return GetModel(name) + } + + first, err := cache.Get("gemma4-audio-cache") + if err != nil { + t.Fatal(err) + } + if !first.capabilitiesCached || !slices.Contains(first.Capabilities(), model.CapabilityAudio) || !first.Gemma4AudioReady { + t.Fatalf("cold state = capabilities:%v ready:%t, want cached audio", first.Capabilities(), first.Gemma4AudioReady) + } + if first.Gemma4AudioConfig == nil || len(first.Gemma4AudioTensors) == 0 { + t.Fatal("cold model did not retain Gemma 4 audio metadata") + } + first.Gemma4AudioConfig.AudioConfig.HiddenSize = 0 + mutatedTensor := false + for name, descriptor := range first.Gemma4AudioTensors { + if len(descriptor.Shape) == 0 { + continue + } + descriptor.Shape[0] = 0 + first.Gemma4AudioTensors[name] = descriptor + delete(first.Gemma4AudioTensors, name) + mutatedTensor = true + break + } + if !mutatedTensor { + t.Fatal("cold model did not retain a shaped Gemma 4 audio tensor") + } + + second, err := cache.Get("gemma4-audio-cache") + if err != nil { + t.Fatal(err) + } + if loadCount != 1 { + t.Fatalf("cache hit load count = %d, want 1", loadCount) + } + if !slices.Contains(second.Capabilities(), model.CapabilityAudio) || !second.Gemma4AudioReady { + t.Fatalf("cached state = capabilities:%v ready:%t, want audio", second.Capabilities(), second.Gemma4AudioReady) + } + if second.Gemma4AudioConfig.AudioConfig.HiddenSize == 0 || len(second.Gemma4AudioTensors) == 0 { + t.Fatal("cached Gemma 4 audio metadata was mutated") + } + for _, descriptor := range second.Gemma4AudioTensors { + if len(descriptor.Shape) > 0 && descriptor.Shape[0] == 0 { + t.Fatal("cached Gemma 4 audio tensor shape was mutated") + } + } + + createSafetensorsTestModel(t, "gemma4-audio-cache", cfg, nil) + third, err := cache.Get("gemma4-audio-cache") + if err != nil { + t.Fatal(err) + } + if loadCount != 2 { + t.Fatalf("invalidated load count = %d, want 2", loadCount) + } + if slices.Contains(third.Capabilities(), model.CapabilityAudio) || third.Gemma4AudioReady { + t.Fatalf("refreshed state = capabilities:%v ready:%t, did not expect audio", third.Capabilities(), third.Gemma4AudioReady) + } +} diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index eefdb653c1c..fd4e44a16d9 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -672,6 +672,7 @@ func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, @@ -752,6 +753,7 @@ func TestInferSafetensorsCapabilitiesGemma4AudioRejectsUnboundedConfig(t *testin TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, HiddenSize: 4, NumAttentionHeads: 2, @@ -863,7 +865,7 @@ func TestInferSafetensorsCapabilitiesGemma4UnifiedAudio(t *testing.T) { func writeGemma4AudioRuntimeConfigs(t *testing.T, dir string) { t.Helper() files := map[string][]byte{ - "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), + "processor_config.json": []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_extractor_type":"Gemma4AudioFeatureExtractor","feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`), "tokenizer_config.json": []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`), "tokenizer.json": []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), } diff --git a/x/models/gemma4/audio_processor.go b/x/models/gemma4/audio_processor.go index 1a751d2b836..e52ab2bf1c9 100644 --- a/x/models/gemma4/audio_processor.go +++ b/x/models/gemma4/audio_processor.go @@ -56,6 +56,7 @@ type gemma4AudioInput struct { func defaultAudioProcessorConfig() AudioProcessorConfig { var cfg AudioProcessorConfig cfg.AudioSequenceLength = 750 + cfg.FeatureExtractor.Type = "Gemma4AudioFeatureExtractor" cfg.FeatureExtractor.FeatureSize = 128 cfg.FeatureExtractor.FFTLength = 512 cfg.FeatureExtractor.FrameLength = 320 @@ -93,7 +94,8 @@ func validateReleasedAudioProcessorConfig(cfg *AudioProcessorConfig) error { } return nil } - if cfg.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || + if f.Type != "Gemma4AudioFeatureExtractor" || cfg.AudioSequenceLength != 750 || + f.FeatureSize != 128 || f.SamplingRate != 16000 || f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || f.MelFloor != 1e-3 || f.Preemphasis != 0 || f.PaddingSide != "right" || diff --git a/x/models/gemma4/audio_processor_test.go b/x/models/gemma4/audio_processor_test.go index 14b37eb5628..cb0c9d76451 100644 --- a/x/models/gemma4/audio_processor_test.go +++ b/x/models/gemma4/audio_processor_test.go @@ -13,8 +13,9 @@ import ( func TestParseReleasedAudioProcessorConfig(t *testing.T) { data := []byte(`{ - "audio_seq_length":750, - "feature_extractor":{ + "audio_seq_length":750, + "feature_extractor":{ + "feature_extractor_type":"Gemma4AudioFeatureExtractor", "dither":0.0,"feature_size":128,"fft_length":512,"fft_overdrive":false, "frame_length":320,"hop_length":160,"input_scale_factor":1.0, "max_frequency":8000.0,"mel_floor":0.001,"min_frequency":0.0, @@ -34,10 +35,15 @@ func TestParseReleasedAudioProcessorConfig(t *testing.T) { if _, err := parseAudioProcessorConfig(bad); err == nil { t.Fatal("1024-point FFT processor: error = nil") } + unknown := bytes.Replace(data, []byte(`"Gemma4AudioFeatureExtractor"`), []byte(`"FutureAudioFeatureExtractor"`), 1) + if _, err := parseAudioProcessorConfig(unknown); err == nil { + t.Fatal("unknown tower processor: error = nil") + } } func TestReleasedAudioProcessorRejectsMalformedDimensions(t *testing.T) { base := defaultAudioProcessorConfig() + base.FeatureExtractor.Type = "Gemma4AudioFeatureExtractor" tests := []struct { name string mutate func(*AudioProcessorConfig) diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index 500bff794ca..08a3d8209c2 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -24,6 +24,7 @@ const releasedGemma4AudioConfig = `{ "vocab_size": 262144 }, "audio_config": { + "model_type": "gemma4_audio", "attention_chunk_size": 12, "attention_context_left": 13, "attention_context_right": 0, @@ -61,6 +62,10 @@ func TestParseReleasedAudioConfig(t *testing.T) { if _, err := parseAudioConfig([]byte(bad)); err == nil { t.Fatal("non-divisible head count: error = nil") } + unknown := strings.Replace(releasedGemma4AudioConfig, `"model_type": "gemma4_audio"`, `"model_type": "future_audio"`, 1) + if _, err := parseAudioConfig([]byte(unknown)); err == nil { + t.Fatal("unknown audio model type: error = nil") + } } func TestParseAudioConfigUsesMetadataPredicate(t *testing.T) { diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index ee53748752d..c3bcf698b21 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -103,7 +103,8 @@ func ValidateAudioRuntimeMetadata(cfg ConfigFile, processorData, tokenizerConfig f.PaddingSide != "right" { return fmt.Errorf("unsupported Gemma 4 unified audio processor configuration") } - } else if processor.AudioSequenceLength != 750 || f.FeatureSize != 128 || f.SamplingRate != 16000 || + } else if f.Type != "Gemma4AudioFeatureExtractor" || processor.AudioSequenceLength != 750 || + f.FeatureSize != 128 || f.SamplingRate != 16000 || f.FrameLength != 320 || f.HopLength != 160 || f.FFTLength != 512 || f.FFTOverdrive || f.Dither != 0 || f.InputScaleFactor != 1 || f.MinFrequency != 0 || f.MaxFrequency != 8000 || f.MelFloor != 1e-3 || f.Preemphasis != 0 || f.PaddingSide != "right" || @@ -232,6 +233,9 @@ func validateAudioConfig(cfg ConfigFile) error { } return nil } + if ac.ModelType != "gemma4_audio" { + return fmt.Errorf("unsupported Gemma 4 audio model type %q", ac.ModelType) + } if err := validateAudioConfigFields(ac); err != nil { return err } diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index 47166f31952..c332069562b 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -13,6 +13,7 @@ func releasedAudioConfig(layers int) ConfigFile { TextConfig: TextConfig{HiddenSize: 2560, VocabSize: 262144}, AudioTokenID: 258881, AudioConfig: &AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 12, AttentionContextLeft: 13, AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 5, HiddenSize: 1024, NumAttentionHeads: 8, @@ -69,7 +70,7 @@ func TestParseAudioConfigEquivalence(t *testing.T) { } func TestValidateAudioRuntimeMetadata(t *testing.T) { - processor := []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`) + processor := []byte(`{"audio_seq_length":750,"feature_extractor":{"feature_extractor_type":"Gemma4AudioFeatureExtractor","feature_size":128,"fft_length":512,"frame_length":320,"hop_length":160,"input_scale_factor":1,"max_frequency":8000,"mel_floor":0.001,"padding_side":"right","sampling_rate":16000}}`) tokens := []byte(`{"boa_token":"<|audio>","audio_token":"<|audio|>","eoa_token":""}`) tokenizerData := []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":258881,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`) cfg := releasedAudioConfig(12) @@ -85,12 +86,14 @@ func TestValidateAudioRuntimeMetadata(t *testing.T) { }{ {"missing processor", nil, tokens, tokenizerData, nil}, {"unsupported processor", []byte(`{"audio_seq_length":749}`), tokens, tokenizerData, nil}, + {"unknown extractor type", []byte(strings.Replace(string(processor), "Gemma4AudioFeatureExtractor", "FutureAudioFeatureExtractor", 1)), tokens, tokenizerData, nil}, {"missing tokens", processor, nil, tokenizerData, nil}, {"incomplete tokens", processor, []byte(`{"audio_token":"<|audio|>"}`), tokenizerData, nil}, {"missing tokenizer", processor, tokens, nil, nil}, {"wrong tokenizer id", processor, tokens, []byte(`{"model":{"type":"BPE","vocab":{},"merges":[]},"added_tokens":[{"id":5,"content":"<|audio>","special":true},{"id":7,"content":"<|audio|>","special":true},{"id":6,"content":"","special":true}]}`), nil}, {"vocab membership is not singleton encoding", processor, tokens, []byte(`{"model":{"type":"BPE","vocab":{"<|audio>":0,"<|audio|>":1,"":2},"merges":[]},"added_tokens":[]}`), func(cfg *ConfigFile) { cfg.AudioTokenID = 1; cfg.TextConfig.VocabSize = 3 }}, {"invalid token id", processor, tokens, tokenizerData, func(cfg *ConfigFile) { cfg.AudioTokenID = cfg.TextConfig.VocabSize }}, + {"unknown model type", processor, tokens, tokenizerData, func(cfg *ConfigFile) { cfg.AudioConfig.ModelType = "future_audio" }}, } { t.Run(tt.name, func(t *testing.T) { candidate := releasedAudioConfig(12) From 0e49cf34ec3b79381663d1ff40b7790a9d71989e Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 07:53:21 +0000 Subject: [PATCH 56/58] gemma4: support ordered MLX media inputs Co-authored-by: Codex --- integration/audio_test.go | 211 ++++++++++++++++++++++++++++++++++++++ openai/openai.go | 34 +++--- openai/openai_test.go | 74 +++++++++++-- server/prompt_test.go | 48 +++++++++ 4 files changed, 345 insertions(+), 22 deletions(-) diff --git a/integration/audio_test.go b/integration/audio_test.go index 3e9b4f96e37..2ab762f95a9 100644 --- a/integration/audio_test.go +++ b/integration/audio_test.go @@ -6,6 +6,7 @@ import ( "bytes" "context" "encoding/base64" + "encoding/binary" "encoding/json" "fmt" "io" @@ -28,6 +29,38 @@ func decodeTestAudio(t *testing.T) api.ImageData { return data } +func silentTestAudio(t *testing.T) api.ImageData { + t.Helper() + const ( + sampleRate = 16_000 + samples = sampleRate / 2 + ) + dataSize := samples * 2 + var out bytes.Buffer + for _, value := range []any{ + []byte("RIFF"), uint32(36 + dataSize), []byte("WAVE"), + []byte("fmt "), uint32(16), uint16(1), uint16(1), + uint32(sampleRate), uint32(sampleRate * 2), uint16(2), uint16(16), + []byte("data"), uint32(dataSize), make([]byte, dataSize), + } { + if err := binary.Write(&out, binary.LittleEndian, value); err != nil { + t.Fatalf("encode silent WAV: %v", err) + } + } + return out.Bytes() +} + +func requireResponseContains(t *testing.T, response string, words ...string) { + t.Helper() + lower := strings.ToLower(response) + for _, word := range words { + if strings.Contains(lower, word) { + return + } + } + t.Fatalf("none of %v found in %q", words, response) +} + // setupAudioModel pulls the model, preloads it, and skips if it doesn't support audio. func setupAudioModel(ctx context.Context, t *testing.T, client *api.Client, model string) { t.Helper() @@ -250,3 +283,181 @@ func runOpenAIChatWithAudio(t *testing.T, models []string) { }) } } + +// TestGemma4MultipleMedia exercises ordered multi-audio, multi-image, mixed +// image/audio, OpenAI interleaving, and retained-history media through MLX. +func TestGemma4MultipleMedia(t *testing.T) { + models := testModels([]string{"gemma4:e2b"}) + for _, model := range models { + t.Run(model, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + defer cancel() + client, endpoint, cleanup := InitServerConnection(ctx, t) + defer cleanup() + + setupAudioModel(ctx, t, client, model) + requireCapability(ctx, t, client, model, "vision") + speech := decodeTestAudio(t) + silence := silentTestAudio(t) + abbeyRoad, docs, _ := decodeTestImages(t) + noThink := &api.ThinkValue{Value: false} + + for _, tc := range []struct { + name string + media []api.ImageData + }{ + {name: "speech_then_silence", media: []api.ImageData{speech, silence}}, + {name: "silence_then_speech", media: []api.ImageData{silence, speech}}, + } { + t.Run(tc.name, func(t *testing.T) { + req := api.ChatRequest{ + Model: model, + Think: noThink, + Messages: []api.Message{{ + Role: "user", + Content: "Two audio clips are attached. Transcribe only the clip containing speech.", + Images: tc.media, + }}, + Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 80}, + } + response := DoChat(ctx, t, client, req, []string{"sky", "blue"}, 90*time.Second, 20*time.Second) + requireResponseContains(t, response.Content, "sky", "blue") + }) + } + + for _, tc := range []struct { + name string + media []api.ImageData + firstWords []string + secondWords []string + }{ + { + name: "abbey_then_docs", media: []api.ImageData{abbeyRoad, docs}, + firstWords: []string{"road", "street", "cross", "walk", "beatles"}, + secondWords: []string{"laptop", "book", "read", "sleep", "documentation", "desk"}, + }, + { + name: "docs_then_abbey", media: []api.ImageData{docs, abbeyRoad}, + firstWords: []string{"laptop", "book", "read", "sleep", "documentation", "desk"}, + secondWords: []string{"road", "street", "cross", "walk", "beatles"}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + req := api.ChatRequest{ + Model: model, + Think: noThink, + Messages: []api.Message{{ + Role: "user", + Content: "Describe both pictures in order. Reply with exactly two labeled lines: " + + "FIRST: the first picture. SECOND: the second picture.", + Images: tc.media, + }}, + Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, + } + response := DoChat(ctx, t, client, req, append(tc.firstWords, tc.secondWords...), 120*time.Second, 20*time.Second) + requireResponseContains(t, response.Content, tc.firstWords...) + requireResponseContains(t, response.Content, tc.secondWords...) + }) + } + + t.Run("mixed_same_message", func(t *testing.T) { + req := api.ChatRequest{ + Model: model, + Think: noThink, + Messages: []api.Message{{ + Role: "user", + Content: "First [img] is a picture and second [img] is audio. Identify the picture subject and transcribe the spoken question.", + Images: []api.ImageData{docs, speech}, + }}, + Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, + } + response := DoChat(ctx, t, client, req, []string{"llama", "alpaca", "sky", "blue"}, 120*time.Second, 20*time.Second) + requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, response.Content, "sky", "blue") + }) + + t.Run("openai_mixed_same_message", func(t *testing.T) { + body, err := json.Marshal(map[string]any{ + "model": model, + "messages": []any{map[string]any{ + "role": "user", + "content": []any{ + map[string]any{"type": "text", "text": "First "}, + map[string]any{"type": "image_url", "image_url": map[string]any{ + "url": "data:image/png;base64," + base64.StdEncoding.EncodeToString(docs), + }}, + map[string]any{"type": "text", "text": " is a picture. Second "}, + map[string]any{"type": "input_audio", "input_audio": map[string]any{ + "data": base64.StdEncoding.EncodeToString(speech), "format": "wav", + }}, + map[string]any{"type": "text", "text": " is audio. Identify the picture subject and transcribe the spoken question."}, + }, + }}, + "temperature": 0, + "seed": 123, + "max_tokens": 200, + "reasoning_effort": "none", + }) + if err != nil { + t.Fatal(err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + fmt.Sprintf("http://%s/v1/chat/completions", endpoint), bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("OpenAI mixed-media request returned %s: %s", resp.Status, responseBody) + } + var result struct { + Choices []struct { + Message struct { + Content string `json:"content"` + Reasoning string `json:"reasoning"` + } `json:"message"` + } `json:"choices"` + } + if err := json.Unmarshal(responseBody, &result); err != nil { + t.Fatal(err) + } + if len(result.Choices) != 1 { + t.Fatalf("OpenAI mixed-media choices = %d, want 1", len(result.Choices)) + } + text := result.Choices[0].Message.Content + " " + result.Choices[0].Message.Reasoning + requireResponseContains(t, text, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, text, "sky", "blue") + }) + + t.Run("mixed_across_history", func(t *testing.T) { + req := api.ChatRequest{ + Model: model, + Think: noThink, + Messages: []api.Message{ + {Role: "user", Content: "Remember this picture.", Images: []api.ImageData{docs}}, + {Role: "assistant", Content: "I will retain the picture for the next instruction."}, + { + Role: "user", + Content: "Use both media inputs. Reply with exactly two labeled lines: " + + "AUDIO: the exact spoken question. IMAGE: the picture subject.", + Images: []api.ImageData{speech}, + }, + }, + Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, + } + response := DoChat(ctx, t, client, req, []string{"llama", "alpaca", "sky", "blue"}, 120*time.Second, 20*time.Second) + requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, response.Content, "sky", "blue") + }) + }) + } +} diff --git a/openai/openai.go b/openai/openai.go index 2d38607dbd5..76999243f2c 100644 --- a/openai/openai.go +++ b/openai/openai.go @@ -564,6 +564,18 @@ func FromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) { } messages = append(messages, api.Message{Role: msg.Role, Content: content, Thinking: msg.Reasoning, ToolCalls: toolCalls, ToolName: toolName, ToolCallID: msg.ToolCallID}) case []any: + toolCalls, err := FromCompletionToolCall(msg.ToolCalls) + if err != nil { + return nil, err + } + converted := api.Message{ + Role: msg.Role, + Thinking: msg.Reasoning, + ToolCalls: toolCalls, + ToolName: toolName, + ToolCallID: msg.ToolCallID, + } + var contentBuilder strings.Builder for _, c := range content { data, ok := c.(map[string]any) if !ok { @@ -575,7 +587,7 @@ func FromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) { if !ok { return nil, errors.New("invalid message format") } - messages = append(messages, api.Message{Role: msg.Role, Content: text}) + contentBuilder.WriteString(text) case "image_url": var url string if urlMap, ok := data["image_url"].(map[string]any); ok { @@ -593,7 +605,8 @@ func FromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) { return nil, err } - messages = append(messages, api.Message{Role: msg.Role, Images: []api.ImageData{img}}) + contentBuilder.WriteString("[img]") + converted.Images = append(converted.Images, img) case "input_audio": audioMap, ok := data["input_audio"].(map[string]any) if !ok { @@ -607,23 +620,14 @@ func FromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) { if err != nil { return nil, fmt.Errorf("invalid input_audio base64 data: %w", err) } - messages = append(messages, api.Message{Role: msg.Role, Images: []api.ImageData{audioBytes}}) + contentBuilder.WriteString("[img]") + converted.Images = append(converted.Images, audioBytes) default: return nil, errors.New("invalid message format") } } - // since we might have added multiple messages above, if we have tools - // calls we'll add them to the last message - if len(messages) > 0 && len(msg.ToolCalls) > 0 { - toolCalls, err := FromCompletionToolCall(msg.ToolCalls) - if err != nil { - return nil, err - } - messages[len(messages)-1].ToolCalls = toolCalls - messages[len(messages)-1].ToolName = toolName - messages[len(messages)-1].ToolCallID = msg.ToolCallID - messages[len(messages)-1].Thinking = msg.Reasoning - } + converted.Content = contentBuilder.String() + messages = append(messages, converted) default: // content is only optional if tool calls are present if msg.ToolCalls == nil { diff --git a/openai/openai_test.go b/openai/openai_test.go index 1f3b258b04b..4d250b81b27 100644 --- a/openai/openai_test.go +++ b/openai/openai_test.go @@ -134,23 +134,83 @@ func TestFromChatRequest_WithImage(t *testing.T) { t.Fatalf("unexpected error: %v", err) } - if len(result.Messages) != 2 { - t.Fatalf("expected 2 messages, got %d", len(result.Messages)) + if len(result.Messages) != 1 { + t.Fatalf("expected 1 message, got %d", len(result.Messages)) } - if result.Messages[0].Content != "Hello" { - t.Errorf("expected first message content 'Hello', got %q", result.Messages[0].Content) + if result.Messages[0].Content != "Hello[img]" { + t.Errorf("expected interleaved message content, got %q", result.Messages[0].Content) } - if len(result.Messages[1].Images) != 1 { - t.Fatalf("expected 1 image, got %d", len(result.Messages[1].Images)) + if len(result.Messages[0].Images) != 1 { + t.Fatalf("expected 1 image, got %d", len(result.Messages[0].Images)) } - if string(result.Messages[1].Images[0]) != string(imgData) { + if string(result.Messages[0].Images[0]) != string(imgData) { t.Error("image data mismatch") } } +func TestFromChatRequest_PreservesMixedMediaOrder(t *testing.T) { + imgData, _ := base64.StdEncoding.DecodeString(image) + audioOne := []byte("audio-one") + audioTwo := []byte("audio-two") + + req := ChatCompletionRequest{ + Model: "test-model", + Messages: []Message{ + { + Role: "user", + Content: []any{ + map[string]any{"type": "text", "text": "before"}, + map[string]any{ + "type": "image_url", + "image_url": map[string]any{"url": prefix + image}, + }, + map[string]any{ + "type": "input_audio", + "input_audio": map[string]any{ + "data": base64.StdEncoding.EncodeToString(audioOne), + "format": "wav", + }, + }, + map[string]any{"type": "text", "text": "between"}, + map[string]any{ + "type": "input_audio", + "input_audio": map[string]any{ + "data": base64.StdEncoding.EncodeToString(audioTwo), + "format": "wav", + }, + }, + map[string]any{"type": "text", "text": "after"}, + }, + }, + }, + } + + result, err := FromChatRequest(req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if got, want := len(result.Messages), 1; got != want { + t.Fatalf("len(Messages) = %d, want %d", got, want) + } + message := result.Messages[0] + if got, want := message.Content, "before[img][img]between[img]after"; got != want { + t.Fatalf("message content = %q, want %q", got, want) + } + wantMedia := [][]byte{imgData, audioOne, audioTwo} + if got, want := len(message.Images), len(wantMedia); got != want { + t.Fatalf("message has %d media items, want %d", got, want) + } + for i := range wantMedia { + if diff := cmp.Diff(wantMedia[i], []byte(message.Images[i])); diff != "" { + t.Fatalf("media %d mismatch (-want +got):\n%s", i, diff) + } + } +} + func TestFromCompleteRequest_Basic(t *testing.T) { temp := float32(0.8) req := CompletionRequest{ diff --git a/server/prompt_test.go b/server/prompt_test.go index 3939e06bc26..2e35474cf56 100644 --- a/server/prompt_test.go +++ b/server/prompt_test.go @@ -9,6 +9,7 @@ import ( "github.com/google/go-cmp/cmp" "github.com/ollama/ollama/api" + "github.com/ollama/ollama/llm" "github.com/ollama/ollama/template" "github.com/ollama/ollama/types/model" ) @@ -565,6 +566,53 @@ func TestChatPromptRendererPreservesExplicitImagePlaceholders(t *testing.T) { } } +func TestChatPromptRendererPreservesMixedMediaOrderAcrossHistory(t *testing.T) { + png := api.ImageData("\x89PNG\r\n\x1a\n") + wav := api.ImageData("RIFF\x00\x00\x00\x00WAVE") + secondPNG := api.ImageData("\x89PNG\r\n\x1a\nsecond") + msgs := []api.Message{ + { + Role: "user", + Content: "compare [img] with [img]", + Images: []api.ImageData{png, wav}, + }, + {Role: "assistant", Content: "I will remember both."}, + { + Role: "user", + Content: "and now [img]", + Images: []api.ImageData{secondPNG}, + }, + } + + m := Model{ + Config: model.ConfigV2{Renderer: "gemma4"}, + ProjectorPaths: []string{"media"}, + } + opts := api.Options{Runner: api.Runner{NumCtx: 8192}} + think := false + prompt, media, err := chatPrompt(t.Context(), &m, mockRunner{}.Tokenize, &opts, msgs, nil, &api.ThinkValue{Value: think}, true) + if err != nil { + t.Fatal(err) + } + + if !strings.Contains(prompt, "compare [img-0] with [img-1]") || !strings.Contains(prompt, "and now [img-2]") { + t.Fatalf("prompt does not preserve media marker order: %q", prompt) + } + if got, want := len(media), 3; got != want { + t.Fatalf("len(media) = %d, want %d", got, want) + } + if diff := cmp.Diff([]llm.MediaKind{llm.MediaKindImage, llm.MediaKindAudio, llm.MediaKindImage}, []llm.MediaKind{ + media[0].Kind, media[1].Kind, media[2].Kind, + }); diff != "" { + t.Fatalf("media kind order mismatch (-want +got):\n%s", diff) + } + for i, item := range media { + if item.ID != i { + t.Fatalf("media[%d].ID = %d, want %d", i, item.ID, i) + } + } +} + func TestRenderPromptResolvesDynamicGemma4Renderer(t *testing.T) { msgs := []api.Message{{Role: "user", Content: "Hello"}} From 1deb636828f9e71e9992fbdbf2b07d10100058f0 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 08:05:22 +0000 Subject: [PATCH 57/58] gemma4: complete quantized shared media support Extend the canonical Gemma 4 media descriptor contract to support quantized vision and audio tensors while preserving dense readiness and runtime validation. Co-authored-by: Codex --- integration/audio_test.go | 44 ++- x/create/create.go | 18 +- x/create/gemma4.go | 72 +++- x/create/gemma4_test.go | 132 +++++-- x/models/gemma4/audio.go | 64 ++-- .../gemma4/audio_forward_reference_test.go | 71 +++- x/models/gemma4/audio_test.go | 1 + x/models/gemma4/gemma4.go | 14 +- x/models/gemma4/media_quant.go | 106 ++++++ x/models/gemma4/media_quant_loader_test.go | 346 ++++++++++++++++++ x/models/gemma4/media_quant_test.go | 113 ++++++ x/models/gemma4/metadata/audio.go | 60 ++- x/models/gemma4/metadata/audio_test.go | 88 +++++ x/models/gemma4/metadata/vision.go | 93 ++++- x/models/gemma4/metadata/vision_test.go | 25 ++ x/models/gemma4/vision.go | 36 +- 16 files changed, 1148 insertions(+), 135 deletions(-) create mode 100644 x/models/gemma4/media_quant.go create mode 100644 x/models/gemma4/media_quant_loader_test.go create mode 100644 x/models/gemma4/media_quant_test.go diff --git a/integration/audio_test.go b/integration/audio_test.go index 2ab762f95a9..9b87201dcca 100644 --- a/integration/audio_test.go +++ b/integration/audio_test.go @@ -12,6 +12,7 @@ import ( "io" "mime/multipart" "net/http" + "regexp" "strings" "testing" "time" @@ -19,6 +20,14 @@ import ( "github.com/ollama/ollama/api" ) +var defaultAudioModels = []string{ + "nemotron3:33b", + "gemma4:e2b", + "gemma4:e4b", +} + +var catResponsePattern = regexp.MustCompile(`\bcats?\b`) + // decodeTestAudio returns the test audio clip ("Why is the sky blue?", 16kHz mono WAV). func decodeTestAudio(t *testing.T) api.ImageData { t.Helper() @@ -54,6 +63,12 @@ func requireResponseContains(t *testing.T, response string, words ...string) { t.Helper() lower := strings.ToLower(response) for _, word := range words { + if word == "cat" { + if catResponsePattern.MatchString(lower) { + return + } + continue + } if strings.Contains(lower, word) { return } @@ -61,6 +76,20 @@ func requireResponseContains(t *testing.T, response string, words ...string) { t.Fatalf("none of %v found in %q", words, response) } +func requireLabeledImageOrder(t *testing.T, response string, firstWords, secondWords []string) { + t.Helper() + lower := strings.ToLower(response) + firstAt := strings.Index(lower, "first:") + secondAt := strings.Index(lower, "second:") + if firstAt < 0 || secondAt <= firstAt { + t.Fatalf("response does not contain ordered FIRST:/SECOND: labels: %q", response) + } + first := response[firstAt+len("first:") : secondAt] + second := response[secondAt+len("second:"):] + requireResponseContains(t, first, firstWords...) + requireResponseContains(t, second, secondWords...) +} + // setupAudioModel pulls the model, preloads it, and skips if it doesn't support audio. func setupAudioModel(ctx context.Context, t *testing.T, client *api.Client, model string) { t.Helper() @@ -333,13 +362,13 @@ func TestGemma4MultipleMedia(t *testing.T) { }{ { name: "abbey_then_docs", media: []api.ImageData{abbeyRoad, docs}, - firstWords: []string{"road", "street", "cross", "walk", "beatles"}, - secondWords: []string{"laptop", "book", "read", "sleep", "documentation", "desk"}, + firstWords: []string{"road", "street", "cross", "walk", "beatles", "stripe", "ollamas"}, + secondWords: []string{"laptop", "book", "read", "sleep", "documentation", "document", "desk", "work", "study", "activity", "office"}, }, { name: "docs_then_abbey", media: []api.ImageData{docs, abbeyRoad}, firstWords: []string{"laptop", "book", "read", "sleep", "documentation", "desk"}, - secondWords: []string{"road", "street", "cross", "walk", "beatles"}, + secondWords: []string{"road", "street", "cross", "walk", "beatles", "stripe"}, }, } { t.Run(tc.name, func(t *testing.T) { @@ -355,8 +384,7 @@ func TestGemma4MultipleMedia(t *testing.T) { Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, } response := DoChat(ctx, t, client, req, append(tc.firstWords, tc.secondWords...), 120*time.Second, 20*time.Second) - requireResponseContains(t, response.Content, tc.firstWords...) - requireResponseContains(t, response.Content, tc.secondWords...) + requireLabeledImageOrder(t, response.Content, tc.firstWords, tc.secondWords) }) } @@ -372,7 +400,7 @@ func TestGemma4MultipleMedia(t *testing.T) { Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, } response := DoChat(ctx, t, client, req, []string{"llama", "alpaca", "sky", "blue"}, 120*time.Second, 20*time.Second) - requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character", "cat") requireResponseContains(t, response.Content, "sky", "blue") }) @@ -434,7 +462,7 @@ func TestGemma4MultipleMedia(t *testing.T) { t.Fatalf("OpenAI mixed-media choices = %d, want 1", len(result.Choices)) } text := result.Choices[0].Message.Content + " " + result.Choices[0].Message.Reasoning - requireResponseContains(t, text, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, text, "llama", "alpaca", "animal", "cartoon", "bear", "character", "cat") requireResponseContains(t, text, "sky", "blue") }) @@ -455,7 +483,7 @@ func TestGemma4MultipleMedia(t *testing.T) { Options: map[string]any{"temperature": 0, "seed": 123, "num_predict": 120}, } response := DoChat(ctx, t, client, req, []string{"llama", "alpaca", "sky", "blue"}, 120*time.Second, 20*time.Second) - requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character") + requireResponseContains(t, response.Content, "llama", "alpaca", "animal", "cartoon", "bear", "character", "cat") requireResponseContains(t, response.Content, "sky", "blue") }) }) diff --git a/x/create/create.go b/x/create/create.go index 4ba0fdff3bb..cf24959006c 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -267,6 +267,19 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string { if !stackedExpert && !ShouldQuantize(name, "") { return "" } + // Vision namespaces remain excluded by the generic policy. Architecture- + // specific policies may explicitly admit their known linear weights before + // calling the eligibility helper below. + if isVision(name) { + return "" + } + return getEligibleTensorQuantization(name, shape, quantize) +} + +// getEligibleTensorQuantization chooses the quantization for a weight after an +// architecture policy has established that its namespace is eligible. +func getEligibleTensorQuantization(name string, shape []int32, quantize string) string { + stackedExpert := isStackedExpertWeight(name) // Quantize standard linear weights (2D). Also allow stacked expert weights (3D), // e.g. qwen switch_mlp / experts combined tensors. @@ -300,11 +313,6 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string { return "" } - // Vision components are too quantization-sensitive; keep source precision. - if isVision(name) { - return "" - } - // MLX quantization requires last dimension to be divisible by group size. if !isAligned(shape, quantNorm) { return "" diff --git a/x/create/gemma4.go b/x/create/gemma4.go index 6f4f16316cd..4f4fdc7df15 100644 --- a/x/create/gemma4.go +++ b/x/create/gemma4.go @@ -7,15 +7,19 @@ import ( ) type gemma4ImportTransform struct { - numLayers int - numExperts int + numLayers int + numExperts int + unifiedAudio bool } // gemma4Config is a minimal subset of the Gemma 4 config.json used for quant decisions. type gemma4Config struct { NumHiddenLayers int `json:"num_hidden_layers"` NumExperts int `json:"num_experts"` - TextConfig struct { + AudioConfig *struct { + ModelType string `json:"model_type"` + } `json:"audio_config"` + TextConfig struct { NumHiddenLayers int `json:"num_hidden_layers"` NumExperts int `json:"num_experts"` } `json:"text_config"` @@ -36,12 +40,41 @@ func newGemma4ImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) numExperts = cfg.TextConfig.NumExperts } - return gemma4ImportTransform{numLayers: numLayers, numExperts: numExperts}, nil + return gemma4ImportTransform{ + numLayers: numLayers, + numExperts: numExperts, + unifiedAudio: cfg.AudioConfig != nil && cfg.AudioConfig.ModelType == "gemma4_unified_audio", + }, nil } func (t gemma4ImportTransform) quantizationType(name string, shape []int32, quantize string) string { base := normalizeQuantType(quantize) switch { + case isGemma4VisionTensor(name) || isGemma4AudioTensor(name): + if t.unifiedAudio && strings.HasSuffix(name, "embed_audio.embedding_projection.weight") { + // The unified architecture projects raw waveform blocks directly into + // language embeddings. Keep this semantic boundary dense. + return "" + } + // Media namespaces include positions, norms, convolutions, and learned + // biases. Only a conventional linear weight is eligible. + if !strings.HasSuffix(name, ".weight") || isGemma4MediaPositionTensor(name) { + return "" + } + eligible := getEligibleTensorQuantization(name, shape, quantize) + if eligible == "" { + return "" + } + if isGemma4MediaBoundaryProjection(name) { + // Four-bit media boundaries use their aligned eight-bit mate or stay + // dense. Never fall back to the requested four-bit representation. + eight := eightBit(base) + if isAligned(shape, eight) { + return eight + } + return "" + } + return eligible case isEmbedTokensWeight(name): // The embedding doubles as the lm_head projection; an 8-bit type keeps // quality close to bf16 (matching GGUF Q6_K) while saving bandwidth. @@ -62,12 +95,37 @@ func (t gemma4ImportTransform) quantizationType(name string, shape []int32, quan } } +func isGemma4MediaPositionTensor(name string) bool { + return strings.Contains(name, "position") || strings.Contains(name, "pos_embed") +} + +func isGemma4VisionTensor(name string) bool { + return isVision(name) +} + +func isGemma4AudioTensor(name string) bool { + return isAudioTower(name) || strings.Contains(name, "audio") +} + +func isGemma4MediaBoundaryProjection(name string) bool { + for _, suffix := range []string{ + "embed_vision.embedding_projection.weight", + "vision_embedder.patch_dense.weight", + "audio_tower.output_proj.weight", + "embed_audio.embedding_projection.weight", + } { + if strings.HasSuffix(name, suffix) { + return true + } + } + return false +} + // isSensitiveProjection reports the value/key/down projections whose precision // most affects quality — attention output (v/k) and the residual stream -// (down). Audio and vision tensors are excluded and follow the generic -// policy. +// (down). Non-text media tensors are handled separately. func (t gemma4ImportTransform) isSensitiveProjection(name string) bool { - if isVision(name) || isAudioTower(name) { + if isGemma4VisionTensor(name) || isGemma4AudioTensor(name) { return false } return strings.Contains(name, ".v_proj") || diff --git a/x/create/gemma4_test.go b/x/create/gemma4_test.go index aedc7cb8339..a801a0a5727 100644 --- a/x/create/gemma4_test.go +++ b/x/create/gemma4_test.go @@ -7,17 +7,19 @@ import ( func TestGemma4UnifiedImportTransformRegistration(t *testing.T) { tests := []struct { - name string - configJSON string - cfg sourceModelConfig - wantErr bool - wantLayers int + name string + configJSON string + cfg sourceModelConfig + wantErr bool + wantLayers int + wantUnifiedAudio bool }{ { - name: "unified conditional generation architecture", - configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48}}`, - cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}}, - wantLayers: 48, + name: "unified conditional generation architecture", + configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48},"audio_config":{"model_type":"gemma4_unified_audio"}}`, + cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}}, + wantLayers: 48, + wantUnifiedAudio: true, }, { name: "unified model type fallback", @@ -55,6 +57,9 @@ func TestGemma4UnifiedImportTransformRegistration(t *testing.T) { if gemmaTransform.numLayers != tt.wantLayers { t.Fatalf("numLayers = %d, want %d", gemmaTransform.numLayers, tt.wantLayers) } + if gemmaTransform.unifiedAudio != tt.wantUnifiedAudio { + t.Fatalf("unifiedAudio = %v, want %v", gemmaTransform.unifiedAudio, tt.wantUnifiedAudio) + } }) } } @@ -64,6 +69,7 @@ func TestGemma4QuantizationType(t *testing.T) { transform26B := gemma4ImportTransform{numLayers: 30, numExperts: 128} // 8-expert model (hypothetical) transform8E := gemma4ImportTransform{numLayers: 30, numExperts: 8} + transformUnified := gemma4ImportTransform{numLayers: 48, unifiedAudio: true} aligned := []int32{2816, 2816} // divisible by 64 (int4/int8 group size) and 16 (nvfp4) @@ -145,30 +151,38 @@ func TestGemma4QuantizationType(t *testing.T) { {"norm", transform26B, "model.layers.0.input_layernorm.weight", []int32{2816}, "int4", ""}, {"router scale", transform26B, "model.layers.0.router.scale", []int32{2816}, "int4", ""}, - // === Audio/vision tower tensors: must pass through unquantized for all quant types === - // These contain .v_proj and down_proj but should NOT be intercepted by - // the sensitive-tensor promotion logic. + // === Audio/vision tensors: only eligible linear weights quantize === {"audio norm int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int4", ""}, {"audio norm nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "nvfp4", ""}, {"audio norm int8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int8", ""}, {"audio norm mxfp8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "mxfp8", ""}, {"audio conv int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "int4", ""}, {"audio conv nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "nvfp4", ""}, - {"audio linear int4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "int4", ""}, - {"audio linear nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "nvfp4", ""}, - // Audio tower v_proj — must NOT be promoted despite containing .v_proj - {"audio v_proj int4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", ""}, - {"audio v_proj nvfp4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", ""}, - // Vision tower — source precision for every quant family. - {"vision v_proj int4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", ""}, - {"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", ""}, - {"vision q_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight", aligned, "nvfp4", ""}, - {"unified vision embedder nvfp4", transform26B, "model.vision_embedder.patch_dense.weight", aligned, "nvfp4", ""}, - {"vision projection nvfp4", transform26B, "model.embed_vision.embedding_projection.linear.weight", aligned, "nvfp4", ""}, - {"embed_vision int4", transform26B, "model.embed_vision.embedding_projection.weight", aligned, "int4", ""}, - // Audio tower down_proj - {"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", ""}, - {"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", ""}, + {"audio position table", transform26B, "model.audio_tower.position_embedding.weight", aligned, "mxfp8", ""}, + {"audio learned bias", transform26B, "model.audio_tower.output_proj.bias", aligned, "mxfp8", ""}, + {"vision position table", transform26B, "model.vision_tower.patch_embedder.position_embedding_table", aligned, "nvfp4", ""}, + {"unified vision position table", transform26B, "model.vision_embedder.pos_embedding", aligned, "nvfp4", ""}, + {"audio linear int4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "int4", "int4"}, + {"audio linear nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "nvfp4", "nvfp4"}, + {"audio linear mxfp4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "mxfp4", "mxfp4"}, + {"audio linear int8", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "int8", "int8"}, + {"audio linear mxfp8", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "mxfp8", "mxfp8"}, + {"audio v_proj int4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", "int8"}, + {"audio v_proj nvfp4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", "mxfp8"}, + {"vision v_proj int4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", "int8"}, + {"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", "mxfp8"}, + {"vision q_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight", aligned, "nvfp4", "nvfp4"}, + {"small media linear", transform26B, "model.audio_tower.small.weight", []int32{16, 16}, "int4", ""}, + {"misaligned media linear", transform26B, "model.audio_tower.proj.weight", []int32{64, 65}, "int4", ""}, + {"embed_vision boundary int4", transform26B, "model.embed_vision.embedding_projection.weight", aligned, "int4", "int8"}, + {"embed_vision boundary nvfp4", transform26B, "model.embed_vision.embedding_projection.weight", aligned, "nvfp4", "mxfp8"}, + {"unified vision boundary mxfp4", transform26B, "model.vision_embedder.patch_dense.weight", aligned, "mxfp4", "mxfp8"}, + {"audio boundary int8", transform26B, "model.audio_tower.output_proj.weight", aligned, "int8", "int8"}, + {"unified audio boundary mxfp8", transform26B, "model.embed_audio.embedding_projection.weight", aligned, "mxfp8", "mxfp8"}, + {"boundary four-bit mate misaligned", transform26B, "model.embed_vision.embedding_projection.weight", []int32{128, 16}, "nvfp4", ""}, + {"unified direct audio stays dense", transformUnified, "model.embed_audio.embedding_projection.weight", aligned, "nvfp4", ""}, + {"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", "int8"}, + {"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", "mxfp8"}, } for _, tt := range tests { @@ -182,7 +196,7 @@ func TestGemma4QuantizationType(t *testing.T) { } } -func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { +func TestGemma4ImportPlanQuantizesMediaLinears(t *testing.T) { policy := gemma4ImportTransform{numLayers: 2} inv := newInventory(sourceModelConfig{}, map[string]string{ "model.embed_tokens.weight": "BF16", @@ -192,6 +206,7 @@ func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { "model.embed_vision.embedding_projection.weight": "BF16", "model.vision_embedder.patch_dense.weight": "BF16", "model.audio_tower.subsample_conv_projection.input_proj_linear.weight": "BF16", + "model.audio_tower.output_proj.weight": "BF16", "model.embed_audio.embedding_projection.weight": "BF16", }) @@ -206,20 +221,61 @@ func TestGemma4ImportPlanKeepsMediaAtSourcePrecision(t *testing.T) { got[tensor.Name] = tensor } } - for _, name := range []string{ - "model.vision_tower.patch_embedder.input_proj.weight", - "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", - "model.embed_vision.embedding_projection.weight", - "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", - "model.embed_audio.embedding_projection.weight", - "model.vision_embedder.patch_dense.weight", - } { + want := map[string]string{ + "model.vision_tower.patch_embedder.input_proj.weight": "nvfp4", + "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight": "mxfp8", + "model.embed_vision.embedding_projection.weight": "mxfp8", + "model.vision_embedder.patch_dense.weight": "mxfp8", + "model.audio_tower.subsample_conv_projection.input_proj_linear.weight": "nvfp4", + "model.audio_tower.output_proj.weight": "mxfp8", + "model.embed_audio.embedding_projection.weight": "mxfp8", + } + for name, quantize := range want { tensor, ok := got[name] if !ok { t.Fatalf("%s missing from plan; got %v", name, specNames(specs)) } - if tensor.Quantize != "" { - t.Fatalf("%s Quantize = %q, want source precision", name, tensor.Quantize) + if tensor.Quantize != quantize { + t.Fatalf("%s Quantize = %q, want %q", name, tensor.Quantize, quantize) + } + } +} + +func TestGemma4UnifiedAudioImportPlanRetainsDirectProjection(t *testing.T) { + policy := gemma4ImportTransform{numLayers: 48, unifiedAudio: true} + inv := newInventory(sourceModelConfig{}, map[string]string{ + "model.embed_tokens.weight": "BF16", + "model.embed_audio.embedding_projection.weight": "BF16", + "model.audio_tower.output_proj.weight": "BF16", + }) + + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "nvfp4"}, policy) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + got := make(map[string]string) + for _, spec := range specs { + for _, tensor := range spec.Tensors { + got[tensor.Name] = tensor.Quantize + } + } + if q := got["model.embed_audio.embedding_projection.weight"]; q != "" { + t.Fatalf("unified direct audio Quantize = %q, want dense", q) + } + if q := got["model.audio_tower.output_proj.weight"]; q != "mxfp8" { + t.Fatalf("ordinary audio boundary Quantize = %q, want mxfp8", q) + } +} + +func TestGenericQuantizationStillExcludesMediaNamespaces(t *testing.T) { + shape := []int32{2048, 2048} + for _, name := range []string{ + "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight", + "model.audio_tower.layers.0.self_attn.q_proj.weight", + "model.embed_audio.embedding_projection.weight", + } { + if got := GetTensorQuantization(name, shape, "mxfp8"); got != "" { + t.Fatalf("GetTensorQuantization(%q) = %q, want generic namespace exclusion", name, got) } } } diff --git a/x/models/gemma4/audio.go b/x/models/gemma4/audio.go index 6d1d0762f9a..5ce6b253e11 100644 --- a/x/models/gemma4/audio.go +++ b/x/models/gemma4/audio.go @@ -109,17 +109,36 @@ type AudioModel struct { Config *AudioConfig } -func hasCompleteGemma4AudioWeights(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32) bool { +func validateGemma4AudioReadiness( + tensors map[string]*mlx.Array, + cfg *AudioConfig, + textHidden int32, + groupSize, bits int, + mode string, + tq map[string]*model.TensorQuantInfo, +) (bool, error) { if cfg == nil { - return false + return false, nil + } + sentinel := tensors["model.audio_tower.subsample_conv_projection.layer0.conv.weight"] + packedSentinel := tensors["model.audio_tower.subsample_conv_projection.layer0.conv.weight_packed"] + if cfg.unified() { + sentinel = tensors["model.embed_audio.embedding_projection.weight"] + packedSentinel = tensors["model.embed_audio.embedding_projection.weight_packed"] } - names := make([]string, 0, len(tensors)) - for name, tensor := range tensors { - if tensor != nil { - names = append(names, name) + if sentinel == nil { + if packedSentinel != nil { + return false, fmt.Errorf("runtime contains source-only Gemma4 audio packed sentinel") } + return false, nil + } + metadataCfg := audioMetadataConfig(cfg, textHidden) + metadataCfg.QuantizationConfig = gemma4metadata.Quantization{Bits: bits, GroupSize: groupSize, Mode: mode} + descriptors := gemma4RuntimeTensorDescriptors(tensors, tq, groupSize, mode) + if err := gemma4metadata.ValidateAudioRuntimeInventory(metadataCfg, descriptors); err != nil { + return false, err } - return gemma4metadata.ValidateAudioTensors(audioMetadataConfig(cfg, textHidden), names) == nil + return true, nil } func audioMetadataConfig(cfg *AudioConfig, textHidden int32) gemma4metadata.ConfigFile { @@ -145,27 +164,8 @@ func audioMetadataConfig(cfg *AudioConfig, textHidden int32) gemma4metadata.Conf } func validateGemma4AudioWeights(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32) error { - required, err := gemma4metadata.RequiredAudioTensorShapes(audioMetadataConfig(cfg, textHidden)) - if err != nil { - return err - } - for name, shape := range required { - tensor := tensors[name] - if tensor == nil { - return fmt.Errorf("missing Gemma4 audio tensor %s", name) - } - want := make([]int, len(shape)) - for i, dim := range shape { - want[i] = int(dim) - } - if !equalIntShape(tensor.Dims(), want) { - return fmt.Errorf("Gemma4 audio tensor %s shape %v, want %v", name, tensor.Dims(), want) - } - if !supportedGemma4AudioDType(tensor.DType()) { - return fmt.Errorf("Gemma4 audio tensor %s has unsupported dtype %s", name, tensor.DType()) - } - } - return nil + _, err := validateGemma4AudioReadiness(tensors, cfg, textHidden, 0, 0, "", nil) + return err } func supportedGemma4AudioDType(dtype mlx.DType) bool { @@ -190,9 +190,13 @@ func equalIntShape(a, b []int) bool { } func loadAudioModel(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden int32, groupSize, bits int, mode string, tq map[string]*model.TensorQuantInfo) (*AudioModel, error) { - if err := validateGemma4AudioWeights(tensors, cfg, textHidden); err != nil { + ready, err := validateGemma4AudioReadiness(tensors, cfg, textHidden, groupSize, bits, mode, tq) + if err != nil { return nil, err } + if !ready { + return nil, errors.New("missing Gemma4 audio tensors") + } const prefix = "model.audio_tower." linears := model.NewLinearFactory(tensors, groupSize, bits, mode, tq) loadConv := func(path string) (*audioConvBlock, error) { @@ -295,7 +299,7 @@ func loadAudioModel(tensors map[string]*mlx.Array, cfg *AudioConfig, textHidden FeedForward1: ff1, FeedForward2: ff2, Attention: &audioAttention{ Q: q, K: k, V: v, Output: attnOut, RelativeK: relativeK, - RelativeKDType: tensors[path+"self_attn.relative_k_proj.weight"].DType(), + RelativeKDType: gemma4LinearComputeDType(tensors, path+"self_attn.relative_k_proj"), PerDimScale: perDimScale, Config: cfg, }, LightConv: &audioLightConv{PreNorm: convPre, ConvNorm: convNorm, Start: convStart, End: convEnd, DepthwiseWeight: depthwise, Config: cfg}, diff --git a/x/models/gemma4/audio_forward_reference_test.go b/x/models/gemma4/audio_forward_reference_test.go index 9fe14216e24..daf4b12ee87 100644 --- a/x/models/gemma4/audio_forward_reference_test.go +++ b/x/models/gemma4/audio_forward_reference_test.go @@ -15,6 +15,11 @@ import ( gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" ) +const ( + gemma4LanguageLogitMaxDiff = 1.25 + gemma4LanguageLogitRMSE = 0.25 +) + func TestAudioForwardReference(t *testing.T) { modelDir := os.Getenv("GEMMA4_AUDIO_MODEL_DIR") refDir := os.Getenv("GEMMA4_AUDIO_REF_DIR") @@ -167,6 +172,9 @@ func TestAudioForwardReference(t *testing.T) { if importedModel := os.Getenv("GEMMA4_AUDIO_MODEL_NAME"); importedModel != "" { fullModel := loadImportedGemma4ReferenceModel(t, importedModel) + // The Python checkpoint records raw model logits before generation-time + // token suppression. Keep this comparison at the same boundary. + fullModel.SuppressLogitBias = nil payload := gemma4MediaPayload{Audio: input, AudioStart: 1, AudioEnd: 1 + input.SoftTokens} modelBatch := &batch.Batch{ InputIDs: inputIDs, @@ -180,7 +188,11 @@ func TestAudioForwardReference(t *testing.T) { logits := fullModel.Unembed(hidden) last := mlx.SliceStartStop(logits, []int32{0, int32(inputIDs.Dim(1) - 1), 0}, []int32{1, int32(inputIDs.Dim(1)), textConfig.VocabSize}) wantLogits := reference.Get("prefill_logits") - compareAudioReference(t, "prefill_logits", last, wantLogits, 0.5, 0.05) + if audioConfig.unified() { + compareLanguageLogitReference(t, "prefill_logits", last, wantLogits) + } else { + compareAudioReference(t, "prefill_logits", last, wantLogits, 0.5, 0.05) + } gotID := last.Argmax(-1, false).AsType(mlx.DTypeInt32) wantID := wantLogits.Argmax(-1, false).AsType(mlx.DTypeInt32) mlx.Eval(gotID, wantID) @@ -255,6 +267,34 @@ func compareAudioMaskReference(t *testing.T, name string, got []bool, want *mlx. } func compareAudioReference(t *testing.T, name string, got, want *mlx.Array, atol, rtol float64) { + t.Helper() + stats := measureAudioReference(t, name, got, want, atol, rtol) + if stats.firstMismatch >= 0 { + t.Errorf("%s[%d] = %g, want %g (diff %g, tolerance %g); max absolute difference %g, mean absolute difference %g, RMSE %g", name, stats.firstMismatch, stats.gotFirst, stats.wantFirst, stats.firstDiff, stats.firstTolerance, stats.maxDiff, stats.meanAbsDiff, stats.rmse) + return + } + t.Logf("%s matched %d values; max absolute difference %g, mean absolute difference %g, RMSE %g", name, stats.count, stats.maxDiff, stats.meanAbsDiff, stats.rmse) +} + +func compareLanguageLogitReference(t *testing.T, name string, got, want *mlx.Array) { + t.Helper() + stats := measureAudioReference(t, name, got, want, math.Inf(1), 0) + if stats.maxDiff > gemma4LanguageLogitMaxDiff || stats.rmse > gemma4LanguageLogitRMSE { + t.Errorf("%s exceeded language-logit bounds: max absolute difference %g (limit %g), mean absolute difference %g, RMSE %g (limit %g)", name, stats.maxDiff, gemma4LanguageLogitMaxDiff, stats.meanAbsDiff, stats.rmse, gemma4LanguageLogitRMSE) + return + } + t.Logf("%s matched language-logit bounds across %d values; max absolute difference %g, mean absolute difference %g, RMSE %g", name, stats.count, stats.maxDiff, stats.meanAbsDiff, stats.rmse) +} + +type audioReferenceStats struct { + count int + firstMismatch int + gotFirst, wantFirst float32 + firstDiff, firstTolerance float64 + maxDiff, meanAbsDiff, rmse float64 +} + +func measureAudioReference(t *testing.T, name string, got, want *mlx.Array, atol, rtol float64) audioReferenceStats { t.Helper() if got == nil || want == nil { t.Fatalf("%s tensor is missing", name) @@ -266,16 +306,31 @@ func compareAudioReference(t *testing.T, name string, got, want *mlx.Array, atol want = want.AsType(mlx.DTypeFloat32) mlx.Eval(got, want) gotValues, wantValues := got.Floats(), want.Floats() - var maxDiff float64 + stats := audioReferenceStats{count: len(wantValues), firstMismatch: -1} + var sumAbs, sumSquared float64 for i := range wantValues { - diff := math.Abs(float64(gotValues[i] - wantValues[i])) - if diff > maxDiff { - maxDiff = diff + gotValue, wantValue := float64(gotValues[i]), float64(wantValues[i]) + if math.IsNaN(gotValue) || math.IsInf(gotValue, 0) || math.IsNaN(wantValue) || math.IsInf(wantValue, 0) { + t.Fatalf("%s[%d] contains a non-finite value: got %g, want %g", name, i, gotValue, wantValue) + } + diff := math.Abs(gotValue - wantValue) + if diff > stats.maxDiff { + stats.maxDiff = diff } + sumAbs += diff + sumSquared += diff * diff tolerance := atol + rtol*math.Abs(float64(wantValues[i])) - if diff > tolerance { - t.Fatalf("%s[%d] = %g, want %g (diff %g, tolerance %g)", name, i, gotValues[i], wantValues[i], diff, tolerance) + if diff > tolerance && stats.firstMismatch < 0 { + stats.firstMismatch = i + stats.gotFirst = gotValues[i] + stats.wantFirst = wantValues[i] + stats.firstDiff = diff + stats.firstTolerance = tolerance } } - t.Logf("%s matched %d values; max absolute difference %g", name, len(wantValues), maxDiff) + if stats.count > 0 { + stats.meanAbsDiff = sumAbs / float64(stats.count) + stats.rmse = math.Sqrt(sumSquared / float64(stats.count)) + } + return stats } diff --git a/x/models/gemma4/audio_test.go b/x/models/gemma4/audio_test.go index 08a3d8209c2..637cc4cd655 100644 --- a/x/models/gemma4/audio_test.go +++ b/x/models/gemma4/audio_test.go @@ -518,6 +518,7 @@ func testGemma4Root(t *testing.T, configData, tokenizerData []byte, extra map[st func tinyAudioConfig() *AudioConfig { return &AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, AttentionContextRight: 0, AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, GradientClipping: 1e4, HiddenSize: 4, NumAttentionHeads: 2, diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index 3ebc9578c49..1d9aa8fd257 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -815,7 +815,10 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { m.LMHead = m.EmbedTokens.AsLinear() } - visionReady, err := validateGemma4VisionWeights(tensors, m.VisionConfig, int(m.HiddenSize), m.TensorQuant) + visionReady, err := validateGemma4VisionWeights( + tensors, m.VisionConfig, int(m.HiddenSize), + m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant, + ) if err != nil { return fmt.Errorf("invalid Gemma4 vision tensors: %w", err) } @@ -839,7 +842,14 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { } m.EmbedVision = embedVision } - if m.AudioConfig != nil && m.AudioProcessorConfig != nil && hasCompleteGemma4AudioWeights(tensors, m.AudioConfig, m.HiddenSize) { + audioReady, err := validateGemma4AudioReadiness( + tensors, m.AudioConfig, m.HiddenSize, + m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant, + ) + if err != nil { + return fmt.Errorf("invalid Gemma4 audio tensors: %w", err) + } + if m.AudioProcessorConfig != nil && audioReady { if !m.AudioConfig.unified() { audio, err := loadAudioModel(tensors, m.AudioConfig, m.HiddenSize, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) if err != nil { diff --git a/x/models/gemma4/media_quant.go b/x/models/gemma4/media_quant.go new file mode 100644 index 00000000000..3dcc2b1219a --- /dev/null +++ b/x/models/gemma4/media_quant.go @@ -0,0 +1,106 @@ +package gemma4 + +import ( + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" +) + +func validateGemma4LinearWeight( + tensors map[string]*mlx.Array, + path string, + want []int, + defaultGroupSize, defaultBits int, + defaultMode string, + tensorQuant map[string]*model.TensorQuantInfo, +) error { + weightName := path + ".weight" + weight := tensors[weightName] + scales := tensors[weightName+"_scale"] + groupSize, bits, mode := defaultGroupSize, defaultBits, defaultMode + if weight != nil && scales != nil { + groupSize, bits, mode = model.ResolveLinearQuantParams( + defaultGroupSize, defaultBits, defaultMode, + tensorQuant, weightName, weight, scales, + ) + } + descriptors := gemma4RuntimeTensorDescriptors(tensors, tensorQuant, defaultGroupSize, defaultMode) + if descriptor, ok := descriptors[weightName]; ok && descriptor.QuantType == "" && scales != nil { + descriptor.QuantType = mode + descriptor.GroupSize = groupSize + descriptors[weightName] = descriptor + } + logical := make([]int32, len(want)) + for i, dim := range want { + logical[i] = int32(dim) + } + cfg := gemma4metadata.ConfigFile{ + QuantizationConfig: gemma4metadata.Quantization{Bits: bits, GroupSize: groupSize, Mode: mode}, + } + return gemma4metadata.ValidateRuntimeLinearDescriptor(cfg, descriptors, path, logical) +} + +func validateGemma4PackedLinearShape(weight, scales, want []int, groupSize, bits int, mode string) error { + descriptors := map[string]gemma4metadata.TensorDescriptor{ + "test.weight": {Dtype: "U32", Shape: intShape32(weight), GroupSize: groupSize}, + "test.weight_scale": {Dtype: "U8", Shape: intShape32(scales)}, + } + if mode == "affine" { + descriptors["test.weight_scale"] = gemma4metadata.TensorDescriptor{Dtype: "BF16", Shape: intShape32(scales)} + descriptors["test.weight_qbias"] = gemma4metadata.TensorDescriptor{Dtype: "BF16", Shape: intShape32(scales)} + } + return gemma4metadata.ValidateRuntimeLinearDescriptor( + gemma4metadata.ConfigFile{QuantizationConfig: gemma4metadata.Quantization{Bits: bits, GroupSize: groupSize, Mode: mode}}, + descriptors, "test", intShape32(want), + ) +} + +func intShape32(shape []int) []int32 { + out := make([]int32, len(shape)) + for i, dim := range shape { + out[i] = int32(dim) + } + return out +} + +func gemma4RuntimeTensorDescriptors( + tensors map[string]*mlx.Array, + tensorQuant map[string]*model.TensorQuantInfo, + defaultGroupSize int, + defaultMode string, +) map[string]gemma4metadata.TensorDescriptor { + descriptors := make(map[string]gemma4metadata.TensorDescriptor, len(tensors)) + for name, tensor := range tensors { + if tensor == nil { + continue + } + descriptor := gemma4metadata.TensorDescriptor{Dtype: tensor.DType().String(), Shape: intShape32(tensor.Dims())} + if quant := tensorQuant[name]; quant != nil { + descriptor.QuantType = quant.QuantType + descriptor.GroupSize = quant.GroupSize + } else if tensor.DType() == mlx.DTypeUint32 && defaultMode != "" { + descriptor.QuantType = defaultMode + descriptor.GroupSize = defaultGroupSize + } + descriptors[name] = descriptor + } + return descriptors +} + +func isGemma4FloatingDType(dtype mlx.DType) bool { + switch dtype { + case mlx.DTypeBFloat16, mlx.DTypeFloat16, mlx.DTypeFloat32: + return true + default: + return false + } +} + +func gemma4LinearComputeDType(tensors map[string]*mlx.Array, path string) mlx.DType { + weightName := path + ".weight" + weight := tensors[weightName] + if weight != nil && tensors[weightName+"_scale"] == nil && isGemma4FloatingDType(weight.DType()) { + return weight.DType() + } + return mlx.DTypeBFloat16 +} diff --git a/x/models/gemma4/media_quant_loader_test.go b/x/models/gemma4/media_quant_loader_test.go new file mode 100644 index 00000000000..bc42bfbe9d0 --- /dev/null +++ b/x/models/gemma4/media_quant_loader_test.go @@ -0,0 +1,346 @@ +package gemma4 + +import ( + "fmt" + "strings" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + gemma4metadata "github.com/ollama/ollama/x/models/gemma4/metadata" + "github.com/ollama/ollama/x/models/nn" +) + +func TestGemma4TowerMediaQuantizedLoadWeights(t *testing.T) { + for _, quantized := range []bool{false, true} { + name := "dense" + if quantized { + name = "mixed_nvfp4" + } + t.Run(name, func(t *testing.T) { + useMLXTestThread(t) + m, tensors := testGemma4TowerMediaModel(t, quantized) + if err := m.LoadWeights(tensors); err != nil { + t.Fatal(err) + } + if m.Vision == nil || m.EmbedVision == nil || m.Audio == nil || m.EmbedAudio == nil { + t.Fatalf("media loaders were suppressed: vision=%v embed_vision=%v audio=%v embed_audio=%v", m.Vision != nil, m.EmbedVision != nil, m.Audio != nil, m.EmbedAudio != nil) + } + + assertGemma4LinearStorage(t, m.Vision.PatchEmbedder.InputProj, quantized) + assertGemma4LinearStorage(t, m.EmbedVision.Projection, quantized) + for _, layer := range m.Vision.Layers { + for _, linear := range []nn.LinearLayer{ + layer.Attention.QProj.Linear, layer.Attention.KProj.Linear, + layer.Attention.VProj.Linear, layer.Attention.OProj.Linear, + layer.MLP.GateProj.Linear, layer.MLP.UpProj.Linear, layer.MLP.DownProj.Linear, + } { + assertGemma4LinearStorage(t, linear, quantized) + } + } + + assertGemma4LinearStorage(t, m.Audio.InputProj, quantized) + assertGemma4LinearStorage(t, m.Audio.OutputProj, quantized) + assertGemma4LinearStorage(t, m.EmbedAudio.Projection, quantized) + if quantized { + assertGemma4QuantizedLinearType(t, m.EmbedVision.Projection, "mxfp8") + assertGemma4QuantizedLinearType(t, m.Audio.OutputProj, "mxfp8") + assertGemma4QuantizedLinearType(t, m.EmbedAudio.Projection, "mxfp8") + } + for _, layer := range m.Audio.Layers { + for _, linear := range []nn.LinearLayer{ + layer.FeedForward1.Up.Linear, layer.FeedForward1.Down.Linear, + layer.FeedForward2.Up.Linear, layer.FeedForward2.Down.Linear, + layer.Attention.Q.Linear, layer.Attention.K.Linear, layer.Attention.V.Linear, + layer.Attention.Output.Linear, layer.Attention.RelativeK, + layer.LightConv.Start.Linear, layer.LightConv.End.Linear, + } { + assertGemma4LinearStorage(t, linear, quantized) + } + } + + if got := m.Vision.PatchEmbedder.PositionEmbeddingTable; got != tensors["model.vision_tower.patch_embedder.position_embedding_table"] { + t.Fatal("vision position table was not retained as a dense tensor") + } + if quantized { + output, ok := m.Audio.OutputProj.(*nn.QuantizedLinear) + if !ok || output.Bias != tensors["model.audio_tower.output_proj.bias"] { + t.Fatal("quantized audio output projection lost its learned bias") + } + } + }) + } +} + +func TestGemma4UnifiedMediaQuantizedLoadWeights(t *testing.T) { + for _, quantized := range []bool{false, true} { + name := "dense" + if quantized { + name = "nvfp4" + } + t.Run(name, func(t *testing.T) { + useMLXTestThread(t) + m, tensors := testGemma4UnifiedMediaModel(quantized) + if err := m.LoadWeights(tensors); err != nil { + t.Fatal(err) + } + if m.UnifiedVision == nil || m.EmbedVision == nil || m.EmbedAudio == nil { + t.Fatalf("unified media loaders were suppressed: vision=%v embed_vision=%v embed_audio=%v", m.UnifiedVision != nil, m.EmbedVision != nil, m.EmbedAudio != nil) + } + if m.Vision != nil || m.Audio != nil { + t.Fatalf("unified model loaded tower encoders: vision=%v audio=%v", m.Vision != nil, m.Audio != nil) + } + + assertGemma4LinearStorage(t, m.UnifiedVision.PatchDense, quantized) + assertGemma4LinearStorage(t, m.EmbedVision.Projection, quantized) + assertGemma4LinearStorage(t, m.EmbedAudio.Projection, quantized) + if quantized { + assertGemma4QuantizedLinearType(t, m.UnifiedVision.PatchDense, "mxfp8") + assertGemma4QuantizedLinearType(t, m.EmbedVision.Projection, "mxfp8") + assertGemma4QuantizedLinearType(t, m.EmbedAudio.Projection, "mxfp8") + projected := m.EmbedAudio.Projection.Forward(mlx.Zeros(mlx.DTypeBFloat16, 1, 640)) + mlx.Eval(projected) + if got, hidden := projected.Dims(), int(m.HiddenSize); len(got) != 2 || got[0] != 1 || got[1] != hidden { + t.Fatalf("quantized unified audio projection shape = %v, want [1 %d]", got, hidden) + } + } + if m.UnifiedVision.PosEmbedding != tensors["model.vision_embedder.pos_embedding"] { + t.Fatal("unified position embedding was not retained as a dense tensor") + } + if quantized { + patch, ok := m.UnifiedVision.PatchDense.(*nn.QuantizedLinear) + if !ok || patch.Bias != tensors["model.vision_embedder.patch_dense.bias"] { + t.Fatal("quantized unified patch projection lost its learned bias") + } + } + }) + } +} + +func TestGemma4UnifiedImporterOutputKeepsDirectAudioDense(t *testing.T) { + useMLXTestThread(t) + m, tensors := testGemma4UnifiedMediaModel(true) + const path = "model.embed_audio.embedding_projection" + weightName := path + ".weight" + tensors[weightName] = testGemma4Array(64, 640) + delete(tensors, weightName+"_scale") + delete(tensors, weightName+"_qbias") + delete(m.TensorQuant, weightName) + + if err := m.LoadWeights(tensors); err != nil { + t.Fatal(err) + } + assertGemma4LinearStorage(t, m.EmbedAudio.Projection, false) + assertGemma4LinearStorage(t, m.UnifiedVision.PatchDense, true) + assertGemma4LinearStorage(t, m.EmbedVision.Projection, true) +} + +func TestGemma4MediaQuantizedLoadUsesSharedGlobalScale(t *testing.T) { + useMLXTestThread(t) + m, tensors := testGemma4UnifiedMediaModel(true) + globalScale := mlx.FromValue(float32(1)) + tensors["model.vision_embedder.patch_dense.weight.global_scale"] = globalScale + if err := m.LoadWeights(tensors); err != nil { + t.Fatal(err) + } + linear, ok := m.UnifiedVision.PatchDense.(*nn.QuantizedLinear) + if !ok { + t.Fatalf("patch projection type = %T, want *nn.QuantizedLinear", m.UnifiedVision.PatchDense) + } + if linear.GlobalScale != globalScale { + t.Fatal("Gemma4 media projection did not reuse the shared global-scale tensor") + } +} + +func testGemma4TowerMediaModel(t *testing.T, quantized bool) (*Model, map[string]*mlx.Array) { + t.Helper() + const hidden = 64 + tensorQuant := make(map[string]*model.TensorQuantInfo) + tensors := testGemma4BaseTensors(hidden) + visionConfig := &VisionConfig{ + ModelType: "gemma4_vision", HiddenSize: hidden, IntermediateSize: 128, + NumHiddenLayers: 1, NumAttentionHeads: 1, NumKeyValueHeads: 1, HeadDim: hidden, + RMSNormEps: 1e-6, DefaultOutputLength: 1, PatchSize: 8, PositionEmbeddingSize: 4, + } + audioConfig := &AudioConfig{ + ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, + AttentionInvalidLogit: -1e9, AttentionLogitCap: 50, ConvKernelSize: 3, + GradientClipping: 1e10, HiddenSize: hidden, NumAttentionHeads: 1, + NumHiddenLayers: 1, OutputProjDims: hidden, ResidualWeight: 0.5, + RMSNormEps: 1e-6, SubsamplingConvChannels: []int32{2, 2}, + } + + putGemma4TestLinear(tensors, tensorQuant, "model.vision_tower.patch_embedder.input_proj", hidden, 8*8*3, quantized) + tensors["model.vision_tower.patch_embedder.position_embedding_table"] = testGemma4Array(2, 4, hidden) + for i := range visionConfig.NumHiddenLayers { + layer := fmt.Sprintf("model.vision_tower.encoder.layers.%d", i) + for _, spec := range []struct { + path string + out, in int + }{ + {layer + ".self_attn.q_proj.linear", hidden, hidden}, + {layer + ".self_attn.k_proj.linear", hidden, hidden}, + {layer + ".self_attn.v_proj.linear", hidden, hidden}, + {layer + ".self_attn.o_proj.linear", hidden, hidden}, + {layer + ".mlp.gate_proj.linear", 128, hidden}, + {layer + ".mlp.up_proj.linear", 128, hidden}, + {layer + ".mlp.down_proj.linear", hidden, 128}, + } { + putGemma4TestLinear(tensors, tensorQuant, spec.path, spec.out, spec.in, quantized) + } + for _, suffix := range []string{ + ".self_attn.q_norm.weight", ".self_attn.k_norm.weight", + ".input_layernorm.weight", ".post_attention_layernorm.weight", + ".pre_feedforward_layernorm.weight", ".post_feedforward_layernorm.weight", + } { + tensors[layer+suffix] = testGemma4Array(hidden) + } + } + putGemma4TestLinear(tensors, tensorQuant, "model.embed_vision.embedding_projection", hidden, hidden, quantized) + + required, err := gemma4metadata.RequiredAudioTensorShapes(audioMetadataConfig(audioConfig, hidden)) + if err != nil { + t.Fatal(err) + } + for name, shape := range required { + if len(shape) == 2 && strings.HasSuffix(name, ".weight") { + putGemma4TestLinear(tensors, tensorQuant, strings.TrimSuffix(name, ".weight"), int(shape[0]), int(shape[1]), quantized) + continue + } + tensors[name] = testGemma4Array32(shape...) + } + + return &Model{ + TextConfig: &TextConfig{ + HiddenSize: hidden, VocabSize: 4, RMSNormEps: 1e-6, + TensorQuant: tensorQuant, + }, + VisionConfig: visionConfig, + AudioConfig: audioConfig, + AudioProcessorConfig: &AudioProcessorConfig{}, + }, tensors +} + +func testGemma4UnifiedMediaModel(quantized bool) (*Model, map[string]*mlx.Array) { + const hidden = 64 + tensorQuant := make(map[string]*model.TensorQuantInfo) + tensors := testGemma4BaseTensors(hidden) + visionConfig := &VisionConfig{ + ModelType: "gemma4_unified_vision", Unified: true, MMEmbedDim: hidden, MMPosembSize: 4, + ModelPatchSize: 8, PatchSize: 8, PoolingKernelSize: 1, RMSNormEps: 1e-6, DefaultOutputLength: 4, + } + audioConfig := &AudioConfig{ + ModelType: "gemma4_unified_audio", AudioEmbedDim: 640, AudioSamplesPerToken: 640, + HiddenSize: 640, OutputProjDims: 640, RMSNormEps: 1e-6, + } + patchDim := 8 * 8 * 3 + for _, name := range []string{"patch_ln1", "patch_ln2", "pos_norm"} { + dim := hidden + if name == "patch_ln1" { + dim = patchDim + } + tensors["model.vision_embedder."+name+".weight"] = testGemma4Array(dim) + tensors["model.vision_embedder."+name+".bias"] = testGemma4Array(dim) + } + tensors["model.vision_embedder.pos_embedding"] = testGemma4Array(4, 2, hidden) + putGemma4TestLinear(tensors, tensorQuant, "model.vision_embedder.patch_dense", hidden, patchDim, quantized) + tensors["model.vision_embedder.patch_dense.bias"] = testGemma4Array(hidden) + putGemma4TestLinear(tensors, tensorQuant, "model.embed_vision.embedding_projection", hidden, hidden, quantized) + putGemma4TestLinear(tensors, tensorQuant, "model.embed_audio.embedding_projection", hidden, 640, quantized) + + return &Model{ + TextConfig: &TextConfig{ + HiddenSize: hidden, VocabSize: 4, RMSNormEps: 1e-6, + TensorQuant: tensorQuant, + }, + VisionConfig: visionConfig, + AudioConfig: audioConfig, + AudioProcessorConfig: &AudioProcessorConfig{}, + }, tensors +} + +func testGemma4BaseTensors(hidden int) map[string]*mlx.Array { + return map[string]*mlx.Array{ + "model.embed_tokens.weight": testGemma4Array(4, hidden), + "model.norm.weight": testGemma4Array(hidden), + } +} + +func putGemma4TestLinear(tensors map[string]*mlx.Array, tensorQuant map[string]*model.TensorQuantInfo, path string, out, in int, quantized bool) { + weightName := path + ".weight" + weight := testGemma4Array(out, in) + if !quantized { + tensors[weightName] = weight + return + } + quantType := "nvfp4" + if strings.Contains(path, ".v_proj") || strings.Contains(path, ".k_proj") || strings.Contains(path, "down_proj") || isGemma4TestMediaBoundaryProjection(path) { + quantType = "mxfp8" + } + groupSize, bits, mode := model.QuantizationParams(quantType) + packed, scales, qbias := mlx.Quantize(weight, groupSize, bits, mode) + if qbias != nil { + mlx.Eval(packed, scales, qbias) + tensors[weightName+"_qbias"] = qbias + } else { + mlx.Eval(packed, scales) + } + tensors[weightName] = packed + tensors[weightName+"_scale"] = scales + tensorQuant[weightName] = &model.TensorQuantInfo{QuantType: quantType, GroupSize: groupSize} +} + +func isGemma4TestMediaBoundaryProjection(path string) bool { + for _, suffix := range []string{ + "embed_vision.embedding_projection", + "vision_embedder.patch_dense", + "audio_tower.output_proj", + "embed_audio.embedding_projection", + } { + if strings.HasSuffix(path, suffix) { + return true + } + } + return false +} + +func testGemma4Array(shape ...int) *mlx.Array { + if len(shape) == 0 { + return mlx.FromValue(float32(0.01)) + } + return mlx.AddScalar(mlx.Zeros(mlx.DTypeBFloat16, shape...), 0.01) +} + +func testGemma4Array32(shape ...int32) *mlx.Array { + dims := make([]int, len(shape)) + for i, dim := range shape { + dims[i] = int(dim) + } + return testGemma4Array(dims...) +} + +func assertGemma4LinearStorage(t *testing.T, linear nn.LinearLayer, quantized bool) { + t.Helper() + if quantized { + if _, ok := linear.(*nn.QuantizedLinear); !ok { + t.Fatalf("linear type = %T, want *nn.QuantizedLinear", linear) + } + return + } + if _, ok := linear.(*nn.Linear); !ok { + t.Fatalf("linear type = %T, want *nn.Linear", linear) + } +} + +func assertGemma4QuantizedLinearType(t *testing.T, linear nn.LinearLayer, quantType string) { + t.Helper() + quantized, ok := linear.(*nn.QuantizedLinear) + if !ok { + t.Fatalf("linear type = %T, want *nn.QuantizedLinear", linear) + } + groupSize, bits, mode := model.QuantizationParams(quantType) + if quantized.GroupSize != groupSize || quantized.Bits != bits || quantized.Mode != mode { + t.Fatalf("quantization params = (%d, %d, %q), want (%d, %d, %q)", + quantized.GroupSize, quantized.Bits, quantized.Mode, groupSize, bits, mode) + } +} diff --git a/x/models/gemma4/media_quant_test.go b/x/models/gemma4/media_quant_test.go new file mode 100644 index 00000000000..30d7d35842b --- /dev/null +++ b/x/models/gemma4/media_quant_test.go @@ -0,0 +1,113 @@ +package gemma4 + +import ( + "math" + "strings" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/models/nn" +) + +func TestValidateGemma4PackedLinearShape(t *testing.T) { + tests := []struct { + name string + weight []int + scales []int + want []int + groupSize int + bits int + mode string + wantErr bool + }{ + {name: "nvfp4", weight: []int{3840, 80}, scales: []int{3840, 40}, want: []int{3840, 640}, groupSize: 16, bits: 4, mode: "nvfp4"}, + {name: "mxfp8", weight: []int{768, 192}, scales: []int{768, 24}, want: []int{768, 768}, groupSize: 32, bits: 8, mode: "mxfp8"}, + {name: "wrong output", weight: []int{767, 96}, scales: []int{768, 48}, want: []int{768, 768}, groupSize: 16, bits: 4, mode: "nvfp4", wantErr: true}, + {name: "wrong packed input", weight: []int{768, 95}, scales: []int{768, 48}, want: []int{768, 768}, groupSize: 16, bits: 4, mode: "nvfp4", wantErr: true}, + {name: "wrong scales", weight: []int{768, 96}, scales: []int{768, 47}, want: []int{768, 768}, groupSize: 16, bits: 4, mode: "nvfp4", wantErr: true}, + {name: "unsupported bits", weight: []int{768, 96}, scales: []int{768, 48}, want: []int{768, 768}, groupSize: 16, bits: 2, mode: "nvfp4", wantErr: true}, + {name: "missing mode", weight: []int{768, 96}, scales: []int{768, 48}, want: []int{768, 768}, groupSize: 16, bits: 4, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateGemma4PackedLinearShape(tt.weight, tt.scales, tt.want, tt.groupSize, tt.bits, tt.mode) + if (err != nil) != tt.wantErr { + t.Fatalf("validateGemma4PackedLinearShape() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestGemma4QuantizedMediaLinearForward(t *testing.T) { + for _, quantType := range []string{"int4", "nvfp4", "mxfp8"} { + t.Run(quantType, func(t *testing.T) { + useMLXTestThread(t) + groupSize, bits, mode := model.QuantizationParams(quantType) + weight := mlx.AddScalar(mlx.Zeros(mlx.DTypeBFloat16, 64, 64), 0.125) + packed, scales, qbias := mlx.Quantize(weight, groupSize, bits, mode) + tensors := map[string]*mlx.Array{ + "media.projection.weight": packed, + "media.projection.weight_scale": scales, + "media.projection.bias": mlx.AddScalar(mlx.Zeros(mlx.DTypeBFloat16, 64), 0.25), + } + if qbias != nil { + tensors["media.projection.weight_qbias"] = qbias + } + tensorQuant := map[string]*model.TensorQuantInfo{ + "media.projection.weight": {QuantType: quantType, GroupSize: groupSize}, + } + if err := validateGemma4LinearWeight(tensors, "media.projection", []int{64, 64}, 0, 0, "", tensorQuant); err != nil { + t.Fatal(err) + } + validScales := scales + if mode == "affine" { + tensors["media.projection.weight_scale"] = mlx.Zeros(mlx.DTypeUint8, scales.Dims()...) + err := validateGemma4LinearWeight(tensors, "media.projection", []int{64, 64}, 0, 0, "", tensorQuant) + if err == nil || !strings.Contains(err.Error(), "media.projection") { + t.Fatalf("affine uint8 scales error = %v", err) + } + tensors["media.projection.weight_scale"] = validScales + delete(tensors, "media.projection.weight_qbias") + err = validateGemma4LinearWeight(tensors, "media.projection", []int{64, 64}, 0, 0, "", tensorQuant) + if err == nil || !strings.Contains(err.Error(), "weight_qbias") { + t.Fatalf("missing affine qbias error = %v", err) + } + tensors["media.projection.weight_qbias"] = qbias + } else { + tensors["media.projection.weight_scale"] = mlx.Zeros(mlx.DTypeBFloat16, scales.Dims()...) + err := validateGemma4LinearWeight(tensors, "media.projection", []int{64, 64}, 0, 0, "", tensorQuant) + if err == nil || !strings.Contains(err.Error(), "media.projection") { + t.Fatalf("%s floating scales error = %v", mode, err) + } + tensors["media.projection.weight_scale"] = validScales + tensors["media.projection.weight_qbias"] = mlx.Zeros(mlx.DTypeBFloat16, scales.Dims()...) + err = validateGemma4LinearWeight(tensors, "media.projection", []int{64, 64}, 0, 0, "", tensorQuant) + if err == nil || !strings.Contains(err.Error(), "unexpected") { + t.Fatalf("%s qbias error = %v", mode, err) + } + delete(tensors, "media.projection.weight_qbias") + } + linear, ok := model.NewLinearFactory(tensors, 0, 0, "", tensorQuant).Make("media.projection").(*nn.QuantizedLinear) + if !ok { + t.Fatal("quantized media linear was not constructed") + } + if linear.Bias != tensors["media.projection.bias"] || linear.QBiases != qbias { + t.Fatal("learned and quantization biases were not kept separate") + } + output := linear.Forward(mlx.AddScalar(mlx.Zeros(mlx.DTypeBFloat16, 1, 64), 1)) + mlx.Eval(output) + if got := output.Dims(); len(got) != 2 || got[0] != 1 || got[1] != 64 { + t.Fatalf("quantized media output shape = %v, want [1 64]", got) + } + floatOutput := output.AsType(mlx.DTypeFloat32) + mlx.Eval(floatOutput) + for i, value := range floatOutput.Floats() { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + t.Fatalf("quantized media output[%d] is non-finite: %g", i, value) + } + } + }) + } +} diff --git a/x/models/gemma4/metadata/audio.go b/x/models/gemma4/metadata/audio.go index c3bcf698b21..d7662534610 100644 --- a/x/models/gemma4/metadata/audio.go +++ b/x/models/gemma4/metadata/audio.go @@ -5,6 +5,8 @@ import ( "fmt" "math" "slices" + "sort" + "strings" "github.com/ollama/ollama/x/tokenizer" ) @@ -163,7 +165,7 @@ func ValidateAudioTensors(cfg ConfigFile, names []string) error { for _, name := range names { present[name] = struct{}{} } - for name := range shapes { + for _, name := range sortedAudioShapeNames(shapes) { if _, ok := present[name]; !ok { return fmt.Errorf("missing %s", name) } @@ -181,7 +183,8 @@ func ValidateAudioSourceInventory(cfg ConfigFile, tensors map[string]TensorDescr if err != nil { return err } - for name, shape := range shapes { + for _, name := range sortedAudioShapeNames(shapes) { + shape := shapes[name] desc, ok := tensors[name] if !ok { return fmt.Errorf("missing %s", name) @@ -196,11 +199,56 @@ func ValidateAudioSourceInventory(cfg ConfigFile, tensors map[string]TensorDescr return nil } -// ValidateAudioInstalledInventory validates the normalized descriptors stored -// in installed tensor layers. Installed audio remains source precision at this -// row, so its descriptor contract is identical to the released source form. +// ValidateAudioInstalledInventory validates normalized descriptors stored in +// installed tensor layers, including dense and quantized linear weights. func ValidateAudioInstalledInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { - return ValidateAudioSourceInventory(cfg, tensors) + return validateAudioNormalizedInventory(cfg, tensors, installedMode) +} + +// ValidateAudioRuntimeInventory validates the materialized runtime descriptor +// inventory. It is the authoritative completeness/readiness gate used before +// the MLX audio loader consumes arrays. +func ValidateAudioRuntimeInventory(cfg ConfigFile, tensors map[string]TensorDescriptor) error { + return validateAudioNormalizedInventory(cfg, tensors, runtimeMode) +} + +func validateAudioNormalizedInventory(cfg ConfigFile, tensors map[string]TensorDescriptor, mode inventoryMode) error { + if err := validateAudioConfig(cfg); err != nil { + return err + } + shapes, err := requiredAudioShapes(cfg) + if err != nil { + return err + } + for _, name := range sortedAudioShapeNames(shapes) { + shape := shapes[name] + desc, ok := tensors[name] + if !ok { + return fmt.Errorf("missing %s", name) + } + if len(shape) == 2 && strings.HasSuffix(name, ".weight") && (!isFloat(desc.Dtype) || hasProducerCompanion(tensors, strings.TrimSuffix(name, ".weight"))) { + if err := requireLinearDescriptor(tensors, strings.TrimSuffix(name, ".weight"), shape, mode, cfg); err != nil { + return err + } + continue + } + if !slices.Equal(desc.Shape, shape) { + return fmt.Errorf("%s shape %v, want %v", name, desc.Shape, shape) + } + if !isFloat(desc.Dtype) { + return fmt.Errorf("%s dtype %s is not floating point", name, desc.Dtype) + } + } + return nil +} + +func sortedAudioShapeNames(shapes map[string][]int32) []string { + names := make([]string, 0, len(shapes)) + for name := range shapes { + names = append(names, name) + } + sort.Strings(names) + return names } // RequiredAudioTensorShapes returns a copy of the normalized released audio diff --git a/x/models/gemma4/metadata/audio_test.go b/x/models/gemma4/metadata/audio_test.go index c332069562b..7923b634c9f 100644 --- a/x/models/gemma4/metadata/audio_test.go +++ b/x/models/gemma4/metadata/audio_test.go @@ -154,6 +154,71 @@ func TestValidateReleasedAudioInventory(t *testing.T) { } } +func TestValidateAudioNormalizedQuantizedInventory(t *testing.T) { + cfg := releasedAudioConfig(1) + dense := completeAudioInventory(cfg) + if err := ValidateAudioInstalledInventory(cfg, dense); err != nil { + t.Fatalf("dense installed inventory: %v", err) + } + if err := ValidateAudioRuntimeInventory(cfg, dense); err != nil { + t.Fatalf("dense runtime inventory: %v", err) + } + + const base = "model.audio_tower.layers.0.self_attn.q_proj.linear" + logical := dense[base+".weight"].Shape + for _, tt := range []struct { + name string + bits int32 + groupSize int32 + }{ + {"mxfp4", 4, 32}, + {"mxfp8", 8, 32}, + } { + t.Run(tt.name, func(t *testing.T) { + installed := maps.Clone(dense) + installed[base+".weight"] = TensorDescriptor{ + Dtype: "U32", Shape: []int32{logical[0], logical[1] / (32 / tt.bits)}, + QuantType: tt.name, GroupSize: int(tt.groupSize), + } + installed[base+".weight.scale"] = TensorDescriptor{Dtype: "U8", Shape: []int32{logical[0], logical[1] / tt.groupSize}} + if err := ValidateAudioInstalledInventory(cfg, installed); err != nil { + t.Fatalf("installed inventory: %v", err) + } + + runtime := maps.Clone(installed) + delete(runtime, base+".weight.scale") + runtime[base+".weight_scale"] = TensorDescriptor{Dtype: "U8", Shape: []int32{logical[0], logical[1] / tt.groupSize}} + if err := ValidateAudioRuntimeInventory(cfg, runtime); err != nil { + t.Fatalf("runtime inventory: %v", err) + } + + delete(runtime, base+".weight_scale") + if err := ValidateAudioRuntimeInventory(cfg, runtime); err == nil { + t.Fatal("incomplete quantized runtime inventory accepted") + } + }) + } + + missing := maps.Clone(dense) + names := sortedAudioShapeNames(mustRequiredAudioShapes(t, cfg)) + delete(missing, names[0]) + delete(missing, names[1]) + for range 3 { + if err := ValidateAudioRuntimeInventory(cfg, missing); err == nil || !strings.Contains(err.Error(), names[0]) { + t.Fatalf("deterministic first missing tensor error = %v, want %s", err, names[0]) + } + } +} + +func mustRequiredAudioShapes(t *testing.T, cfg ConfigFile) map[string][]int32 { + t.Helper() + shapes, err := requiredAudioShapes(cfg) + if err != nil { + t.Fatal(err) + } + return shapes +} + func TestValidateAudioTensorsExactNormalizedNames(t *testing.T) { cfg := releasedAudioConfig(1) shapes, err := RequiredAudioTensorShapes(cfg) @@ -210,6 +275,29 @@ func TestValidateReleasedUnifiedAudioInventory(t *testing.T) { if err := ValidateAudioSourceInventory(cfg, bad); err == nil { t.Fatal("bad unified projection shape: error = nil") } + + quantized := maps.Clone(tensors) + quantized["model.embed_audio.embedding_projection.weight"] = TensorDescriptor{ + Dtype: "U32", Shape: []int32{3840, 160}, QuantType: "mxfp8", GroupSize: 32, + } + quantized["model.embed_audio.embedding_projection.weight.scale"] = TensorDescriptor{ + Dtype: "U8", Shape: []int32{3840, 20}, + } + if err := ValidateAudioInstalledInventory(cfg, quantized); err != nil { + t.Fatalf("quantized unified installed inventory: %v", err) + } + runtime := maps.Clone(quantized) + delete(runtime, "model.embed_audio.embedding_projection.weight.scale") + runtime["model.embed_audio.embedding_projection.weight_scale"] = TensorDescriptor{ + Dtype: "U8", Shape: []int32{3840, 20}, + } + if err := ValidateAudioRuntimeInventory(cfg, runtime); err != nil { + t.Fatalf("quantized unified runtime inventory: %v", err) + } + delete(runtime, "model.embed_audio.embedding_projection.weight_scale") + if err := ValidateAudioRuntimeInventory(cfg, runtime); err == nil { + t.Fatal("incomplete quantized unified runtime inventory accepted") + } } func TestValidateAudioConfig(t *testing.T) { diff --git a/x/models/gemma4/metadata/vision.go b/x/models/gemma4/metadata/vision.go index 204e937570f..978391f70ea 100644 --- a/x/models/gemma4/metadata/vision.go +++ b/x/models/gemma4/metadata/vision.go @@ -347,6 +347,14 @@ func ValidateVisionRuntimeInventory(cfg ConfigFile, tensors map[string]TensorDes return validateInventory(cfg, tensors, runtimeMode) } +// ValidateRuntimeLinearDescriptor validates one materialized linear against +// the same normalized runtime contract used by the complete media inventory +// validators. Loaders use this only for checks that require concrete arrays; +// readiness and completeness remain owned by the full inventory validators. +func ValidateRuntimeLinearDescriptor(cfg ConfigFile, tensors map[string]TensorDescriptor, path string, logical []int32) error { + return requireLinearDescriptor(tensors, path, logical, runtimeMode, cfg) +} + type inventoryMode int const ( @@ -678,24 +686,62 @@ func validateNormalizedQuant(t map[string]TensorDescriptor, base string, weight return fmt.Errorf("missing or invalid normalized scale %s", scaleName) } if strings.EqualFold(scale.Dtype, "U8") { - if !strings.EqualFold(weight.Dtype, "U32") || !packedWeightShape(weight.Shape, logical, 8) || !packedScaleShape(scale, logical, 16) { - return fmt.Errorf("invalid normalized NVFP4 descriptors for %s", base) + contractWeight := weight + if contractWeight.QuantType == "" { + quantization := sourceQuant(cfg) + contractWeight.QuantType = quantization.Mode + if contractWeight.GroupSize == 0 { + contractWeight.GroupSize = quantization.GroupSize + } + } + bits, groupSize, quantType, ok := normalizedFloatQuantContract(contractWeight) + if weight.QuantType == "" { + quantization := sourceQuant(cfg) + if quantization.Bits != 0 && quantization.Bits != bits { + ok = false + } + if quantization.GroupSize != 0 && quantization.GroupSize != groupSize { + ok = false + } + if quantization.Bits != 0 && quantization.GroupSize != 0 && quantization.Mode == "" { + ok = false + } + if quantization.Mode != "" && quantization.Mode != quantType { + ok = false + } + } + if !ok || !strings.EqualFold(weight.Dtype, "U32") || + !packedWeightShape(weight.Shape, logical, int32(32/bits)) || + !packedScaleShape(scale, logical, int32(groupSize)) { + return fmt.Errorf("invalid normalized floating-point quantization descriptors for %s", base) } if global, present := t[base+".weight.global_scale"]; present && (!isScalar(global.Shape) || !strings.EqualFold(global.Dtype, "F32")) { return fmt.Errorf("invalid normalized global scale for %s", base) } if _, present := t[biasName]; present { - return fmt.Errorf("unexpected affine bias for normalized NVFP4 %s", base) + return fmt.Errorf("unexpected affine bias for normalized %s %s", quantType, base) } return nil } bits, groupSize, ok := inferAffineContract(weight, scale, logical) - q := sourceQuant(cfg) - if q.Bits > 0 && q.Bits != bits { - ok = false - } - if q.GroupSize > 0 && q.GroupSize != groupSize { - ok = false + if weight.QuantType != "" { + switch strings.ToLower(weight.QuantType) { + case "int4": + ok = ok && bits == 4 + case "int8": + ok = ok && bits == 8 + case "affine": + default: + ok = false + } + } else { + q := sourceQuant(cfg) + if q.Bits > 0 && q.Bits != bits { + ok = false + } + if q.GroupSize > 0 && q.GroupSize != groupSize { + ok = false + } } if weight.GroupSize > 0 && weight.GroupSize != groupSize { ok = false @@ -703,8 +749,8 @@ func validateNormalizedQuant(t map[string]TensorDescriptor, base string, weight if !ok || !isFloat(scale.Dtype) || !strings.EqualFold(weight.Dtype, "U32") { return fmt.Errorf("invalid normalized affine descriptors for %s", base) } - if bias, present := t[biasName]; present && (!isFloat(bias.Dtype) || !slices.Equal(bias.Shape, scale.Shape)) { - return fmt.Errorf("invalid normalized affine bias for %s", base) + if bias, present := t[biasName]; !present || !isFloat(bias.Dtype) || !slices.Equal(bias.Shape, scale.Shape) { + return fmt.Errorf("missing or invalid normalized affine bias %s", biasName) } if _, present := t[base+".weight.global_scale"]; present { return fmt.Errorf("unexpected global scale for normalized affine %s", base) @@ -712,6 +758,31 @@ func validateNormalizedQuant(t map[string]TensorDescriptor, base string, weight return nil } +func normalizedFloatQuantContract(weight TensorDescriptor) (bits, groupSize int, quantType string, ok bool) { + quantType = strings.ToLower(weight.QuantType) + switch quantType { + case "": + // Preserve the accepted normalized-NVFP4 representation, which + // predates explicit per-tensor quantization metadata. + return 4, 16, "nvfp4", weight.GroupSize == 0 || weight.GroupSize == 16 + case "nvfp4": + groupSize = 16 + bits = 4 + case "mxfp4": + groupSize = 32 + bits = 4 + case "mxfp8": + groupSize = 32 + bits = 8 + default: + return 0, 0, quantType, false + } + if weight.GroupSize != 0 && weight.GroupSize != groupSize { + return 0, 0, quantType, false + } + return bits, groupSize, quantType, true +} + func inferAffineContract(weight, scale TensorDescriptor, logical []int32) (bits, groupSize int, ok bool) { if len(weight.Shape) != 2 || len(scale.Shape) != 2 || len(logical) != 2 || weight.Shape[0] != logical[0] || scale.Shape[0] != logical[0] || weight.Shape[1] <= 0 || scale.Shape[1] <= 0 { return 0, 0, false diff --git a/x/models/gemma4/metadata/vision_test.go b/x/models/gemma4/metadata/vision_test.go index 7959303256e..b6bb1b66725 100644 --- a/x/models/gemma4/metadata/vision_test.go +++ b/x/models/gemma4/metadata/vision_test.go @@ -246,6 +246,31 @@ func TestVisionInventoryProducerAndNormalizationMatrix(t *testing.T) { if err := ValidateVisionRuntimeInventory(cfg, runtime); err != nil { t.Fatalf("normalized runtime error = %v", err) } + for _, tt := range []struct { + name string + bits int32 + groupSize int32 + }{ + {"mxfp4", 4, 32}, + {"mxfp8", 8, 32}, + } { + t.Run("normalized "+tt.name, func(t *testing.T) { + const floatTarget = "model.vision_tower.encoder.layers.0.mlp.down_proj.linear" + candidate := cloneDescriptors(dense) + candidate[floatTarget+".weight"] = TensorDescriptor{ + Dtype: "U32", Shape: []int32{16, 32 / (32 / tt.bits)}, + QuantType: tt.name, GroupSize: int(tt.groupSize), + } + candidate[floatTarget+".weight.scale"] = TensorDescriptor{Dtype: "U8", Shape: []int32{16, 32 / tt.groupSize}} + if err := ValidateVisionInstalledInventory(cfg, candidate); err != nil { + t.Fatalf("installed %s error = %v", tt.name, err) + } + delete(candidate, floatTarget+".weight.scale") + if err := ValidateVisionInstalledInventory(cfg, candidate); err == nil { + t.Fatalf("incomplete installed %s inventory accepted", tt.name) + } + }) + } for _, missing := range []string{target + ".weight.scale"} { partial := cloneDescriptors(normalized) delete(partial, missing) diff --git a/x/models/gemma4/vision.go b/x/models/gemma4/vision.go index ce90c4e3e75..0e052252e59 100644 --- a/x/models/gemma4/vision.go +++ b/x/models/gemma4/vision.go @@ -236,7 +236,7 @@ func validateVisionConfig(cfg *VisionConfig) error { } // This validates vision-only call sites before the enclosing text model is // available. Executable inventory validation supplies the real text width. - _, err := gemma4metadata.ProjectVisionArchitecture(metadataConfigFromVision(cfg, 1)) + _, err := gemma4metadata.ProjectVisionArchitecture(metadataConfigFromVision(cfg, 1, 0, 0, "")) return err } @@ -298,7 +298,7 @@ func parseGemma4MediaTokens(data []byte, fallback gemma4MediaTokens) gemma4Media return fallback } -func metadataConfigFromVision(cfg *VisionConfig, textHidden int) gemma4metadata.ConfigFile { +func metadataConfigFromVision(cfg *VisionConfig, textHidden, groupSize, bits int, mode string) gemma4metadata.ConfigFile { if cfg == nil { return gemma4metadata.ConfigFile{} } @@ -316,10 +316,20 @@ func metadataConfigFromVision(cfg *VisionConfig, textHidden int) gemma4metadata. v.ModelType = "gemma4_unified_vision" } v.RopeParameters.RopeTheta = float64(cfg.RopeParameters.RopeTheta) - return gemma4metadata.ConfigFile{TextConfig: gemma4metadata.TextConfig{HiddenSize: textHidden}, VisionConfig: v} + return gemma4metadata.ConfigFile{ + TextConfig: gemma4metadata.TextConfig{HiddenSize: textHidden}, + VisionConfig: v, + QuantizationConfig: gemma4metadata.Quantization{Bits: bits, GroupSize: groupSize, Mode: mode}, + } } -func validateGemma4VisionWeights(tensors map[string]*mlx.Array, cfg *VisionConfig, textHidden int, tq map[string]*model.TensorQuantInfo) (bool, error) { +func validateGemma4VisionWeights( + tensors map[string]*mlx.Array, + cfg *VisionConfig, + textHidden, groupSize, bits int, + mode string, + tq map[string]*model.TensorQuantInfo, +) (bool, error) { sentinel := firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight", "model.vision_tower.patch_embedder.input_proj.weight") packedSentinel := firstNonNil(tensors, "vision_tower.patch_embedder.input_proj.weight_packed", "model.vision_tower.patch_embedder.input_proj.weight_packed") if cfg != nil && cfg.unified() { @@ -332,22 +342,8 @@ func validateGemma4VisionWeights(tensors map[string]*mlx.Array, cfg *VisionConfi } return false, nil } - descriptors := make(map[string]gemma4metadata.TensorDescriptor, len(tensors)) - for name, tensor := range tensors { - if tensor != nil { - shape := make([]int32, tensor.NumDims()) - for i, d := range tensor.Dims() { - shape[i] = int32(d) - } - d := gemma4metadata.TensorDescriptor{Dtype: tensor.DType().String(), Shape: shape} - if q := tq[name]; q != nil { - d.QuantType = q.QuantType - d.GroupSize = q.GroupSize - } - descriptors[name] = d - } - } - if err := gemma4metadata.ValidateVisionRuntimeInventory(metadataConfigFromVision(cfg, textHidden), descriptors); err != nil { + descriptors := gemma4RuntimeTensorDescriptors(tensors, tq, groupSize, mode) + if err := gemma4metadata.ValidateVisionRuntimeInventory(metadataConfigFromVision(cfg, textHidden, groupSize, bits, mode), descriptors); err != nil { return false, err } return true, nil From 1a1e866d69ac6419e998d8968a45e0fc828cc741 Mon Sep 17 00:00:00 2001 From: Philipp Date: Sun, 16 Aug 2026 08:22:44 +0000 Subject: [PATCH 58/58] gemma4: hydrate unified vision manifest tensors Include unified vision_embedder tensor layers in the canonical installed descriptor inventory so capability reporting validates the complete model rather than suppressing a valid unified vision model. Co-authored-by: Codex --- server/images.go | 97 ++++++++++++++++++-- server/images_test.go | 129 +++++++++++++++++++++++---- server/model_inference_cache_test.go | 7 +- server/model_list_cache.go | 4 +- server/model_list_cache_test.go | 38 ++++++++ x/create/client/create_test.go | 28 ++++++ 6 files changed, 273 insertions(+), 30 deletions(-) diff --git a/server/images.go b/server/images.go index 7c842c33132..f704249bacb 100644 --- a/server/images.go +++ b/server/images.go @@ -509,9 +509,12 @@ func suppressVisionCapability(m *Model) bool { } func suppressGemma4SafetensorsVisionCapability(m *Model) bool { - if m == nil || !isLocalGemma4SafetensorsConfig(m.Config) { + if m == nil || !isGemma4SafetensorsConfig(m.Config) { return false } + if !isLocalGemma4SafetensorsConfig(m.Config) { + return true + } return m.Gemma4VisionConfig == nil || gemma4metadata.ValidateVisionInstalledInventory(*m.Gemma4VisionConfig, m.Gemma4VisionTensors) != nil } @@ -555,12 +558,15 @@ func isNemotron3NanoSafetensorsConfig(cfg model.ConfigV2) bool { } func isLocalGemma4SafetensorsConfig(cfg model.ConfigV2) bool { - return cfg.ModelFormat == "safetensors" && - isGemma4Renderer(cfg.Renderer) && + return isGemma4SafetensorsConfig(cfg) && cfg.RemoteHost == "" && cfg.RemoteModel == "" } +func isGemma4SafetensorsConfig(cfg model.ConfigV2) bool { + return cfg.ModelFormat == "safetensors" && isGemma4Renderer(cfg.Renderer) +} + func hasGemma4VisionTensorLayers(cfg gemma4metadata.ConfigFile, layers []manifest.Layer) bool { tensors, err := gemma4VisionTensorDescriptors(layers) return err == nil && gemma4metadata.ValidateVisionInstalledInventory(cfg, tensors) == nil @@ -618,7 +624,12 @@ func gemma4AudioTensorDescriptors(layers []manifest.Layer) (map[string]gemma4met return nil, fmt.Errorf("Gemma4 audio tensor inventory exceeds descriptor work limit %d", maxGemma4VisionDescriptorWork) } descriptorWork += work - tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + descriptor, err := gemma4TensorDescriptor(tensor, ext.metadata) + if err != nil { + ext.Close() + return nil, err + } + tensors[name] = descriptor } if err := ext.Close(); err != nil { return nil, err @@ -634,7 +645,9 @@ func gemma4VisionTensorDescriptors(layers []manifest.Layer) (map[string]gemma4me if layer.MediaType != manifest.MediaTypeImageTensor { continue } - if !strings.Contains(layer.Name, "vision_tower.") && !strings.Contains(layer.Name, "embed_vision.") { + if !strings.Contains(layer.Name, "vision_tower.") && + !strings.Contains(layer.Name, "vision_embedder.") && + !strings.Contains(layer.Name, "embed_vision.") { continue } filename, err := manifest.BlobsPath(layer.Digest) @@ -666,7 +679,12 @@ func gemma4VisionTensorDescriptors(layers []manifest.Layer) (map[string]gemma4me return nil, fmt.Errorf("Gemma4 vision tensor inventory exceeds descriptor work limit %d", maxGemma4VisionDescriptorWork) } descriptorWork += work - tensors[name] = gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + descriptor, err := gemma4TensorDescriptor(tensor, ext.metadata) + if err != nil { + ext.Close() + return nil, err + } + tensors[name] = descriptor } if err := ext.Close(); err != nil { return nil, err @@ -675,7 +693,12 @@ func gemma4VisionTensorDescriptors(layers []manifest.Layer) (map[string]gemma4me return tensors, nil } -func openGemma4TensorLayer(filename string) (extractor *safetensors.TensorExtractor, err error) { +type gemma4TensorLayer struct { + *safetensors.TensorExtractor + metadata map[string]string +} + +func openGemma4TensorLayer(filename string) (layer *gemma4TensorLayer, err error) { f, err := os.Open(filename) if err != nil { return nil, err @@ -708,16 +731,72 @@ func openGemma4TensorLayer(filename string) (extractor *safetensors.TensorExtrac f.Close() return nil, err } + metadata, err := gemma4SafetensorsMetadata(header) + if err != nil { + f.Close() + return nil, err + } if err := f.Close(); err != nil { return nil, err } defer func() { if recovered := recover(); recovered != nil { - extractor = nil + layer = nil err = fmt.Errorf("invalid safetensors header: %v", recovered) } }() - return safetensors.OpenForExtraction(filename) + extractor, err := safetensors.OpenForExtraction(filename) + if err != nil { + return nil, err + } + return &gemma4TensorLayer{TensorExtractor: extractor, metadata: metadata}, nil +} + +func gemma4SafetensorsMetadata(header []byte) (map[string]string, error) { + var entries map[string]json.RawMessage + if err := json.Unmarshal(header, &entries); err != nil { + return nil, fmt.Errorf("parse safetensors header: %w", err) + } + raw, ok := entries["__metadata__"] + if !ok { + return nil, nil + } + var metadata map[string]string + if err := json.Unmarshal(raw, &metadata); err != nil { + return nil, fmt.Errorf("parse safetensors metadata: %w", err) + } + return metadata, nil +} + +func gemma4TensorDescriptor(tensor *safetensors.TensorData, metadata map[string]string) (gemma4metadata.TensorDescriptor, error) { + descriptor := gemma4metadata.TensorDescriptor{Dtype: tensor.Dtype, Shape: slices.Clone(tensor.Shape)} + if !strings.EqualFold(tensor.Dtype, "U32") { + return descriptor, nil + } + + quantType, quantSpecific := metadata[tensor.Name+".quant_type"] + groupValue, groupSpecific := metadata[tensor.Name+".group_size"] + if quantSpecific != groupSpecific { + return gemma4metadata.TensorDescriptor{}, fmt.Errorf("incomplete quantization metadata for %s", tensor.Name) + } + if !quantSpecific { + quantType, quantSpecific = metadata["quant_type"] + groupValue, groupSpecific = metadata["group_size"] + if quantSpecific != groupSpecific { + return gemma4metadata.TensorDescriptor{}, fmt.Errorf("incomplete quantization metadata for %s", tensor.Name) + } + } + if !quantSpecific { + return descriptor, nil + } + + groupSize, err := strconv.ParseInt(groupValue, 10, 32) + if err != nil || groupSize <= 0 { + return gemma4metadata.TensorDescriptor{}, fmt.Errorf("invalid quantization group_size %q for %s", groupValue, tensor.Name) + } + descriptor.QuantType = strings.ToLower(quantType) + descriptor.GroupSize = int(groupSize) + return descriptor, nil } type gemma4SafetensorInfo struct { diff --git a/server/images_test.go b/server/images_test.go index 5a2cde8cd0e..1ce180975b3 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -15,6 +15,7 @@ import ( "os" "path/filepath" "slices" + "strconv" "strings" "testing" "time" @@ -614,6 +615,24 @@ func TestModelCapabilities(t *testing.T) { }, expectedCaps: []model.Capability{model.CapabilityAudio}, }, + { + name: "remote gemma4 safetensors suppresses local media", + model: Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + RemoteHost: "https://example.invalid", + Capabilities: []string{"vision", "audio"}, + }, + Gemma4VisionConfig: gemma4VisionConfig(2), + Gemma4VisionTensors: testGemma4VisionTensorDescriptors(2), + Gemma4AudioConfig: gemma4AudioConfig(1), + Gemma4AudioTensors: testGemma4AudioTensorDescriptors(1), + Gemma4AudioReady: true, + Template: chatTemplate, + }, + expectedCaps: nil, + }, } // compare two slices of model.Capability regardless of order @@ -710,6 +729,34 @@ func TestGemma4SafetensorsVisionCapabilityRequiresTensorLayers(t *testing.T) { } } +func TestGemma4SafetensorsVisionCapabilityReadsInstalledQuantMetadata(t *testing.T) { + setTestHome(t, t.TempDir()) + const projection = "model.embed_vision.embedding_projection" + descriptors := testGemma4VisionTensorDescriptorsForGeometry(1, 32, 64, 32, 2, 32, 32) + descriptors[projection+".weight"] = gemma4metadata.TensorDescriptor{ + Dtype: "U32", Shape: []int32{32, 8}, QuantType: "mxfp8", GroupSize: 32, + } + descriptors[projection+".weight.scale"] = gemma4metadata.TensorDescriptor{Dtype: "U8", Shape: []int32{32, 1}} + config := []byte(`{"text_config":{"hidden_size":32},"vision_config":{"hidden_size":32,"intermediate_size":64,"num_hidden_layers":1,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":32,"default_output_length":1,"patch_size":2,"position_embedding_size":32,"pooling_kernel_size":1}}`) + layers := gemma4VisionManifestLayersFromDescriptors(t, descriptors, config) + + createSafetensorsTestModel(t, "gemma4-quantized-vision", model.ConfigV2{ + ModelFormat: "safetensors", Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "vision"}, + }, layers) + m, err := GetModel("gemma4-quantized-vision") + if err != nil { + t.Fatal(err) + } + if !slices.Contains(m.Capabilities(), model.CapabilityVision) { + t.Fatalf("quantized capabilities = %v, want vision", m.Capabilities()) + } + got := m.Gemma4VisionTensors[projection+".weight"] + if got.QuantType != "mxfp8" || got.GroupSize != 32 { + t.Fatalf("installed quantization descriptor = %+v, want mxfp8/group32", got) + } +} + func TestGemma4SafetensorsAudioCapabilityRequiresCompleteInventory(t *testing.T) { complete := Model{ Config: model.ConfigV2{ @@ -755,15 +802,15 @@ func TestGemma4SafetensorsAudioCapabilityRequiresCompleteInventory(t *testing.T) remoteCfg := complete.Config remoteCfg.RemoteHost = "https://example.invalid" - remoteCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityAudio}, remoteCfg) - if slices.Contains(remoteCaps, model.CapabilityAudio) { - t.Fatal("remote Gemma 4 list capability exposed unsupported local audio") + remoteCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityAudio}, remoteCfg) + if slices.Contains(remoteCaps, model.CapabilityVision) || slices.Contains(remoteCaps, model.CapabilityAudio) { + t.Fatal("remote Gemma 4 list capability exposed unsupported local media") } otherCfg := remoteCfg otherCfg.Renderer = "other" - otherCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityAudio}, otherCfg) - if !slices.Contains(otherCaps, model.CapabilityAudio) { - t.Fatal("remote non-Gemma audio capability was suppressed") + otherCaps := filterUnsupportedModelListCapabilities([]model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityAudio}, otherCfg) + if !slices.Contains(otherCaps, model.CapabilityVision) || !slices.Contains(otherCaps, model.CapabilityAudio) { + t.Fatal("remote non-Gemma media capability was suppressed") } } @@ -1038,9 +1085,9 @@ func gemma4AudioFixtureLayer(t *testing.T, manifestName, internalName string, de if err != nil { t.Fatal(err) } - built, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{ + built, err := io.ReadAll(safetensors.BuildPackedSafetensorsReaderWithMetadata([]*safetensors.TensorData{ safetensors.NewTensorDataFromBytes(internalName, descriptor.Dtype, descriptor.Shape, make([]byte, int(size))), - })) + }, gemma4DescriptorMetadata(descriptor))) if err != nil { t.Fatal(err) } @@ -1178,9 +1225,44 @@ func TestGemma4UnifiedAudioCapabilityRequiresProjection(t *testing.T) { func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { t.Helper() + config := []byte(`{"text_config":{"hidden_size":6},"vision_config":{"hidden_size":4,"intermediate_size":8,"num_hidden_layers":2,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":4,"default_output_length":1,"patch_size":2,"position_embedding_size":16,"pooling_kernel_size":1}}`) + layers := gemma4VisionManifestLayersFromDescriptors(t, testGemma4VisionTensorDescriptors(2), config) + descriptors, err := gemma4VisionTensorDescriptors(layers) + if err != nil { + t.Fatal(err) + } + if err := gemma4metadata.ValidateVisionInstalledInventory(*gemma4VisionConfig(2), descriptors); err != nil { + t.Fatalf("validate Gemma4 descriptor fixture: %v", err) + } + return layers +} - layers := make([]manifest.Layer, 0, len(gemma4VisionTensorNames(2))+1) - for name, descriptor := range testGemma4VisionTensorDescriptors(2) { +func gemma4UnifiedVisionManifestLayers(t *testing.T, complete bool) []manifest.Layer { + t.Helper() + descriptors := map[string]gemma4metadata.TensorDescriptor{ + "model.vision_embedder.patch_ln1.weight": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_ln1.bias": {Dtype: "F32", Shape: []int32{12}}, + "model.vision_embedder.patch_dense.weight": {Dtype: "F32", Shape: []int32{3, 12}}, + "model.vision_embedder.patch_dense.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.patch_ln2.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_embedding": {Dtype: "F32", Shape: []int32{4, 2, 3}}, + "model.vision_embedder.pos_norm.weight": {Dtype: "F32", Shape: []int32{3}}, + "model.vision_embedder.pos_norm.bias": {Dtype: "F32", Shape: []int32{3}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "F32", Shape: []int32{5, 3}}, + } + if !complete { + delete(descriptors, "model.vision_embedder.pos_norm.bias") + } + config := []byte(`{"architectures":["Gemma4UnifiedForConditionalGeneration"],"model_type":"gemma4_unified","text_config":{"hidden_size":5},"vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3,"mm_posemb_size":4,"model_patch_size":2,"num_soft_tokens":2,"patch_size":1,"pooling_kernel_size":2}}`) + return gemma4VisionManifestLayersFromDescriptors(t, descriptors, config) +} + +func gemma4VisionManifestLayersFromDescriptors(t *testing.T, descriptors map[string]gemma4metadata.TensorDescriptor, config []byte) []manifest.Layer { + t.Helper() + + layers := make([]manifest.Layer, 0, len(descriptors)+1) + for name, descriptor := range descriptors { shape := make([]int64, len(descriptor.Shape)) for i, dim := range descriptor.Shape { shape[i] = int64(dim) @@ -1189,9 +1271,9 @@ func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { if err != nil { t.Fatal(err) } - data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{ + data, err := io.ReadAll(safetensors.BuildPackedSafetensorsReaderWithMetadata([]*safetensors.TensorData{ safetensors.NewTensorDataFromBytes(name, descriptor.Dtype, descriptor.Shape, make([]byte, int(payloadSize))), - })) + }, gemma4DescriptorMetadata(descriptor))) if err != nil { t.Fatal(err) } @@ -1203,7 +1285,6 @@ func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { Name: name, }) } - config := []byte(`{"text_config":{"hidden_size":6},"vision_config":{"hidden_size":4,"intermediate_size":8,"num_hidden_layers":2,"num_attention_heads":1,"num_key_value_heads":1,"head_dim":4,"default_output_length":1,"patch_size":2,"position_embedding_size":16,"pooling_kernel_size":1}}`) configDigest := createTestBlob(t, config) layers = append(layers, manifest.Layer{ MediaType: "application/vnd.ollama.image.json", @@ -1211,14 +1292,24 @@ func gemma4VisionManifestLayers(t *testing.T) []manifest.Layer { Size: int64(len(config)), Name: "config.json", }) - if descriptors, err := gemma4VisionTensorDescriptors(layers); err != nil { + if got, err := gemma4VisionTensorDescriptors(layers); err != nil { t.Fatalf("read Gemma4 descriptor fixture: %v", err) - } else if err := gemma4metadata.ValidateVisionInstalledInventory(*gemma4VisionConfig(2), descriptors); err != nil { - t.Fatalf("validate Gemma4 descriptor fixture: %v", err) + } else if len(got) != len(descriptors) { + t.Fatalf("hydrated descriptor count = %d, want %d", len(got), len(descriptors)) } return layers } +func gemma4DescriptorMetadata(descriptor gemma4metadata.TensorDescriptor) map[string]string { + if descriptor.QuantType == "" || descriptor.GroupSize <= 0 { + return nil + } + return map[string]string{ + "quant_type": descriptor.QuantType, + "group_size": strconv.Itoa(descriptor.GroupSize), + } +} + func TestOpenGemma4TensorLayerRejectsUnsafeHeaders(t *testing.T) { tests := []struct { name string @@ -1321,8 +1412,10 @@ func gemma4VisionConfig(layers int) *gemma4metadata.ConfigFile { func gemma4AudioConfig(layers int) *gemma4metadata.ConfigFile { return &gemma4metadata.ConfigFile{ - TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, - AudioTokenID: 7, + Architectures: []string{"Gemma4ForConditionalGeneration"}, + ModelType: "gemma4", + TextConfig: gemma4metadata.TextConfig{HiddenSize: 5, VocabSize: 32}, + AudioTokenID: 7, AudioConfig: &gemma4metadata.AudioConfig{ ModelType: "gemma4_audio", AttentionChunkSize: 2, AttentionContextLeft: 2, diff --git a/server/model_inference_cache_test.go b/server/model_inference_cache_test.go index 106fe8d0ccb..d03658ddfd0 100644 --- a/server/model_inference_cache_test.go +++ b/server/model_inference_cache_test.go @@ -204,6 +204,8 @@ func TestInferenceModelCacheGemma4AudioCapabilities(t *testing.T) { t.Fatal("cold model did not retain Gemma 4 audio metadata") } first.Gemma4AudioConfig.AudioConfig.HiddenSize = 0 + first.Gemma4AudioConfig.Architectures[0] = "mutated" + first.Gemma4AudioConfig.AudioConfig.SubsamplingConvChannels[0] = 0 mutatedTensor := false for name, descriptor := range first.Gemma4AudioTensors { if len(descriptor.Shape) == 0 { @@ -229,7 +231,10 @@ func TestInferenceModelCacheGemma4AudioCapabilities(t *testing.T) { if !slices.Contains(second.Capabilities(), model.CapabilityAudio) || !second.Gemma4AudioReady { t.Fatalf("cached state = capabilities:%v ready:%t, want audio", second.Capabilities(), second.Gemma4AudioReady) } - if second.Gemma4AudioConfig.AudioConfig.HiddenSize == 0 || len(second.Gemma4AudioTensors) == 0 { + if second.Gemma4AudioConfig.AudioConfig.HiddenSize == 0 || + second.Gemma4AudioConfig.Architectures[0] == "mutated" || + second.Gemma4AudioConfig.AudioConfig.SubsamplingConvChannels[0] == 0 || + len(second.Gemma4AudioTensors) == 0 { t.Fatal("cached Gemma 4 audio metadata was mutated") } for _, descriptor := range second.Gemma4AudioTensors { diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 0dbaa691052..04ff8290c9f 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -412,10 +412,10 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum } func filterUnsupportedModelListCapabilities(capabilities []model.Capability, cfg model.ConfigV2) []model.Capability { - if cfg.ModelFormat == "safetensors" && isGemma4Renderer(cfg.Renderer) && + if isGemma4SafetensorsConfig(cfg) && (cfg.RemoteHost != "" || cfg.RemoteModel != "") { capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { - return c == model.CapabilityAudio + return c == model.CapabilityVision || c == model.CapabilityAudio }) } if isNemotron3NanoSafetensorsConfig(cfg) { diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index b937cf9379c..7d64b2af0fa 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -178,6 +178,44 @@ func TestModelListSummaryGemma4SafetensorsVisionRequiresTensorLayers(t *testing. } } +func TestModelListSummaryGemma4UnifiedVisionRequiresTensorLayers(t *testing.T) { + setTestHome(t, t.TempDir()) + cfg := model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: gemma4RendererLarge, + Capabilities: []string{"completion", "vision", "audio", "tools", "thinking"}, + } + for _, tt := range []struct { + name string + layers []manifest.Layer + wantVision bool + }{ + {"complete", gemma4UnifiedVisionManifestLayers(t, true), true}, + {"partial", gemma4UnifiedVisionManifestLayers(t, false), false}, + } { + t.Run(tt.name, func(t *testing.T) { + modelName := "list-gemma4-unified-vision-" + tt.name + createSafetensorsTestModel(t, modelName, cfg, tt.layers) + mf, err := manifest.ParseNamedManifest(model.ParseName(modelName)) + if err != nil { + t.Fatal(err) + } + summary, err := buildModelListSummary(model.ParseName(modelName), mf) + if err != nil { + t.Fatal(err) + } + if got := slices.Contains(summary.Capabilities, model.CapabilityVision); got != tt.wantVision { + t.Fatalf("vision capability = %v, want %v (%v)", got, tt.wantVision, summary.Capabilities) + } + for _, capability := range []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking} { + if !slices.Contains(summary.Capabilities, capability) { + t.Fatalf("capabilities = %v, want preserved %s", summary.Capabilities, capability) + } + } + }) + } +} + func TestModelListSummaryGemma4AudioRequiresRuntimeMetadataAndTensors(t *testing.T) { setTestHome(t, t.TempDir()) cfg := model.ConfigV2{ModelFormat: "safetensors", Renderer: gemma4RendererSmall, Capabilities: []string{"completion", "audio"}} diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index fd4e44a16d9..9b2f438376d 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -655,6 +655,34 @@ func TestInferSafetensorsCapabilitiesGemma4UnifiedVision(t *testing.T) { check(t, wrong, false) } +func TestInferSafetensorsCapabilitiesGemma4Unified12B(t *testing.T) { + const configJSON = `{ + "architectures":["Gemma4UnifiedForConditionalGeneration"],"model_type":"gemma4_unified", + "text_config":{"hidden_size":3840}, + "vision_config":{"model_type":"gemma4_unified_vision","mm_embed_dim":3840,"mm_posemb_size":1120,"model_patch_size":48,"num_soft_tokens":280,"output_proj_dims":3840,"patch_size":16,"pooling_kernel_size":3,"rms_norm_eps":1e-6} + }` + tensors := map[string]gemma4metadata.TensorDescriptor{ + "model.vision_embedder.patch_ln1.weight": {Dtype: "BF16", Shape: []int32{6912}}, + "model.vision_embedder.patch_ln1.bias": {Dtype: "BF16", Shape: []int32{6912}}, + "model.vision_embedder.patch_dense.weight": {Dtype: "BF16", Shape: []int32{3840, 6912}}, + "model.vision_embedder.patch_dense.bias": {Dtype: "BF16", Shape: []int32{3840}}, + "model.vision_embedder.patch_ln2.weight": {Dtype: "BF16", Shape: []int32{3840}}, + "model.vision_embedder.patch_ln2.bias": {Dtype: "BF16", Shape: []int32{3840}}, + "model.vision_embedder.pos_embedding": {Dtype: "BF16", Shape: []int32{1120, 2, 3840}}, + "model.vision_embedder.pos_norm.weight": {Dtype: "BF16", Shape: []int32{3840}}, + "model.vision_embedder.pos_norm.bias": {Dtype: "BF16", Shape: []int32{3840}}, + "model.embed_vision.embedding_projection.weight": {Dtype: "BF16", Shape: []int32{3840, 3840}}, + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { + t.Fatal(err) + } + writeClientSafetensorDescriptors(t, dir, tensors) + if got := inferSafetensorsCapabilities(dir, ""); !slices.Contains(got, "vision") { + t.Fatalf("frozen unified-12B capabilities = %v, want vision", got) + } +} + func TestInferSafetensorsCapabilitiesGemma4AudioInventory(t *testing.T) { identities := []struct { name string