mlxrunner: add structured output support

The MLX runner accepted the API's format field but did not enforce it:
requests asking for JSON or a JSON Schema got unconstrained text, and
clients had no way to tell.

Enforce format with xgrammar: each sampling step masks the logits to
the tokens the grammar allows, so every emitted token and the end of
generation are valid under the constraint. Sampling, penalties, and
logprobs see the constrained distribution, and "json" yields a JSON
object, as the API documents and the llama-server path already
enforces. Only sampling waits on the mask; the forward pass is
dispatched before it, so constrained decoding stays pipelined.

The grammar engine is a dynamic library alongside MLX; when it is
missing, plain inference is unaffected and structured requests fail
with an explicit error. Constrained requests decode without
speculative decoding for now.

Decoding 256 tokens of a book-list schema on qwen3.8:27b-mlx (M5 Max,
seed 42, thinking off); pre-decode is the request time spent before
the first token:

    unconstrained              ~65 tok/s   pre-decode ~70 ms
    unconstrained, no draft    ~32 tok/s   pre-decode ~70 ms
    JSON schema                ~32 tok/s   pre-decode ~70 ms

Schema and draft-less decoding are equal to within 0.1 tok/s in
paired adjacent requests, and a cold grammar compile adds nothing
measurable to pre-decode. The gap to unconstrained decoding is the
disabled draft model.

Fixes #16563
Co-authored-by: Daniel Hiltgen <daniel@ollama.com>

.github/workflows/test.yaml

@@ -52,6 +52,8 @@ jobs:
             'ml/backend/ggml/ggml/**/*' \
             'x/imagegen/mlx/**' \
             'x/imagegen/mlx/**/*' \
+            'x/mlxrunner/xgrammar/native/**' \
+            'x/mlxrunner/xgrammar/native/**/*' \
             '.github/**/*') | tee -a $GITHUB_OUTPUT
           echo app_changed=$(changed 'app/**' 'app/**/*') | tee -a $GITHUB_OUTPUT
           echo enginehash=$(cat LLAMA_CPP_VERSION)-$(cat MLX_VERSION)-$(cat MLX_C_VERSION) | tee -a $GITHUB_OUTPUT
@@ -116,7 +118,7 @@ jobs:
             superbuild_target: ollama-mlx-cuda_v13
             superbuild_dir: build/local-superbuild-mlx-cuda_v13
             superbuild_args: '-DOLLAMA_MLX_BACKENDS=cuda_v13 -DCMAKE_CUDA_ARCHITECTURES=87 -DMLX_CUDA_ARCHITECTURES=80-virtual -DBLAS_INCLUDE_DIRS=/usr/include/x86_64-linux-gnu -DLAPACK_INCLUDE_DIRS=/usr/include/x86_64-linux-gnu'
-            expected_payload: lib/ollama/mlx_cuda_v13/libmlx.so
+            expected_payload: lib/ollama/mlx_cuda_v13/libmlx.so lib/ollama/mlx_cuda_v13/libollama_xgrammar.so
             install-go: true
     runs-on: linux
     container: ${{ matrix.container }}
@@ -158,7 +160,9 @@ jobs:
         run: |
           cmake -S . -B "${{ matrix.superbuild_dir }}" ${{ matrix.superbuild_args }}
           CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) cmake --build "${{ matrix.superbuild_dir }}" --target "${{ matrix.superbuild_target }}" -- -l $(nproc)
-          test -e "${{ matrix.superbuild_dir }}/${{ matrix.expected_payload }}"
+          for f in ${{ matrix.expected_payload }}; do
+            test -e "${{ matrix.superbuild_dir }}/$f"
+          done
       - name: Verify local superbuild install
         if: matrix.superbuild_target == 'ollama-local'
         run: |
@@ -214,7 +218,7 @@ jobs:
             superbuild_target: ollama-mlx-cuda_v13
             superbuild_dir: build\local-superbuild-mlx-cuda_v13
             superbuild_args: '-DOLLAMA_MLX_BACKENDS=cuda_v13 -DCMAKE_CUDA_ARCHITECTURES=80 -DMLX_CUDA_ARCHITECTURES=80-virtual'
-            expected_payload: lib\ollama\mlx_cuda_v13\mlx.dll
+            expected_payload: lib\ollama\mlx_cuda_v13\mlx.dll lib\ollama\mlx_cuda_v13\ollama_xgrammar.dll
             install-go: true
             cuda-components:
               - '"cudart"'
@@ -334,8 +338,10 @@ jobs:
           cmake -S . -B "${{ matrix.superbuild_dir }}" ${{ matrix.superbuild_args }}
           $env:CMAKE_BUILD_PARALLEL_LEVEL = [Environment]::ProcessorCount
           cmake --build "${{ matrix.superbuild_dir }}" --target "${{ matrix.superbuild_target }}" -- -l $([Environment]::ProcessorCount)
-          if (!(Test-Path "${{ matrix.superbuild_dir }}\${{ matrix.expected_payload }}")) {
-            throw "missing ${{ matrix.expected_payload }}"
+          foreach ($f in "${{ matrix.expected_payload }}".Split(' ')) {
+            if (!(Test-Path "${{ matrix.superbuild_dir }}\$f")) {
+              throw "missing $f"
+            }
           }
         env:
           CMAKE_GENERATOR: Ninja

Dockerfile

@@ -215,6 +215,7 @@ COPY CMakeLists.txt CMakePresets.json .
 COPY cmake cmake
 COPY mlx mlx
 COPY x/mlxrunner/mlx x/mlxrunner/mlx
+COPY x/mlxrunner/xgrammar/native x/mlxrunner/xgrammar/native
 COPY go.mod go.sum .
 COPY MLX_VERSION MLX_C_VERSION .
 RUN curl -fsSL https://golang.org/dl/go$(awk '/^go/ { print $2 }' go.mod).linux-$(case $(uname -m) in x86_64) echo amd64 ;; aarch64) echo arm64 ;; esac).tar.gz | tar xz -C /usr/local

cmake/local.cmake

@@ -204,6 +204,18 @@ if(OLLAMA_MLX_BACKENDS)
             USES_TERMINAL_PATCH TRUE)
         list(APPEND _mlx_source_targets ollama-mlx-c-source)
     endif()
+    # XGrammar has no pre-fetch: without an override each variant's build
+    # clones it via FetchContent.
+    if(DEFINED FETCHCONTENT_SOURCE_DIR_XGRAMMAR AND NOT "${FETCHCONTENT_SOURCE_DIR_XGRAMMAR}" STREQUAL "")
+        get_filename_component(OLLAMA_XGRAMMAR_SOURCE_DIR
+            "${FETCHCONTENT_SOURCE_DIR_XGRAMMAR}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}")
+        message(STATUS "Using XGrammar source override: ${OLLAMA_XGRAMMAR_SOURCE_DIR}")
+    elseif(DEFINED ENV{OLLAMA_XGRAMMAR_SOURCE})
+        get_filename_component(OLLAMA_XGRAMMAR_SOURCE_DIR
+            "$ENV{OLLAMA_XGRAMMAR_SOURCE}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}")
+        message(STATUS "Using local XGrammar source: ${OLLAMA_XGRAMMAR_SOURCE_DIR}")
+    endif()
+
     # Refresh the vendored MLX-C headers once the sources are present. Every MLX
     # backend variant shares this destination in the source tree, so the copy has
     # to happen here rather than in each variant's build.
@@ -493,6 +505,10 @@ function(ollama_add_mlx_build name)
         ${ARG_CMAKE_ARGS}
         ${_mlx_cache_args}
     )
+    if(OLLAMA_XGRAMMAR_SOURCE_DIR)
+        list(APPEND _cmake_args
+            -DFETCHCONTENT_SOURCE_DIR_XGRAMMAR=${OLLAMA_XGRAMMAR_SOURCE_DIR})
+    endif()
     foreach(_arg IN ITEMS
             BLAS_INCLUDE_DIRS
             LAPACK_INCLUDE_DIRS
@@ -534,6 +550,7 @@ function(ollama_add_mlx_build name)
             ${OLLAMA_NATIVE_CONFIG_ARG}
             ${OLLAMA_NATIVE_BUILD_TARGET_ARG} mlx
             ${OLLAMA_NATIVE_BUILD_TARGET_ARG} mlxc
+            ${OLLAMA_NATIVE_BUILD_TARGET_ARG} ollama_xgrammar
         INSTALL_COMMAND ${CMAKE_COMMAND} --install <BINARY_DIR>
             ${OLLAMA_NATIVE_CONFIG_ARG}
             --component MLX

cmake/mlx/CMakeLists.txt

@@ -61,6 +61,45 @@ foreach(_cudnn_var CUDNN_INCLUDE_PATH CUDNN_LIBRARY_PATH)
 endforeach()
 add_subdirectory(${OLLAMA_SOURCE_DIR}/x/mlxrunner/mlx ${CMAKE_BINARY_DIR}/x/mlxrunner/mlx)
 
+include(FetchContent)
+set(XGRAMMAR_VERSION v0.2.5)
+FetchContent_Declare(xgrammar
+    GIT_REPOSITORY "https://github.com/mlc-ai/xgrammar.git"
+    GIT_TAG ${XGRAMMAR_VERSION}
+    GIT_SHALLOW TRUE
+    GIT_SUBMODULES 3rdparty/dlpack
+    # Do not add XGrammar's Python-oriented CMake project.
+    SOURCE_SUBDIR cmake/ollama)
+FetchContent_MakeAvailable(xgrammar)
+
+file(GLOB_RECURSE XGRAMMAR_SOURCES CONFIGURE_DEPENDS "${xgrammar_SOURCE_DIR}/cpp/*.cc")
+list(FILTER XGRAMMAR_SOURCES EXCLUDE REGEX "/cpp/tvm_ffi/.*\\.cc$")
+add_library(xgrammar STATIC ${XGRAMMAR_SOURCES})
+set_target_properties(xgrammar PROPERTIES
+    POSITION_INDEPENDENT_CODE ON
+    CXX_VISIBILITY_PRESET hidden
+    VISIBILITY_INLINES_HIDDEN ON)
+target_include_directories(xgrammar PUBLIC "${xgrammar_SOURCE_DIR}/include")
+target_include_directories(xgrammar SYSTEM PUBLIC
+    "${xgrammar_SOURCE_DIR}/3rdparty/picojson"
+    "${xgrammar_SOURCE_DIR}/3rdparty/dlpack/include")
+target_compile_definitions(xgrammar PUBLIC
+    XGRAMMAR_ENABLE_CPPTRACE=0
+    XGRAMMAR_ENABLE_INTERNAL_CHECK=0)
+
+add_library(ollama_xgrammar SHARED
+    "${OLLAMA_SOURCE_DIR}/x/mlxrunner/xgrammar/native/xgrammar.cpp")
+target_include_directories(ollama_xgrammar PRIVATE
+    "${OLLAMA_SOURCE_DIR}/x/mlxrunner/xgrammar/native")
+target_compile_definitions(ollama_xgrammar PRIVATE
+    OLLAMA_XGRAMMAR_BUILD=1
+    OLLAMA_XGRAMMAR_VERSION="${XGRAMMAR_VERSION}")
+# Export only the ollama_xgrammar_* API; xgrammar's own C++ symbols stay hidden.
+set_target_properties(ollama_xgrammar PROPERTIES
+    CXX_VISIBILITY_PRESET hidden
+    VISIBILITY_INLINES_HIDDEN ON)
+target_link_libraries(ollama_xgrammar PRIVATE xgrammar Threads::Threads)
+
 # Find CUDA toolkit if MLX is built with CUDA support.
 find_package(CUDAToolkit)
 
@@ -98,12 +137,39 @@ endif()
 
 # Keep mlx/mlxc targets separate from runtime dependencies so --strip only
 # applies to the binaries we build, not vendor DLLs/libs.
-install(TARGETS mlx mlxc
+install(TARGETS mlx mlxc ollama_xgrammar
     RUNTIME_DEPENDENCY_SET mlx_runtime_deps
     RUNTIME DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX
     LIBRARY DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX
     FRAMEWORK DESTINATION ${OLLAMA_INSTALL_DIR} COMPONENT MLX
 )
+install(FILES
+    "${xgrammar_SOURCE_DIR}/LICENSE"
+    DESTINATION ${OLLAMA_INSTALL_DIR}
+    RENAME XGRAMMAR_LICENSE
+    COMPONENT MLX)
+install(FILES
+    "${xgrammar_SOURCE_DIR}/NOTICE"
+    DESTINATION ${OLLAMA_INSTALL_DIR}
+    RENAME XGRAMMAR_NOTICE
+    COMPONENT MLX)
+install(FILES
+    "${xgrammar_SOURCE_DIR}/3rdparty/dlpack/LICENSE"
+    DESTINATION ${OLLAMA_INSTALL_DIR}
+    RENAME DLPACK_LICENSE
+    COMPONENT MLX)
+file(READ "${xgrammar_SOURCE_DIR}/3rdparty/picojson/picojson.h" _picojson_header LIMIT 4096)
+string(FIND "${_picojson_header}" "*/" _picojson_license_end)
+if(_picojson_license_end EQUAL -1)
+    message(FATAL_ERROR "picojson license header not found")
+endif()
+math(EXPR _picojson_license_end "${_picojson_license_end} + 2")
+string(SUBSTRING "${_picojson_header}" 0 ${_picojson_license_end} _picojson_license)
+file(WRITE "${CMAKE_BINARY_DIR}/PICOJSON_LICENSE" "${_picojson_license}\n")
+install(FILES
+    "${CMAKE_BINARY_DIR}/PICOJSON_LICENSE"
+    DESTINATION ${OLLAMA_INSTALL_DIR}
+    COMPONENT MLX)
 install(RUNTIME_DEPENDENCY_SET mlx_runtime_deps
     DIRECTORIES ${MLX_RUNTIME_DIRS}
     PRE_INCLUDE_REGEXES ${MLX_INCLUDE_REGEXES}

integration/reg_release_test.go

@@ -2,6 +2,8 @@
 
 package integration
 
+import "testing"
+
 var (
 	releaseUnicodeInputModel    = integrationModel{Name: "deepseek-coder-v2:16b-lite-instruct-q2_K", MinVRAMGB: 12}
 	releaseUnicodeOutputModel   = "gemma2:2b"
@@ -64,6 +66,10 @@ var (
 
 const releaseSplitBatchVisionModel = "qwen3.5:2b"
 
+func TestStructuredOutput(t *testing.T) {
+	runIntegrationGroup(t, "structured-output")
+}
+
 func init() {
 	// Fixed release regression cases
 	registerIntegrationCases(
@@ -119,6 +125,7 @@ func init() {
 		integrationTestCase("create-gguf", "", runCreateGGUF),
 		integrationTestCase("quantization", "qwen2.5:0.5b-instruct-fp16", runQuantization),
 	)
+	registerStructuredOutputCases()
 
 	// Model-parametric cases
 	registerModelMinVRAM([]integrationModel{releaseUnicodeInputModel, releaseParallelHistoryModel})

integration/structured_output_test.go

@@ -0,0 +1,220 @@
+//go:build integration
+
+package integration
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"math"
+	"testing"
+
+	"github.com/ollama/ollama/api"
+)
+
+var structuredOutputSchema = json.RawMessage(`{
+  "type": "object",
+  "properties": {
+    "color": {"type": "string", "enum": ["blue", "violet"]},
+    "count": {"type": "integer", "minimum": 1, "maximum": 3}
+  },
+  "required": ["color", "count"],
+  "additionalProperties": false
+}`)
+
+func validateStructuredObject(t *testing.T, content string) {
+	t.Helper()
+	if !json.Valid([]byte(content)) {
+		t.Fatalf("response is not valid JSON: %q", content)
+	}
+	var object map[string]json.RawMessage
+	if err := json.Unmarshal([]byte(content), &object); err != nil {
+		t.Fatal(err)
+	}
+	if len(object) != 2 || object["color"] == nil || object["count"] == nil {
+		t.Fatalf("response has the wrong fields: %s", content)
+	}
+	var color string
+	if err := json.Unmarshal(object["color"], &color); err != nil || (color != "blue" && color != "violet") {
+		t.Fatalf("color = %q, %v", color, err)
+	}
+	var count int
+	if err := json.Unmarshal(object["count"], &count); err != nil || count < 1 || count > 3 {
+		t.Fatalf("count = %d, %v", count, err)
+	}
+}
+
+func validateConstrainedLogprobs(t *testing.T, logprobs []api.Logprob) {
+	t.Helper()
+	if len(logprobs) == 0 {
+		t.Fatal("constrained response did not include logprobs")
+	}
+	for i, entry := range logprobs {
+		if math.IsInf(entry.Logprob, 0) || math.IsNaN(entry.Logprob) {
+			t.Fatalf("logprob[%d] is not finite: %v", i, entry.Logprob)
+		}
+		for j, alternative := range entry.TopLogprobs {
+			if math.IsInf(alternative.Logprob, 0) || math.IsNaN(alternative.Logprob) {
+				t.Fatalf("logprob[%d].top_logprobs[%d] is not finite: %v", i, j, alternative.Logprob)
+			}
+		}
+	}
+}
+
+const structuredOutputMLXModel = "qwen3.5:2b-nvfp4"
+
+func registerStructuredOutputCases() {
+	registerModelMinVRAM([]integrationModel{{Name: structuredOutputMLXModel, MinVRAMGB: 4}})
+	registerModelIntegrationCases("structured-output", testModels([]string{smol, structuredOutputMLXModel}), runStructuredOutput)
+}
+
+func runStructuredOutput(t *testing.T, model string) {
+	skipRegisteredMinVRAM(t, model)
+	ctx, cancel := context.WithTimeout(context.Background(), apiTestTimeout)
+	defer cancel()
+	client, _, cleanup := InitServerConnection(ctx, t)
+	defer cleanup()
+	pullOrSkip(ctx, t, client, model)
+	noThink := api.ThinkValue{Value: false}
+	preloadGenerateModel(ctx, t, client, api.GenerateRequest{
+		Model:  model,
+		Prompt: "Respond with one word.",
+		Think:  &noThink,
+		Options: map[string]any{
+			"temperature": 0,
+			"num_predict": 1,
+		},
+	})
+
+	t.Run("generate schema adversarial prompt", func(t *testing.T) {
+		// The prompt asks for prose, so only grammar enforcement can make
+		// the response satisfy the schema.
+		stream := false
+		req := api.GenerateRequest{
+			Model:  model,
+			Prompt: "Say hi",
+			Stream: &stream,
+			Format: structuredOutputSchema,
+			Think:  &noThink,
+			Options: map[string]any{
+				"temperature": 0,
+				"seed":        17,
+				"num_predict": 96,
+			},
+		}
+		var content bytes.Buffer
+		if err := client.Generate(ctx, &req, func(response api.GenerateResponse) error {
+			content.WriteString(response.Response)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+		validateStructuredObject(t, content.String())
+	})
+
+	t.Run("generate builtin JSON streaming", func(t *testing.T) {
+		stream := true
+		req := api.GenerateRequest{
+			Model:  model,
+			Prompt: "Return the smallest possible JSON value. Output JSON only.",
+			Stream: &stream,
+			Format: json.RawMessage(`"json"`),
+			Think:  &noThink,
+			Options: map[string]any{
+				"temperature": 0,
+				"seed":        17,
+				"num_predict": 96,
+			},
+		}
+		var content bytes.Buffer
+		if err := client.Generate(ctx, &req, func(response api.GenerateResponse) error {
+			content.WriteString(response.Response)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+		if !json.Valid(content.Bytes()) {
+			t.Fatalf("response is not valid JSON: %q", content.String())
+		}
+	})
+
+	t.Run("generate schema greedy with logprobs", func(t *testing.T) {
+		stream := false
+		req := api.GenerateRequest{
+			Model:       model,
+			Prompt:      "Return an object with a color and a small count. Output JSON only.",
+			Stream:      &stream,
+			Format:      structuredOutputSchema,
+			Think:       &noThink,
+			Logprobs:    true,
+			TopLogprobs: 20,
+			Options: map[string]any{
+				"temperature": 0,
+				"seed":        23,
+				"num_predict": 96,
+			},
+		}
+		var content bytes.Buffer
+		var logprobs []api.Logprob
+		if err := client.Generate(ctx, &req, func(response api.GenerateResponse) error {
+			content.WriteString(response.Response)
+			logprobs = append(logprobs, response.Logprobs...)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+		validateStructuredObject(t, content.String())
+		validateConstrainedLogprobs(t, logprobs)
+	})
+
+	t.Run("generate schema sampled streaming", func(t *testing.T) {
+		stream := true
+		req := api.GenerateRequest{
+			Model:  model,
+			Prompt: "Return a color and count as a JSON object. Output JSON only.",
+			Stream: &stream,
+			Format: structuredOutputSchema,
+			Think:  &noThink,
+			Options: map[string]any{
+				"temperature": 0.7,
+				"seed":        29,
+				"num_predict": 96,
+			},
+		}
+		var content bytes.Buffer
+		if err := client.Generate(ctx, &req, func(response api.GenerateResponse) error {
+			content.WriteString(response.Response)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+		validateStructuredObject(t, content.String())
+	})
+
+	t.Run("chat schema greedy", func(t *testing.T) {
+		stream := false
+		req := api.ChatRequest{
+			Model: model,
+			Messages: []api.Message{{
+				Role:    "user",
+				Content: "Return an object with a color and a small count. Output JSON only.",
+			}},
+			Stream: &stream,
+			Format: structuredOutputSchema,
+			Think:  &noThink,
+			Options: map[string]any{
+				"temperature": 0,
+				"seed":        31,
+				"num_predict": 96,
+			},
+		}
+		var content bytes.Buffer
+		if err := client.Chat(ctx, &req, func(response api.ChatResponse) error {
+			content.WriteString(response.Message.Content)
+			return nil
+		}); err != nil {
+			t.Fatal(err)
+		}
+		validateStructuredObject(t, content.String())
+	})
+}

scripts/build_darwin.sh

@@ -115,7 +115,7 @@ _merge_darwin_payload() {
         [ -d "$AMD_VARIANT" ] || AMD_VARIANT=dist/darwin-amd64/lib/ollama
         mkdir -p "$DEST"
 
-        for LIB in libmlx.dylib libmlxc.dylib; do
+        for LIB in libmlx.dylib libmlxc.dylib libollama_xgrammar.dylib; do
             if [ -f "$AMD_VARIANT/$LIB" ] && [ -f "$VARIANT$LIB" ]; then
                 lipo -create -output "$DEST/$LIB" "$AMD_VARIANT/$LIB" "$VARIANT$LIB"
             elif [ -f "$VARIANT$LIB" ]; then
@@ -128,7 +128,7 @@ _merge_darwin_payload() {
         for F in "$VARIANT"*; do
             [ -f "$F" ] && [ ! -L "$F" ] || continue
             case "$(basename "$F")" in
-                libmlx.dylib|libmlxc.dylib) continue ;;
+                libmlx.dylib|libmlxc.dylib|libollama_xgrammar.dylib) continue ;;
             esac
             cp "$F" "$DEST/"
         done

x/mlxrunner/client.go

@@ -113,6 +113,7 @@ func (c *Client) WaitUntilRunning(ctx context.Context) error {
 type CompletionRequest struct {
 	Prompt      string
 	Media       []llm.MediaData
+	Format      json.RawMessage
 	Options     api.Options
 	Logprobs    bool
 	TopLogprobs int
@@ -157,6 +158,7 @@ func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn f
 	creq := CompletionRequest{
 		Prompt:      req.Prompt,
 		Media:       req.Media,
+		Format:      req.Format,
 		Logprobs:    req.Logprobs,
 		TopLogprobs: req.TopLogprobs,
 	}

x/mlxrunner/dflash_test.go

@@ -288,7 +288,7 @@ func TestDecodeBlockDraft(t *testing.T) {
 		t.Fatalf("open rejected a block-draft request")
 	}
 	pinDraftLimit(spec, 4)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}

x/mlxrunner/grammar.go

@@ -0,0 +1,414 @@
+package mlxrunner
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"math"
+	"net/http"
+	"path/filepath"
+	"slices"
+	"sync"
+	"unicode/utf8"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+	"github.com/ollama/ollama/x/mlxrunner/xgrammar"
+	"github.com/ollama/ollama/x/tokenizer"
+)
+
+const maxGrammarVocabSize = 1 << 20
+
+const (
+	maxGrammarSchemaBytes = 1 << 20
+	maxGrammarSchemaDepth = 128
+	// The token cap bounds grammar compile cost, which the serial runner
+	// pays as head-of-line blocking of the request queue.
+	maxGrammarSchemaTokens = 1 << 14
+)
+
+const (
+	grammarCompileThreads    = 8
+	grammarCompileCacheBytes = 128 << 20
+)
+
+// grammarEngine is the runner's structured-output subsystem: xgrammar bound
+// to the model's vocabulary, plus the pinned lookup table for expanding
+// packed token masks on the device.
+type grammarEngine struct {
+	// compileMu is the single compile slot: one native compile at a time
+	// bounds the engine's compile threads and memory, and close takes it to
+	// order the compiler's release against in-flight compiles.
+	compileMu sync.Mutex
+	compiler  *xgrammar.Compiler
+
+	// words is the width of one packed mask row, ceil(vocab/32).
+	words int
+
+	maskTable  *mlx.Array
+	byteShifts *mlx.Array
+}
+
+func newGrammarEngine(logitsWidth int, tokenizer *tokenizer.Tokenizer) *grammarEngine {
+	library, err := mlx.LoadedLibraryPath()
+	if err != nil {
+		slog.Warn("Structured output is unavailable", "error", err)
+		return nil
+	}
+	if err := validateGrammarVocab(logitsWidth, tokenizer.VocabSize()); err != nil {
+		slog.Warn("Structured output is unavailable", "error", err)
+		return nil
+	}
+	pieces := make([]string, logitsWidth)
+	for id := range logitsWidth {
+		pieces[id] = tokenizer.Decode([]int32{int32(id)})
+	}
+	stops := slices.DeleteFunc(slices.Clone(tokenizer.EOSTokens()), func(id int32) bool {
+		return id < 0 || int(id) >= logitsWidth
+	})
+	compiler, err := xgrammar.New(filepath.Dir(library), pieces, logitsWidth, stops, grammarCompileThreads, grammarCompileCacheBytes)
+	if err != nil {
+		slog.Warn("Structured output is unavailable", "error", err)
+		return nil
+	}
+	e := &grammarEngine{compiler: compiler}
+	e.initMask(logitsWidth)
+	slog.Info("Structured output initialized", "library", "xgrammar", "version", compiler.Version(), "vocab_size", logitsWidth, "path", compiler.Path())
+	return e
+}
+
+// The grammar vocabulary is always the logits width. A tokenizer longer
+// than the model's head (input-only tokens, e.g. Llama 3.2 Vision's image
+// tokens) is fine: those ids can never be sampled, so they stay out of the
+// grammar's vocabulary.
+func validateGrammarVocab(logitsWidth, tokenizerSize int) error {
+	if tokenizerSize <= 0 {
+		return fmt.Errorf("invalid tokenizer vocabulary size %d", tokenizerSize)
+	}
+	if logitsWidth <= 0 {
+		return fmt.Errorf("invalid model logits width %d", logitsWidth)
+	}
+	if logitsWidth > maxGrammarVocabSize {
+		return fmt.Errorf("model logits width %d exceeds structured output limit %d", logitsWidth, maxGrammarVocabSize)
+	}
+	return nil
+}
+
+// initMask builds the byte-to-mask lookup table for expanding packed token
+// masks on the device, pinned for the runner's lifetime: row v holds, for
+// each of the byte value v's eight bits low to high, 0 where the bit is set
+// (token allowed) and -inf where it is clear.
+func (e *grammarEngine) initMask(vocabSize int) {
+	e.words = (vocabSize + 31) / 32
+	vals := make([]float32, 256*8)
+	for v := range 256 {
+		for bit := range 8 {
+			if v>>bit&1 == 0 {
+				vals[v*8+bit] = float32(math.Inf(-1))
+			}
+		}
+	}
+	e.maskTable = mlx.FromValues(vals, 256, 8)
+	e.byteShifts = mlx.FromValues([]int32{0, 8, 16, 24}, 4)
+	mlx.Pin(e.maskTable, e.byteShifts)
+}
+
+func (e *grammarEngine) close() {
+	e.compileMu.Lock()
+	defer e.compileMu.Unlock()
+	if e.compiler != nil {
+		e.compiler.Close()
+		e.compiler = nil
+	}
+	mlx.Unpin(e.maskTable, e.byteShifts)
+	e.maskTable, e.byteShifts = nil, nil
+}
+
+// prepare parses a request format and, when it asks for structured output,
+// launches and returns the grammar compilation; a format that asks for none
+// returns nil. Safe on a nil subsystem, which reports structured output
+// unavailable.
+func (e *grammarEngine) prepare(format json.RawMessage) (*grammarCompilation, error) {
+	spec, err := parseGrammar(format)
+	if err != nil || spec == nil {
+		return nil, err
+	}
+	if e == nil {
+		return nil, api.StatusError{StatusCode: http.StatusNotImplemented, ErrorMessage: "structured output is unavailable"}
+	}
+	return e.compile(spec), nil
+}
+
+type grammarSpec struct {
+	kind   xgrammar.Kind
+	source string
+}
+
+func parseGrammar(format json.RawMessage) (*grammarSpec, error) {
+	if len(format) > 0 {
+		switch string(format) {
+		case `null`, `""`:
+			return nil, nil
+		case `"json"`:
+			// The API documents "json" as producing a JSON object; the engine's
+			// builtin JSON grammar would also admit arrays and bare values.
+			return &grammarSpec{kind: xgrammar.JSONSchema, source: `{"type":"object"}`}, nil
+		default:
+			if format[0] != '{' {
+				return nil, errors.New("invalid format: expected \"json\" or a valid JSON Schema object")
+			}
+			if err := validateGrammarSchema(format); err != nil {
+				return nil, fmt.Errorf("invalid JSON Schema: %w", err)
+			}
+			return &grammarSpec{kind: xgrammar.JSONSchema, source: string(format)}, nil
+		}
+	}
+	return nil, nil
+}
+
+func validateGrammarSchema(schema []byte) error {
+	if len(schema) > maxGrammarSchemaBytes {
+		return fmt.Errorf("schema is %d bytes; limit is %d", len(schema), maxGrammarSchemaBytes)
+	}
+	if !utf8.Valid(schema) {
+		return errors.New("schema is not valid UTF-8")
+	}
+
+	decoder := json.NewDecoder(bytes.NewReader(schema))
+	decoder.UseNumber()
+	first, err := decoder.Token()
+	if err != nil {
+		return fmt.Errorf("invalid JSON: %w", err)
+	}
+	if first != json.Delim('{') {
+		return errors.New("schema must be a JSON object")
+	}
+
+	tokens, depth := 1, 1
+	for depth > 0 {
+		token, err := decoder.Token()
+		if err != nil {
+			if err == io.EOF {
+				return errors.New("unexpected end of JSON")
+			}
+			return fmt.Errorf("invalid JSON: %w", err)
+		}
+		tokens++
+		if tokens > maxGrammarSchemaTokens {
+			return fmt.Errorf("schema contains more than %d JSON tokens", maxGrammarSchemaTokens)
+		}
+
+		if delim, ok := token.(json.Delim); ok {
+			switch delim {
+			case '{', '[':
+				depth++
+				if depth > maxGrammarSchemaDepth {
+					return fmt.Errorf("schema nesting exceeds %d levels", maxGrammarSchemaDepth)
+				}
+			case '}', ']':
+				depth--
+			}
+		}
+	}
+
+	if _, err := decoder.Token(); err != io.EOF {
+		if err != nil {
+			return fmt.Errorf("invalid JSON after schema object: %w", err)
+		}
+		return errors.New("schema contains more than one JSON value")
+	}
+	return nil
+}
+
+func (e *grammarEngine) compile(spec *grammarSpec) *grammarCompilation {
+	c := &grammarCompilation{done: make(chan struct{})}
+	go func() {
+		defer close(c.done)
+
+		e.compileMu.Lock()
+		defer e.compileMu.Unlock()
+		// A request cancelled while queued for the slot never compiles.
+		c.mu.Lock()
+		abandoned := c.abandoned
+		c.mu.Unlock()
+		if abandoned {
+			return
+		}
+		if e.compiler == nil {
+			c.err = errors.New("grammar engine closed")
+			return
+		}
+		matcher, err := e.compiler.Compile(spec.kind, spec.source)
+		if err != nil {
+			// A schema can pass validateGrammarSchema yet be rejected by
+			// the engine (e.g. an empty enum); that is still a request error.
+			c.err = api.StatusError{
+				StatusCode:   http.StatusBadRequest,
+				ErrorMessage: fmt.Sprintf("invalid structured output grammar: %v", err),
+			}
+			return
+		}
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		if c.abandoned {
+			matcher.Close()
+			return
+		}
+		c.grammar = &grammar{m: matcher}
+	}()
+	return c
+}
+
+// A grammarCompilation runs concurrently with prompt processing. resolve
+// blocks until it finishes; close releases the grammar without waiting, so a
+// cancelled request never holds the serial request loop through a compile.
+type grammarCompilation struct {
+	done chan struct{}
+	err  error // written by the compile goroutine, read only after done
+
+	// mu orders the compile's finish against close: whichever runs second
+	// frees the matcher.
+	mu        sync.Mutex
+	abandoned bool
+	grammar   *grammar
+}
+
+// A nil compilation resolves to no grammar.
+func (c *grammarCompilation) resolve(ctx context.Context) (*grammar, error) {
+	if c == nil {
+		return nil, nil
+	}
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-c.done:
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	// A Closed compilation must never resolve to an unconstrained decode.
+	if c.abandoned {
+		return nil, errors.New("grammar compilation abandoned")
+	}
+	return c.grammar, c.err
+}
+
+func (c *grammarCompilation) close() {
+	if c == nil {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.abandoned = true
+	c.grammar.close()
+	c.grammar = nil
+}
+
+// grammar is a request's compiled grammar, the runner's seam to the engine.
+// Methods are safe on a nil grammar, which never constrains.
+type grammar struct {
+	m *xgrammar.Matcher
+}
+
+// constraining reports whether sampling is currently constrained, read from
+// the matcher's state: a grammar constrains from its first token until its
+// state machine terminates. Kinds that trigger mid-response will decide
+// this from richer matcher state.
+func (g *grammar) constraining() bool {
+	return g != nil && !g.m.Terminated()
+}
+
+func (g *grammar) close() {
+	if g != nil {
+		g.m.Close()
+	}
+}
+
+// hasGrammar reports whether any row carries a grammar — the read that
+// decides whether a step takes the deferred shape: grammar work needs the
+// committed token values on the host, so the step's sample cannot fuse
+// onto the forward's chain.
+func (e *grammarEngine) hasGrammar(grammars []*grammar) bool {
+	for _, g := range grammars {
+		if g != nil {
+			return true
+		}
+	}
+	return false
+}
+
+// accept advances a batch's grammars over its newly committed tokens —
+// every row's grammar accepts the row's token, constraining or not, with
+// grammars[i] and committed[i] aligned. Each token was sampled under its
+// own matcher's mask, so a rejection is an engine fault; errs[i] carries
+// row i's fault, which ends that request rather than the runner, and errs
+// is nil when every row succeeded.
+func (e *grammarEngine) accept(grammars []*grammar, committed []int32) []error {
+	var errs []error
+	for i, g := range grammars {
+		if g == nil {
+			continue
+		}
+		if err := g.m.Accept(committed[i]); err != nil {
+			if errs == nil {
+				errs = make([]error, len(grammars))
+			}
+			errs[i] = fmt.Errorf("grammar: accept sampled token %d: %w", committed[i], err)
+		}
+	}
+	return errs
+}
+
+// mask fills each constraining row's packed token mask and applies them to
+// the batch's logits in one device op, logits row i masked under
+// grammars[i]. Unconstrained, terminated, and failed rows keep their logits
+// (all-ones mask rows add zero). errs as in accept.
+func (e *grammarEngine) mask(grammars []*grammar, logits *mlx.Array) (*mlx.Array, []error) {
+	var errs []error
+	packed := make([]int32, len(grammars)*e.words)
+	for i := range packed {
+		packed[i] = -1
+	}
+	apply := false
+	for i, g := range grammars {
+		if !g.constraining() {
+			continue
+		}
+		row := packed[i*e.words : (i+1)*e.words]
+		constrained, err := g.m.Fill(row)
+		if err != nil {
+			if errs == nil {
+				errs = make([]error, len(grammars))
+			}
+			errs[i] = fmt.Errorf("grammar: fill token mask: %w", err)
+			continue
+		}
+		// An all-zero mask would send every logit to -inf and sampling to NaN.
+		if constrained && !slices.ContainsFunc(row, func(w int32) bool { return w != 0 }) {
+			if errs == nil {
+				errs = make([]error, len(grammars))
+			}
+			errs[i] = errors.New("grammar: token mask rejects every vocabulary token")
+			continue
+		}
+		apply = apply || constrained
+	}
+	if !apply {
+		return logits, errs
+	}
+	return e.apply(logits, mlx.FromValues(packed, len(grammars), e.words)), errs
+}
+
+// apply masks logits under packed token masks, one row per sequence: logits
+// is [B, V], packed is [B, words], and disallowed tokens come back -inf.
+// Each mask word is split into bytes, each byte gathers its eight mask
+// values from the table, and the flattened rows are added to the logits.
+func (e *grammarEngine) apply(logits, packed *mlx.Array) *mlx.Array {
+	maskBytes := packed.ExpandDims(-1).RightShift(e.byteShifts).BitwiseAnd(mlx.FromValue(255))
+	mask := e.maskTable.TakeAxis(maskBytes, 0).Flatten(1, 3)
+	mask = mask.Slice(mlx.Slice(), mlx.Slice(0, logits.Dim(1))).AsType(logits.DType())
+	return logits.Add(mask)
+}

x/mlxrunner/grammar_mask_test.go

@@ -0,0 +1,48 @@
+package mlxrunner
+
+import (
+	"math"
+	"testing"
+
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+)
+
+func TestApplyTokenMask(t *testing.T) {
+	skipIfNoMLX(t)
+	const (
+		bitsPerMaskWord       = 32
+		firstTokenID          = 0                   // Least-significant bit of the first mask word.
+		interiorTokenID       = 7                   // Last bit of the first mask word's low byte.
+		lastIDInFirstMaskWord = bitsPerMaskWord - 1 // Sign bit of the int32-backed first mask word.
+		lastVocabID           = 40                  // Final valid ID in a partially used second mask word.
+		vocabSize             = lastVocabID + 1
+	)
+	allowedIDs := []int{firstTokenID, interiorTokenID, lastIDInFirstMaskWord, lastVocabID}
+	word0 := uint32(1)<<firstTokenID |
+		uint32(1)<<interiorTokenID |
+		uint32(1)<<lastIDInFirstMaskWord
+	packed := []int32{
+		int32(word0),
+		int32(uint32(1) << (lastVocabID - bitsPerMaskWord)),
+	}
+	e := &grammarEngine{}
+	e.initMask(vocabSize)
+	logits := mlx.Zeros(mlx.DTypeFloat32, 1, vocabSize)
+	masked := e.apply(logits, mlx.FromValues(packed, 1, len(packed)))
+	mlx.Eval(masked)
+	got := masked.Floats()
+	for id := range vocabSize {
+		allowed := false
+		for _, a := range allowedIDs {
+			if id == a {
+				allowed = true
+			}
+		}
+		if allowed && got[id] != 0 {
+			t.Fatalf("allowed token %d masked to %v", id, got[id])
+		}
+		if !allowed && !math.IsInf(float64(got[id]), -1) {
+			t.Fatalf("disallowed token %d = %v, want -Inf", id, got[id])
+		}
+	}
+}

x/mlxrunner/grammar_test.go

@@ -0,0 +1,228 @@
+package mlxrunner
+
+import (
+	"encoding/json"
+	"errors"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/x/mlxrunner/batch"
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+	"github.com/ollama/ollama/x/mlxrunner/xgrammar"
+)
+
+func TestParseGrammar(t *testing.T) {
+	tests := []struct {
+		name    string
+		format  string
+		kind    xgrammar.Kind
+		source  string
+		want    bool
+		wantErr bool
+	}{
+		{name: "unset"},
+		{name: "null", format: `null`},
+		{name: "empty", format: `""`},
+		{name: "json", format: `"json"`, kind: xgrammar.JSONSchema, source: `{"type":"object"}`, want: true},
+		{name: "schema", format: `{"type":"integer"}`, kind: xgrammar.JSONSchema, source: `{"type":"integer"}`, want: true},
+		{name: "unsupported string", format: `"xml"`, wantErr: true},
+		{name: "whitespace JSON", format: ` "json" `, wantErr: true},
+		{name: "whitespace schema", format: ` {"type":"integer"} `, wantErr: true},
+		{name: "array", format: `[]`, wantErr: true},
+		{name: "invalid json", format: `{`, wantErr: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseGrammar(json.RawMessage(tt.format))
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("parseGrammar error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if tt.wantErr {
+				return
+			}
+			if (got != nil) != tt.want {
+				t.Fatalf("parseGrammar = %#v, want present %v", got, tt.want)
+			}
+			if got != nil && (got.kind != tt.kind || got.source != tt.source) {
+				t.Errorf("parseGrammar = %#v, want kind %v source %q", got, tt.kind, tt.source)
+			}
+		})
+	}
+}
+
+func TestParseGrammarDoesNotEchoOversizedInput(t *testing.T) {
+	format := json.RawMessage("{" + strings.Repeat("x", maxGrammarSchemaBytes))
+	_, err := parseGrammar(format)
+	if err == nil || !strings.Contains(err.Error(), "schema is 1048577 bytes") {
+		t.Fatalf("parseGrammar error = %v, want bounded size error", err)
+	}
+	if len(err.Error()) > 256 {
+		t.Fatalf("parseGrammar echoed oversized input in %d-byte error", len(err.Error()))
+	}
+}
+
+func nestedGrammarSchema(depth int) []byte {
+	return []byte(`{"value":` + strings.Repeat("[", depth-1) + `0` + strings.Repeat("]", depth-1) + `}`)
+}
+
+func grammarSchemaArray(entries int) []byte {
+	var b strings.Builder
+	b.WriteString(`{"enum":[`)
+	for i := range entries {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		b.WriteByte('0')
+	}
+	b.WriteString(`]}`)
+	return []byte(b.String())
+}
+
+func TestValidateGrammarSchema(t *testing.T) {
+	invalidUTF8 := append([]byte(`{"value":"`), 0xff)
+	invalidUTF8 = append(invalidUTF8, []byte(`"}`)...)
+	tests := []struct {
+		name    string
+		schema  []byte
+		wantErr string
+	}{
+		{name: "object", schema: []byte(`{"type":"object","properties":{"answer":{"type":"string"}}}`)},
+		{name: "maximum depth", schema: nestedGrammarSchema(maxGrammarSchemaDepth)},
+		{name: "too deep", schema: nestedGrammarSchema(maxGrammarSchemaDepth + 1), wantErr: "nesting exceeds"},
+		{name: "too many tokens", schema: grammarSchemaArray(maxGrammarSchemaTokens), wantErr: "more than 16384 JSON tokens"},
+		{name: "too large", schema: []byte(`{"value":"` + strings.Repeat("x", maxGrammarSchemaBytes) + `"}`), wantErr: "limit is 1048576"},
+		{name: "invalid UTF-8", schema: invalidUTF8, wantErr: "not valid UTF-8"},
+		{name: "trailing value", schema: []byte(`{} {}`), wantErr: "more than one JSON value"},
+		{name: "malformed", schema: []byte(`{"type":`), wantErr: "unexpected end"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := validateGrammarSchema(tt.schema)
+			if tt.wantErr == "" {
+				if err != nil {
+					t.Fatal(err)
+				}
+				return
+			}
+			if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
+				t.Fatalf("validateGrammarSchema error = %v, want containing %q", err, tt.wantErr)
+			}
+		})
+	}
+}
+
+func BenchmarkValidateGrammarSchema(b *testing.B) {
+	schema := []byte(`{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`)
+	b.ReportAllocs()
+	for b.Loop() {
+		if err := validateGrammarSchema(schema); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func FuzzParseGrammar(f *testing.F) {
+	for _, format := range []string{
+		"",
+		`"json"`,
+		`{"type":"integer"}`,
+		`{`,
+	} {
+		f.Add(format)
+	}
+
+	f.Fuzz(func(t *testing.T, format string) {
+		spec, err := parseGrammar(json.RawMessage(format))
+		if err != nil || spec == nil {
+			return
+		}
+		switch spec.kind {
+		case xgrammar.JSONSchema:
+		default:
+			t.Fatalf("parseGrammar kind = %d", spec.kind)
+		}
+	})
+}
+
+func TestValidateGrammarVocab(t *testing.T) {
+	for _, tt := range []struct {
+		name      string
+		logits    int
+		tokenizer int
+		wantErr   bool
+	}{
+		{name: "exact fit", logits: 32, tokenizer: 32},
+		{name: "padded model head", logits: 40, tokenizer: 32},
+		{name: "input-only tokens past the head", logits: 31, tokenizer: 32},
+		{name: "invalid tokenizer", logits: 32, tokenizer: -1, wantErr: true},
+		{name: "invalid logits width", logits: 0, tokenizer: 32, wantErr: true},
+		{name: "allocation bound", logits: maxGrammarVocabSize + 1, tokenizer: 32, wantErr: true},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			err := validateGrammarVocab(tt.logits, tt.tokenizer)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("validateGrammarVocab(%d, %d) error = %v, wantErr %v", tt.logits, tt.tokenizer, err, tt.wantErr)
+			}
+		})
+	}
+}
+
+// resolvedGrammarCompilation wraps an already-built matcher as a finished
+// compilation.
+func resolvedGrammarCompilation(m *xgrammar.Matcher) *grammarCompilation {
+	c := &grammarCompilation{done: make(chan struct{}), grammar: &grammar{m: m}}
+	close(c.done)
+	return c
+}
+
+func TestPrepareGrammarUnavailable(t *testing.T) {
+	r := &Runner{
+		Model:         textOnlyModel{},
+		Tokenizer:     newTestTokenizer(t, []int32{7}),
+		contextLength: 32,
+	}
+	request := &Request{CompletionRequest: CompletionRequest{
+		Prompt: "0",
+		Format: json.RawMessage(`"json"`),
+	}}
+	err := r.Prepare(request)
+	var statusErr api.StatusError
+	if !errors.As(err, &statusErr) {
+		t.Fatalf("Prepare error = %T %v, want api.StatusError", err, err)
+	}
+	if statusErr.StatusCode != http.StatusNotImplemented {
+		t.Fatalf("status = %d, want %d", statusErr.StatusCode, http.StatusNotImplemented)
+	}
+}
+
+type grammarTestDrafter struct{}
+
+func (grammarTestDrafter) open([]any) draftSession { return grammarTestDraftSession{} }
+func (grammarTestDrafter) draftLimit() int         { return 0 }
+
+type grammarTestDraftSession struct{}
+
+func (grammarTestDraftSession) propose(*mlx.Array, int) *draftCandidates { return nil }
+func (grammarTestDraftSession) committed(*mlx.Array, *mlx.Array, int, []batch.MediaItem) {
+}
+func (grammarTestDraftSession) settle(*mlx.Array) {}
+func (grammarTestDraftSession) close()            {}
+
+func TestGrammarParksSpeculation(t *testing.T) {
+	s := &speculation{drafter: grammarTestDrafter{}, depth: newDepthController()}
+	constrained := s.open(Request{Grammar: resolvedGrammarCompilation(&xgrammar.Matcher{})}, nil)
+	defer constrained.close()
+	if constrained.enabled {
+		t.Fatal("structured output request enabled speculative decoding")
+	}
+
+	unconstrained := s.open(Request{}, nil)
+	defer unconstrained.close()
+	if !unconstrained.enabled {
+		t.Fatal("ordinary request unexpectedly disabled speculative decoding")
+	}
+}

x/mlxrunner/mlx/ops.go

@@ -71,6 +71,12 @@ func (t *Array) AsStrided(shape []int, strides []int, offset int) *Array {
 	return out
 }
 
+func (t *Array) BitwiseAnd(other *Array) *Array {
+	out := New("BITWISE_AND")
+	C.mlx_bitwise_and(&out.ctx, t.ctx, other.ctx, DefaultStream().ctx)
+	return out
+}
+
 func (t *Array) Concatenate(axis int, others ...*Array) *Array {
 	if len(others) == 0 {
 		return t.Clone()
@@ -214,6 +220,12 @@ func (t *Array) Reshape(axes ...int) *Array {
 	return out
 }
 
+func (t *Array) RightShift(other *Array) *Array {
+	out := New("RIGHT_SHIFT")
+	C.mlx_right_shift(&out.ctx, t.ctx, other.ctx, DefaultStream().ctx)
+	return out
+}
+
 func (t *Array) Sigmoid() *Array {
 	out := New("SIGMOID")
 	C.mlx_sigmoid(&out.ctx, t.ctx, DefaultStream().ctx)

x/mlxrunner/mtp_test.go

@@ -398,7 +398,7 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	d := testDecoder(r, req, caches, []int32{1}, position)
+	d := testDecoder(t, r, req, caches, []int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -464,7 +464,7 @@ func TestRunMTPDecodeSampled(t *testing.T) {
 		t.Fatalf("open rejected a sampled request")
 	}
 	pinDraftLimit(spec, 4)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -509,7 +509,7 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 	// hidden row, leaving the drafter ready to propose from slot 1.
 	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -573,7 +573,7 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
 	spec.limit = 4
 	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -620,7 +620,7 @@ func TestDecodePlain(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	d := testDecoder(r, req, caches, []int32{1}, position)
+	d := testDecoder(t, r, req, caches, []int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -678,7 +678,7 @@ func TestDecodeCancelledMidStream(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	d := testDecoder(r, req, caches, []int32{1}, position)
+	d := testDecoder(t, r, req, caches, []int32{1}, position)
 	err := r.decode(ctx, req, session, d, 0)
 	d.close()
 	if !errors.Is(err, context.Canceled) {
@@ -723,7 +723,7 @@ func TestLayoutRidesEveryForward(t *testing.T) {
 	}
 	spec := r.spec.open(req, []any{"layout"})
 	pinDraftLimit(spec, 4)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 1)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 1, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -761,14 +761,14 @@ func pinDraftLimit(spec *speculationSession, limit int) {
 // testDecoder builds the decoder TextGenerationPipeline would construct for
 // this request, with the draft length pinned to a fixed width; tests close it
 // explicitly so close-time effects are visible to assertions.
-func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
+func testDecoder(t *testing.T, r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
 	if spec := r.spec.open(req, nil); spec != nil {
 		if spec.enabled {
 			pinDraftLimit(spec, 4)
 		}
-		return spec.decoder(mlx.FromValues(seed, len(seed)), position)
+		return spec.decoder(mlx.FromValues(seed, len(seed)), position, nil)
 	}
-	return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position, nil)
+	return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position, nil, nil)
 }
 
 func TestDecodeKVDraft(t *testing.T) {
@@ -803,7 +803,7 @@ func TestDecodeKVDraft(t *testing.T) {
 	}
 	pinDraftLimit(spec, 4)
 	defer spec.close()
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -885,7 +885,7 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 	spec := r.spec.open(req, nil)
 	pinDraftLimit(spec, 4)
 	defer spec.close()
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -956,7 +956,7 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 	if spec == nil || spec.enabled {
 		t.Fatalf("want a permanent-park speculationSession, got %+v", spec)
 	}
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -1136,7 +1136,7 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
 	}
 	spec := r.spec.open(req, nil)
 	pinDraftLimit(spec, 4)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
@@ -1193,7 +1193,7 @@ func TestDecodeParkedDraftResume(t *testing.T) {
 		t.Fatalf("want a drafting speculationSession, got %+v", spec)
 	}
 	pinDraftLimit(spec, 0)
-	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0).(*speculativeDecoder)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0, nil).(*speculativeDecoder)
 
 	// Two parked calls arrive pipelined, one token each.
 	for _, want := range []int32{2, 3} {

x/mlxrunner/pipeline.go

@@ -26,11 +26,25 @@ 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) error {
+func (r *Runner) Prepare(request *Request) (err error) {
 	if r.Model == nil {
 		return errors.New("model not loaded")
 	}
 
+	// Launched first so the compile overlaps tokenization and media
+	// preparation as well as prefill.
+	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
 	if len(request.Media) == 0 {
@@ -113,11 +127,16 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 	// Register the sampler after prefill completes.
 	r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs)
 
+	grammar, err := request.Grammar.resolve(ctx)
+	if err != nil {
+		return err
+	}
+
 	var d decoder
 	if spec != nil {
-		d = spec.decoder(seed, position)
+		d = spec.decoder(seed, position, grammar)
 	} else {
-		d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position, media.rowLayout())
+		d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position, media.rowLayout(), grammar)
 	}
 	defer d.close()
 	return r.decode(ctx, request, session, d, promptEval)
@@ -324,31 +343,94 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess
 	}
 }
 
-// pipelinedDecoder decodes one token per call, one call ahead of emission:
-// the next token's chain is dispatched before the returned one is
-// synchronized, so the device runs ahead of host emission.
+// pipelinedDecoder decodes one token per row per call, one call ahead of
+// emission: the forward for the next tokens is dispatched before the
+// returned ones are synchronized, so the device runs ahead of host
+// emission. While no grammar constrains, the next sample is fused onto the
+// forward's chain. A constraining grammar's sample needs a token mask that
+// depends on the returned token's value, so only the forward runs ahead
+// and the host's grammar work overlaps it.
 type pipelinedDecoder struct {
 	r *Runner
 	// spec, when non-nil, receives every forwarded token and settles its
 	// drafter at close, keeping a non-drafting session's draft KV level.
 	spec     *speculationSession
 	caches   []cache.Cache
-	layout   []any // the request's per-row layout state, stamped on every forward
+	layout   []any      // the request's per-row layout state, stamped on every forward
+	grammars []*grammar // row i's grammar; nil rows are unconstrained
 	position int
-	sample   sampler.Result // in flight: sampled, not yet forwarded
+	pending  sampler.Result // in flight: sampled, not yet forwarded
+	// Steps run ahead asynchronously: when one faults, its token is already
+	// forwarded and still has to be returned, so err waits for the next call.
+	err error
 }
 
-func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int, layout []any) *pipelinedDecoder {
-	t := &pipelinedDecoder{r: r, spec: spec, caches: caches, layout: layout, position: position}
-	t.sample = t.dispatch(seed)
+func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int, layout []any, g *grammar) *pipelinedDecoder {
+	t := &pipelinedDecoder{
+		r: r, spec: spec, caches: caches, layout: layout, position: position,
+		grammars: []*grammar{g},
+	}
+	logits := t.forward(seed)
+	mlx.Pin(logits)
+	defer mlx.Unpin(logits)
+
+	if r.grammarEngine.hasGrammar(t.grammars) {
+		// Dispatch the forward before the host builds the first masks. The
+		// first sample commits nothing, so there is nothing to accept. A mask
+		// fault here is a step fault like any other: the seed is already
+		// forwarded, so the error waits for the first call.
+		mlx.Sweep()
+		mlx.AsyncEval(logits)
+		var errs []error
+		logits, errs = r.grammarEngine.mask(t.grammars, logits)
+		t.err = t.failRows(errs)
+	}
+	t.pending = t.sample(logits)
 	return t
 }
 
-// dispatch builds one forward-and-sample chain without reading the token's
-// value, so it is in flight before the previous token is synchronized.
-func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result {
-	r := t.r
-	hidden, auxHidden := r.Model.Forward(&batch.Batch{
+// failRows clears the grammars of rows whose grammar work faulted — a dead
+// row does no further grammar work — and joins their errors.
+func (t *pipelinedDecoder) failRows(errs []error) error {
+	for i, err := range errs {
+		if err != nil {
+			t.grammars[i] = nil
+		}
+	}
+	return errors.Join(errs...)
+}
+
+func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
+	if t.err != nil {
+		return nil, t.err
+	}
+	out := t.pending
+	logits := t.forward(out.Token.ExpandDims(-1))
+	mlx.Pin(logits)
+	defer mlx.Unpin(logits)
+
+	if t.r.grammarEngine.hasGrammar(t.grammars) {
+		// Dispatch the forward before the host's grammar work.
+		mlx.Sweep()
+		mlx.AsyncEval(logits)
+
+		err := t.failRows(t.r.grammarEngine.accept(t.grammars, out.Token.Ints()))
+
+		var errs []error
+		logits, errs = t.r.grammarEngine.mask(t.grammars, logits)
+		t.err = errors.Join(err, t.failRows(errs))
+	}
+
+	t.pending = t.sample(logits)
+
+	mlx.Unpin(out.Arrays()...)
+	return []sampler.Result{out}, nil
+}
+
+// forward runs the model one step over token, shaped [B, L], and returns the
+// final position's logits, still lazy.
+func (t *pipelinedDecoder) forward(token *mlx.Array) *mlx.Array {
+	hidden, auxHidden := t.r.Model.Forward(&batch.Batch{
 		InputIDs:     token,
 		SeqOffsets:   []int32{int32(t.position)},
 		SeqQueryLens: []int32{int32(token.Dim(1))},
@@ -356,33 +438,33 @@ func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result {
 	}, t.caches)
 	t.spec.committed(token, auxHidden, t.position, nil)
 	t.position += token.Dim(1)
-	logits := r.Model.Unembed(hidden)
-	next := r.Sampler.Sample([]int{pipelineSlot}, logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1))
+	logits := t.r.Model.Unembed(hidden)
+	return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1)
+}
+
+// sample dispatches the batched sample over the decoder's rows. On an
+// unconstrained step it fuses onto the forward's chain, so the whole step
+// is in flight before the previous tokens are synchronized.
+func (t *pipelinedDecoder) sample(logits *mlx.Array) sampler.Result {
+	next := t.r.Sampler.Sample([]int{pipelineSlot}, logits)
 	mlx.Pin(next.Arrays()...)
 	mlx.Sweep()
 	mlx.AsyncEval(next.Arrays()...)
 	return next
 }
 
-func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
-	out := t.sample
-	t.sample = t.dispatch(out.Token.ExpandDims(-1))
-	mlx.Unpin(out.Arrays()...)
-	return []sampler.Result{out}, nil
-}
-
 // drain ends production: it returns the in-flight sample (sampled but never
 // forwarded) and the position its forward would have taken. The decoder
 // keeps the sample for close.
 func (t *pipelinedDecoder) drain() ([]sampler.Result, int) {
-	return []sampler.Result{t.sample}, t.position
+	return []sampler.Result{t.pending}, t.position
 }
 
 func (t *pipelinedDecoder) close() {
 	// The in-flight sample's forward was never dispatched; its report settles
 	// the drafter level with the caches' resting offset.
-	t.spec.settle(t.sample.Token)
-	mlx.Unpin(t.sample.Arrays()...)
+	t.spec.settle(t.pending.Token)
+	mlx.Unpin(t.pending.Arrays()...)
 }
 
 // detokenizer serializes sampled tokens into response chunks, holding bytes
@@ -397,10 +479,24 @@ type detokenizer struct {
 	wantTopLogprobs int
 }
 
+// clampLogprobs floors logprobs at -9999, the OpenAI-compatible bound;
+// tokens a grammar masked out would otherwise report -Inf, which JSON
+// cannot encode.
+func clampLogprobs(logprobs []llm.Logprob) {
+	for i := range logprobs {
+		logprobs[i].Logprob = max(logprobs[i].Logprob, -9999)
+		for j := range logprobs[i].TopLogprobs {
+			logprobs[i].TopLogprobs[j].Logprob = max(logprobs[i].TopLogprobs[j].Logprob, -9999)
+		}
+	}
+}
+
 func (d *detokenizer) detokenize(res sampler.Result) (CompletionResponse, bool) {
 	output := res.Token.Int()
 	d.buf.WriteString(d.tokenizer.Decode([]int32{output}))
-	d.logprobs = append(d.logprobs, buildLogprob(res, d.wantLogprobs, d.wantTopLogprobs, d.tokenizer.Decode)...)
+	logprobs := buildLogprob(res, d.wantLogprobs, d.wantTopLogprobs, d.tokenizer.Decode)
+	clampLogprobs(logprobs)
+	d.logprobs = append(d.logprobs, logprobs...)
 
 	content := flushValidUTF8Prefix(&d.buf)
 	if content == "" {

x/mlxrunner/runner.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/x/internal/mlxthread"
+	"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"
@@ -35,6 +36,7 @@ type Request struct {
 	MediaItems  []mediaItem
 	Layout      any // opaque PrepareMedia layout state, stamped on every batch
 	SamplerOpts sample.Options
+	Grammar     *grammarCompilation
 }
 
 type Runner struct {
@@ -45,6 +47,9 @@ type Runner struct {
 	cache         *prefixCache
 	contextLength int
 	mlxThread     *mlxthread.Thread
+	// grammarEngine is the structured-output subsystem; nil when the grammar
+	// library or vocabulary failed to load.
+	grammarEngine *grammarEngine
 	// spec is the speculative-decoding subsystem. Nil when the model ships no
 	// draft head.
 	spec *speculation
@@ -115,12 +120,20 @@ func (r *Runner) Load(modelName string) error {
 	r.cache = newPrefixCache(slices.Concat(caches, draftCaches))
 	r.Sampler = sample.New(r.contextLength)
 	r.spec = newSpeculation(r, draftModel, caches, draftCaches)
+	r.grammarEngine = newGrammarEngine(logitsWidth(m), r.Tokenizer)
 
 	mlx.EnableCompile()
 
 	return nil
 }
 
+func (r *Runner) Close() {
+	if r.grammarEngine != nil {
+		r.grammarEngine.close()
+		r.grammarEngine = nil
+	}
+}
+
 // newDraftCaches returns nil when the model ships no draft.
 func newDraftCaches(draft base.DraftModel) []cache.Cache {
 	if draft == nil {
@@ -129,6 +142,29 @@ func newDraftCaches(draft base.DraftModel) []cache.Cache {
 	return draft.NewCaches()
 }
 
+// logitsWidth reads a model's logits width off a one-token forward's static
+// shape — the same Forward and Unembed path decode logits take. Nothing is
+// evaluated, and the probe's caches and graph are released before returning,
+// which sweeps every unpinned array: call this only at load, after the
+// model's weights are pinned.
+func logitsWidth(m base.Model) int {
+	caches := m.NewCaches()
+	hidden, _ := m.Forward(&batch.Batch{
+		InputIDs:     mlx.FromValues([]int32{0}, 1, 1),
+		SeqOffsets:   []int32{0},
+		SeqQueryLens: []int32{1},
+	}, caches)
+	logits := m.Unembed(hidden)
+	width := logits.Dim(logits.NumDims() - 1)
+	for _, c := range caches {
+		if c != nil {
+			c.Free()
+		}
+	}
+	mlx.Sweep()
+	return width
+}
+
 func configureWiredMemory() {
 	if !mlx.GPUIsAvailable() {
 		return
@@ -257,6 +293,7 @@ func (r *Runner) Run(host, port string, mux http.Handler) error {
 }
 
 func (r *Runner) runRequest(request Request) error {
+	defer request.Grammar.close()
 	if r.mlxThread == nil {
 		return request.Pipeline(request.Ctx, request)
 	}

x/mlxrunner/server.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
@@ -13,6 +14,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/ollama/ollama/api"
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/logutil"
 	"github.com/ollama/ollama/x/internal/mlxthread"
@@ -68,6 +70,7 @@ func Execute(args []string) error {
 	}); err != nil {
 		return err
 	}
+	defer runner.Close()
 
 	readMemory := func() (uint64, error) {
 		return uint64(mlx.ActiveMemory() + mlx.CacheMemory()), nil
@@ -143,7 +146,12 @@ func Execute(args []string) error {
 		}
 
 		if err := runner.Prepare(&request); err != nil {
-			http.Error(w, err.Error(), http.StatusBadRequest)
+			var statusErr api.StatusError
+			if errors.As(err, &statusErr) {
+				http.Error(w, statusErr.ErrorMessage, statusErr.StatusCode)
+			} else {
+				http.Error(w, err.Error(), http.StatusBadRequest)
+			}
 			return
 		}
 
@@ -153,6 +161,8 @@ func Execute(args []string) error {
 
 		select {
 		case <-r.Context().Done():
+			// Never queued, so the runner will not close the grammar.
+			request.Grammar.close()
 			return
 		case runner.Requests <- request:
 		}

x/mlxrunner/speculate.go

@@ -118,7 +118,7 @@ func (s *speculation) open(request Request, layout []any) *speculationSession {
 	// Logprobs are not yet supported, so a logprobs request keeps a speculationSession
 	// only to maintain a draft cache in lockstep (permanently parked).
 	opts := request.SamplerOpts
-	enabled := !opts.Logprobs && opts.TopLogprobs == 0
+	enabled := request.Grammar == nil && !opts.Logprobs && opts.TopLogprobs == 0
 
 	spec := &speculationSession{spec: s, drafter: d, layout: layout, enabled: enabled, prevDrafts: -1, roundDrafts: -1}
 	if enabled {
@@ -128,8 +128,13 @@ func (s *speculation) open(request Request, layout []any) *speculationSession {
 }
 
 // beginRound records the previous round's cost sample (its wall time runs to
-// this round's start) and starts timing the new one.
+// this round's start) and starts timing the new one. A session that cannot
+// draft records nothing: its parked rounds carry grammar or logprobs work
+// that would skew the shared depth-0 cost.
 func (s *speculationSession) beginRound() {
+	if !s.enabled {
+		return
+	}
 	now := time.Now()
 	if !s.lastRoundStart.IsZero() && s.roundDrafts >= 0 {
 		if s.roundDrafts == s.prevDrafts {
@@ -194,15 +199,18 @@ type speculativeDecoder struct {
 	position int
 	current  sampler.Result    // emitted (or the seed), not yet forwarded
 	inner    *pipelinedDecoder // pipelines plain tokens while parked; nil while drafting
+	// grammar reaches sampling through the parked inner decoder; a
+	// constrained session never drafts.
+	grammar *grammar
 }
 
 // decoder returns the decoder for this engine's session. A speculationSession that
-// cannot draft (logprobs) has no depth controller and permanently parks,
-// running the inner pipelined decoder whose reports keep the draft KV level.
-func (s *speculationSession) decoder(seed *mlx.Array, position int) decoder {
+// cannot draft (a grammar, logprobs) has no depth controller and permanently
+// parks, running the inner pipelined decoder whose reports keep the draft KV level.
+func (s *speculationSession) decoder(seed *mlx.Array, position int, grammar *grammar) decoder {
 	current := sampler.Result{Token: seed}
 	mlx.Pin(current.Arrays()...)
-	return &speculativeDecoder{s: s, position: position, current: current}
+	return &speculativeDecoder{s: s, position: position, current: current, grammar: grammar}
 }
 
 func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) {
@@ -276,7 +284,7 @@ func (st *speculativeDecoder) resume() []sampler.Result {
 func (st *speculativeDecoder) park(remaining int) ([]sampler.Result, error) {
 	s := st.s
 	if st.inner == nil {
-		st.inner = s.spec.r.pipelinedDecoder(s, s.spec.targets, st.current.Token.ExpandDims(-1), st.position, s.layout)
+		st.inner = s.spec.r.pipelinedDecoder(s, s.spec.targets, st.current.Token.ExpandDims(-1), st.position, s.layout, st.grammar)
 	}
 	return st.inner.next(remaining)
 }

x/mlxrunner/xgrammar/dynamic.c

@@ -0,0 +1,146 @@
+#include "dynamic.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#ifdef _WIN32
+#include <windows.h>
+
+#ifndef LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
+#define LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR 0x00000100
+#endif
+#ifndef LOAD_LIBRARY_SEARCH_DEFAULT_DIRS
+#define LOAD_LIBRARY_SEARCH_DEFAULT_DIRS 0x00001000
+#endif
+
+static void* open_library(const char* path) {
+    int length = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, path, -1, NULL, 0);
+    if (length == 0) {
+        return NULL;
+    }
+    wchar_t* wide_path = malloc((size_t)length * sizeof(wchar_t));
+    if (wide_path == NULL) {
+        return NULL;
+    }
+    if (MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, path, -1, wide_path, length) == 0) {
+        free(wide_path);
+        return NULL;
+    }
+    HMODULE module = LoadLibraryExW(
+        wide_path,
+        NULL,
+        LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR | LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
+    free(wide_path);
+    return (void*)module;
+}
+
+#define OPEN(path) open_library(path)
+#define CLOSE(handle) FreeLibrary((HMODULE)(handle))
+#define SYMBOL(handle, name) ((void*)GetProcAddress((HMODULE)(handle), name))
+#else
+#include <dlfcn.h>
+#define OPEN(path) dlopen(path, RTLD_NOW | RTLD_LOCAL)
+#define CLOSE(handle) dlclose(handle)
+#define SYMBOL(handle, name) dlsym(handle, name)
+#endif
+
+static const char* (*version_fn)(void);
+static const char* (*last_error_fn)(void);
+static int (*compiler_new_fn)(const char*, size_t, const uint64_t*, size_t, int32_t, const int32_t*, size_t, int32_t, int64_t, ollama_xgrammar_compiler**);
+static void (*compiler_free_fn)(ollama_xgrammar_compiler*);
+static int (*matcher_new_fn)(ollama_xgrammar_compiler*, ollama_xgrammar_kind, const char*, size_t, ollama_xgrammar_matcher**);
+static void (*matcher_free_fn)(ollama_xgrammar_matcher*);
+static int (*matcher_fill_fn)(ollama_xgrammar_matcher*, int32_t*, size_t, int*);
+static int (*matcher_accept_fn)(ollama_xgrammar_matcher*, int32_t, int*);
+static const char* load_error;
+
+static void clear_symbols(void) {
+    version_fn = NULL;
+    last_error_fn = NULL;
+    compiler_new_fn = NULL;
+    compiler_free_fn = NULL;
+    matcher_new_fn = NULL;
+    matcher_free_fn = NULL;
+    matcher_fill_fn = NULL;
+    matcher_accept_fn = NULL;
+}
+
+#define LOAD(handle, field, name) do { \
+    *(void**)(&field) = SYMBOL((handle)->ctx, name); \
+    if ((field) == NULL) { load_error = "xgrammar library is missing symbol " name; goto fail; } \
+} while (0)
+
+int ollama_xgrammar_dynamic_load(ollama_xgrammar_dynamic_handle* handle, const char* path) {
+    clear_symbols();
+    if (handle == NULL || path == NULL) {
+        load_error = "invalid xgrammar library path";
+        return 1;
+    }
+    handle->ctx = OPEN(path);
+    if (handle->ctx == NULL) {
+        load_error = "unable to open xgrammar library";
+        return 1;
+    }
+    LOAD(handle, version_fn, "ollama_xgrammar_version");
+    LOAD(handle, last_error_fn, "ollama_xgrammar_last_error");
+    LOAD(handle, compiler_new_fn, "ollama_xgrammar_compiler_new");
+    LOAD(handle, compiler_free_fn, "ollama_xgrammar_compiler_free");
+    LOAD(handle, matcher_new_fn, "ollama_xgrammar_matcher_new");
+    LOAD(handle, matcher_free_fn, "ollama_xgrammar_matcher_free");
+    LOAD(handle, matcher_fill_fn, "ollama_xgrammar_matcher_fill");
+    LOAD(handle, matcher_accept_fn, "ollama_xgrammar_matcher_accept");
+    load_error = NULL;
+    return 0;
+
+fail:
+    CLOSE(handle->ctx);
+    handle->ctx = NULL;
+    clear_symbols();
+    return 1;
+}
+
+const char* ollama_xgrammar_dynamic_version(void) {
+    return version_fn == NULL ? "" : version_fn();
+}
+
+const char* ollama_xgrammar_dynamic_error(void) {
+    if (last_error_fn != NULL) {
+        const char* message = last_error_fn();
+        if (message != NULL && message[0] != '\0') {
+            return message;
+        }
+    }
+    return load_error == NULL ? "xgrammar library error" : load_error;
+}
+
+static int capture_error(int result, char** error) {
+    if (error != NULL) {
+        *error = NULL;
+    }
+    if (result == 0 || error == NULL) {
+        return result;
+    }
+    const char* message = ollama_xgrammar_dynamic_error();
+    size_t size = strlen(message) + 1;
+    *error = malloc(size);
+    if (*error != NULL) {
+        memcpy(*error, message, size);
+    }
+    return result;
+}
+
+int ollama_xgrammar_dynamic_compiler_new(const char* d, size_t ds, const uint64_t* o, size_t n, int32_t v, const int32_t* s, size_t ns, int32_t mt, int64_t cb, ollama_xgrammar_compiler** m, char** error) {
+    return capture_error(compiler_new_fn(d, ds, o, n, v, s, ns, mt, cb, m), error);
+}
+void ollama_xgrammar_dynamic_compiler_free(ollama_xgrammar_compiler* m) { compiler_free_fn(m); }
+int ollama_xgrammar_dynamic_matcher_new(ollama_xgrammar_compiler* m, ollama_xgrammar_kind k, const char* s, size_t n, ollama_xgrammar_matcher** out, char** error) {
+    return capture_error(matcher_new_fn(m, k, s, n, out), error);
+}
+void ollama_xgrammar_dynamic_matcher_free(ollama_xgrammar_matcher* m) { matcher_free_fn(m); }
+int ollama_xgrammar_dynamic_matcher_fill(ollama_xgrammar_matcher* m, int32_t* b, size_t n, int* a, char** error) {
+    return capture_error(matcher_fill_fn(m, b, n, a), error);
+}
+int ollama_xgrammar_dynamic_matcher_accept(ollama_xgrammar_matcher* m, int32_t t, int* a, char** error) {
+    return capture_error(matcher_accept_fn(m, t, a), error);
+}

x/mlxrunner/xgrammar/dynamic.h

@@ -0,0 +1,25 @@
+#ifndef OLLAMA_XGRAMMAR_DYNAMIC_H
+#define OLLAMA_XGRAMMAR_DYNAMIC_H
+
+#include "native/xgrammar.h"
+
+typedef struct ollama_xgrammar_dynamic_handle {
+    void* ctx;
+} ollama_xgrammar_dynamic_handle;
+
+int ollama_xgrammar_dynamic_load(ollama_xgrammar_dynamic_handle* handle, const char* path);
+const char* ollama_xgrammar_dynamic_version(void);
+const char* ollama_xgrammar_dynamic_error(void);
+
+int ollama_xgrammar_dynamic_compiler_new(
+    const char*, size_t, const uint64_t*, size_t, int32_t,
+    const int32_t*, size_t, int32_t, int64_t, ollama_xgrammar_compiler**, char**);
+void ollama_xgrammar_dynamic_compiler_free(ollama_xgrammar_compiler*);
+int ollama_xgrammar_dynamic_matcher_new(
+    ollama_xgrammar_compiler*, ollama_xgrammar_kind, const char*, size_t,
+    ollama_xgrammar_matcher**, char**);
+void ollama_xgrammar_dynamic_matcher_free(ollama_xgrammar_matcher*);
+int ollama_xgrammar_dynamic_matcher_fill(ollama_xgrammar_matcher*, int32_t*, size_t, int*, char**);
+int ollama_xgrammar_dynamic_matcher_accept(ollama_xgrammar_matcher*, int32_t, int*, char**);
+
+#endif

x/mlxrunner/xgrammar/native/xgrammar.cpp

@@ -0,0 +1,228 @@
+#include "xgrammar.h"
+
+#include <dlpack/dlpack.h>
+#include <xgrammar/xgrammar.h>
+
+#include <algorithm>
+#include <exception>
+#include <memory>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace {
+
+thread_local std::string last_error;
+
+// Errors travel into API responses and logs: strip xgrammar's console
+// decoration — "[HH:MM:SS] file:line: message\n" — and cap the length,
+// since messages can embed schema fragments.
+constexpr size_t kMaxErrorSize = 512;
+
+void sanitize_error(std::string& s) {
+    if (s.size() >= 11 && s[0] == '[' && s[9] == ']' && s[10] == ' ') {
+        for (size_t i = 11; (i = s.find(':', i)) != std::string::npos; ++i) {
+            size_t j = i + 1;
+            while (j < s.size() && s[j] >= '0' && s[j] <= '9') {
+                ++j;
+            }
+            if (j > i + 1 && j + 1 < s.size() && s[j] == ':' && s[j + 1] == ' ') {
+                s.erase(0, j + 2);
+                break;
+            }
+        }
+    }
+    while (!s.empty() && s.back() == '\n') {
+        s.pop_back();
+    }
+    if (s.size() > kMaxErrorSize) {
+        s.resize(kMaxErrorSize);
+        // Do not leave a torn UTF-8 sequence at the cut.
+        while (!s.empty() && (static_cast<unsigned char>(s.back()) & 0xc0) == 0x80) {
+            s.pop_back();
+        }
+        if (!s.empty() && static_cast<unsigned char>(s.back()) >= 0xc0) {
+            s.pop_back();
+        }
+        s += "...";
+    }
+}
+
+template <typename F>
+int protect(F&& fn) {
+    try {
+        last_error.clear();
+        fn();
+        return 0;
+    } catch (const std::exception& e) {
+        last_error = e.what();
+        sanitize_error(last_error);
+    } catch (...) {
+        last_error = "unknown xgrammar library error";
+    }
+    return 1;
+}
+
+}  // namespace
+
+struct ollama_xgrammar_compiler {
+    int32_t vocab_size;
+    xgrammar::TokenizerInfo tokenizer;
+    xgrammar::GrammarCompiler compiler;
+
+    ollama_xgrammar_compiler(std::vector<std::string> vocab, int32_t size, std::vector<int32_t> stops,
+                          int32_t max_threads, int64_t cache_bytes)
+        : vocab_size(size),
+          tokenizer(vocab, xgrammar::VocabType::RAW, size, std::move(stops)),
+          compiler(tokenizer, max_threads, /*cache_enabled=*/cache_bytes > 0,
+                   /*max_memory_bytes=*/cache_bytes > 0 ? cache_bytes : 0) {}
+};
+
+struct ollama_xgrammar_matcher {
+    int32_t vocab_size;
+    xgrammar::GrammarMatcher matcher;
+
+    ollama_xgrammar_matcher(int32_t size, xgrammar::CompiledGrammar grammar)
+        : vocab_size(size), matcher(grammar) {}
+};
+
+#ifndef OLLAMA_XGRAMMAR_VERSION
+#define OLLAMA_XGRAMMAR_VERSION "unknown"
+#endif
+
+extern "C" {
+
+const char* ollama_xgrammar_version(void) {
+    return OLLAMA_XGRAMMAR_VERSION;
+}
+
+const char* ollama_xgrammar_last_error(void) {
+    return last_error.c_str();
+}
+
+int ollama_xgrammar_compiler_new(
+    const char* token_data,
+    size_t token_data_size,
+    const uint64_t* token_offsets,
+    size_t token_count,
+    int32_t vocab_size,
+    const int32_t* stop_token_ids,
+    size_t stop_token_count,
+    int32_t max_threads,
+    int64_t cache_bytes,
+    ollama_xgrammar_compiler** compiler) {
+    return protect([&] {
+        if (compiler == nullptr) {
+            throw std::invalid_argument("compiler output is null");
+        }
+        *compiler = nullptr;
+        if (token_count > 0 && token_offsets == nullptr) {
+            throw std::invalid_argument("token offsets are null");
+        }
+        if (token_data_size > 0 && token_data == nullptr) {
+            throw std::invalid_argument("token data is null");
+        }
+        if (stop_token_count > 0 && stop_token_ids == nullptr) {
+            throw std::invalid_argument("stop token ids are null");
+        }
+
+        std::vector<std::string> vocab;
+        vocab.reserve(token_count);
+        const char* base = token_data == nullptr ? "" : token_data;
+        uint64_t begin = 0;
+        for (size_t i = 0; i < token_count; ++i) {
+            uint64_t end = token_offsets[i];
+            if (end < begin || end > token_data_size) {
+                throw std::invalid_argument("invalid token offsets");
+            }
+            vocab.emplace_back(base + begin, static_cast<size_t>(end - begin));
+            begin = end;
+        }
+        if (begin != token_data_size) {
+            throw std::invalid_argument("token offsets do not consume token data");
+        }
+
+        std::vector<int32_t> stops(stop_token_ids, stop_token_ids + stop_token_count);
+        *compiler = new ollama_xgrammar_compiler(
+            std::move(vocab), vocab_size, std::move(stops), max_threads, cache_bytes);
+    });
+}
+
+int ollama_xgrammar_matcher_new(
+    ollama_xgrammar_compiler* compiler,
+    ollama_xgrammar_kind kind,
+    const char* source,
+    size_t source_size,
+    ollama_xgrammar_matcher** matcher) {
+    return protect([&] {
+        if (compiler == nullptr || matcher == nullptr) {
+            throw std::invalid_argument("compiler or matcher output is null");
+        }
+        *matcher = nullptr;
+        if (source_size > 0 && source == nullptr) {
+            throw std::invalid_argument("grammar source is null");
+        }
+        const char* source_data = source == nullptr ? "" : source;
+
+        xgrammar::CompiledGrammar compiled = [&]() -> xgrammar::CompiledGrammar {
+            switch (kind) {
+            case OLLAMA_XGRAMMAR_JSON_SCHEMA:
+                return compiler->compiler.CompileJSONSchema(std::string(source_data, source_size));
+            default:
+                throw std::invalid_argument("unknown grammar kind");
+            }
+        }();
+        *matcher = new ollama_xgrammar_matcher(compiler->vocab_size, std::move(compiled));
+    });
+}
+
+int ollama_xgrammar_matcher_fill(
+    ollama_xgrammar_matcher* matcher,
+    int32_t* bitmask,
+    size_t bitmask_words,
+    int* needs_apply) {
+    return protect([&] {
+        if (matcher == nullptr || bitmask == nullptr || needs_apply == nullptr) {
+            throw std::invalid_argument("matcher, bitmask, or result is null");
+        }
+        size_t expected = static_cast<size_t>(xgrammar::GetBitmaskSize(matcher->vocab_size));
+        if (bitmask_words != expected) {
+            throw std::invalid_argument("incorrect bitmask size");
+        }
+
+        int64_t shape[] = {static_cast<int64_t>(bitmask_words)};
+        DLTensor tensor{};
+        tensor.data = bitmask;
+        tensor.device = DLDevice{kDLCPU, 0};
+        tensor.ndim = 1;
+        tensor.dtype = xgrammar::GetBitmaskDLType();
+        tensor.shape = shape;
+        tensor.strides = nullptr;
+        tensor.byte_offset = 0;
+        *needs_apply = matcher->matcher.FillNextTokenBitmask(&tensor) ? 1 : 0;
+    });
+}
+
+int ollama_xgrammar_matcher_accept(
+    ollama_xgrammar_matcher* matcher,
+    int32_t token_id,
+    int* accepted) {
+    return protect([&] {
+        if (matcher == nullptr || accepted == nullptr) {
+            throw std::invalid_argument("matcher or result is null");
+        }
+        *accepted = matcher->matcher.AcceptToken(token_id) ? 1 : 0;
+    });
+}
+
+void ollama_xgrammar_matcher_free(ollama_xgrammar_matcher* matcher) {
+    delete matcher;
+}
+
+void ollama_xgrammar_compiler_free(ollama_xgrammar_compiler* compiler) {
+    delete compiler;
+}
+
+}  // extern "C"

x/mlxrunner/xgrammar/native/xgrammar.h

@@ -0,0 +1,68 @@
+#ifndef OLLAMA_XGRAMMAR_H
+#define OLLAMA_XGRAMMAR_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#if defined(_WIN32)
+#if defined(OLLAMA_XGRAMMAR_BUILD)
+#define OLLAMA_XGRAMMAR_API __declspec(dllexport)
+#else
+#define OLLAMA_XGRAMMAR_API __declspec(dllimport)
+#endif
+#else
+#define OLLAMA_XGRAMMAR_API __attribute__((visibility("default")))
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct ollama_xgrammar_compiler ollama_xgrammar_compiler;
+typedef struct ollama_xgrammar_matcher ollama_xgrammar_matcher;
+
+typedef enum ollama_xgrammar_kind {
+    OLLAMA_XGRAMMAR_JSON_SCHEMA = 0,
+} ollama_xgrammar_kind;
+
+// The pinned xgrammar release the library was built from.
+OLLAMA_XGRAMMAR_API const char* ollama_xgrammar_version(void);
+
+OLLAMA_XGRAMMAR_API const char* ollama_xgrammar_last_error(void);
+
+OLLAMA_XGRAMMAR_API int ollama_xgrammar_compiler_new(
+    const char* token_data,
+    size_t token_data_size,
+    const uint64_t* token_offsets,
+    size_t token_count,
+    int32_t vocab_size,
+    const int32_t* stop_token_ids,
+    size_t stop_token_count,
+    int32_t max_threads,
+    int64_t cache_bytes,
+    ollama_xgrammar_compiler** compiler);
+OLLAMA_XGRAMMAR_API void ollama_xgrammar_compiler_free(ollama_xgrammar_compiler* compiler);
+
+OLLAMA_XGRAMMAR_API int ollama_xgrammar_matcher_new(
+    ollama_xgrammar_compiler* compiler,
+    ollama_xgrammar_kind kind,
+    const char* source,
+    size_t source_size,
+    ollama_xgrammar_matcher** matcher);
+OLLAMA_XGRAMMAR_API void ollama_xgrammar_matcher_free(ollama_xgrammar_matcher* matcher);
+
+OLLAMA_XGRAMMAR_API int ollama_xgrammar_matcher_fill(
+    ollama_xgrammar_matcher* matcher,
+    int32_t* bitmask,
+    size_t bitmask_words,
+    int* needs_apply);
+OLLAMA_XGRAMMAR_API int ollama_xgrammar_matcher_accept(
+    ollama_xgrammar_matcher* matcher,
+    int32_t token_id,
+    int* accepted);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

x/mlxrunner/xgrammar/xgrammar.go

@@ -0,0 +1,262 @@
+package xgrammar
+
+// #cgo linux LDFLAGS: -ldl
+// #include "dynamic.h"
+// #include <stdlib.h>
+import "C"
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"slices"
+	"sync"
+	"unsafe"
+)
+
+type Kind int
+
+const (
+	JSONSchema Kind = C.OLLAMA_XGRAMMAR_JSON_SCHEMA
+)
+
+// The native library is loaded once per process and never unloaded.
+var (
+	loadOnce sync.Once
+	loadDir  string
+	loadPath string
+	loadErr  error
+)
+
+func loadLibrary(dir string) (string, error) {
+	loadOnce.Do(func() {
+		loadDir = dir
+		loadPath, loadErr = openLibrary(dir)
+	})
+	if loadErr != nil {
+		return "", loadErr
+	}
+	if dir != loadDir {
+		return "", fmt.Errorf("xgrammar library already loaded from %s", loadDir)
+	}
+	return loadPath, nil
+}
+
+func openLibrary(dir string) (string, error) {
+	var name string
+	switch runtime.GOOS {
+	case "darwin":
+		name = "libollama_xgrammar.dylib"
+	case "linux":
+		name = "libollama_xgrammar.so"
+	case "windows":
+		name = "ollama_xgrammar.dll"
+	default:
+		return "", fmt.Errorf("xgrammar library is not supported on %s", runtime.GOOS)
+	}
+	path, err := filepath.Abs(filepath.Join(dir, name))
+	if err != nil {
+		return "", fmt.Errorf("resolve xgrammar library path: %w", err)
+	}
+	if _, err := os.Stat(path); err != nil {
+		return "", fmt.Errorf("xgrammar library not found at %s: %w", path, err)
+	}
+	cPath := C.CString(path)
+	defer C.free(unsafe.Pointer(cPath))
+	var handle C.ollama_xgrammar_dynamic_handle
+	if C.ollama_xgrammar_dynamic_load(&handle, cPath) != 0 {
+		return "", fmt.Errorf("load xgrammar library %s: %s", path, C.GoString(C.ollama_xgrammar_dynamic_error()))
+	}
+	return path, nil
+}
+
+// Compiler compiles grammars for one vocabulary.
+type Compiler struct {
+	ctx       *C.ollama_xgrammar_compiler
+	path      string
+	vocabSize int
+	stops     []int32
+}
+
+// New loads the native library from dir — once per process, never unloaded —
+// and binds a grammar compiler to the vocabulary. cacheBytes bounds the
+// engine's compiled-grammar cache; <= 0 disables it.
+func New(dir string, pieces []string, vocabSize int, stopIDs []int32, threads int, cacheBytes int64) (*Compiler, error) {
+	path, err := loadLibrary(dir)
+	if err != nil {
+		return nil, err
+	}
+	if vocabSize <= 0 || len(pieces) > vocabSize {
+		return nil, fmt.Errorf("invalid vocabulary size %d for %d token pieces", vocabSize, len(pieces))
+	}
+	if len(stopIDs) == 0 {
+		return nil, errors.New("tokenizer has no stop tokens")
+	}
+	for _, id := range stopIDs {
+		if id < 0 || int(id) >= vocabSize {
+			return nil, fmt.Errorf("stop token %d is outside the vocabulary", id)
+		}
+	}
+
+	var data []byte
+	offsets := make([]C.uint64_t, len(pieces))
+	for i, piece := range pieces {
+		data = append(data, piece...)
+		offsets[i] = C.uint64_t(len(data))
+	}
+	cStops := make([]C.int32_t, len(stopIDs))
+	for i, id := range stopIDs {
+		cStops[i] = C.int32_t(id)
+	}
+
+	var dataPtr *C.char
+	if len(data) > 0 {
+		dataPtr = (*C.char)(unsafe.Pointer(&data[0]))
+	}
+	var offsetPtr *C.uint64_t
+	if len(offsets) > 0 {
+		offsetPtr = &offsets[0]
+	}
+	var ctx *C.ollama_xgrammar_compiler
+	var cError *C.char
+	if C.ollama_xgrammar_dynamic_compiler_new(
+		dataPtr, C.size_t(len(data)), offsetPtr, C.size_t(len(offsets)), C.int32_t(vocabSize),
+		&cStops[0], C.size_t(len(cStops)), C.int32_t(threads), C.int64_t(cacheBytes), &ctx, &cError,
+	) != 0 {
+		return nil, nativeError("create grammar compiler", cError)
+	}
+	return &Compiler{ctx: ctx, path: path, vocabSize: vocabSize, stops: slices.Clone(stopIDs)}, nil
+}
+
+// Path returns the loaded native library's location.
+func (c *Compiler) Path() string {
+	return c.path
+}
+
+// Version returns the pinned xgrammar release the library was built from.
+func (c *Compiler) Version() string {
+	return C.GoString(C.ollama_xgrammar_dynamic_version())
+}
+
+func (c *Compiler) Compile(kind Kind, source string) (*Matcher, error) {
+	if c == nil {
+		return nil, errors.New("grammar compiler is unavailable")
+	}
+	if c.ctx == nil {
+		return nil, errors.New("grammar compiler is closed")
+	}
+
+	var sourcePtr *C.char
+	if len(source) > 0 {
+		sourcePtr = (*C.char)(unsafe.Pointer(unsafe.StringData(source)))
+	}
+	var ctx *C.ollama_xgrammar_matcher
+	var cError *C.char
+	if C.ollama_xgrammar_dynamic_matcher_new(
+		c.ctx, C.ollama_xgrammar_kind(kind), sourcePtr, C.size_t(len(source)), &ctx, &cError,
+	) != 0 {
+		return nil, nativeError("compile grammar", cError)
+	}
+	return &Matcher{ctx: ctx, vocabSize: c.vocabSize, stops: c.stops}, nil
+}
+
+func (c *Compiler) Close() {
+	if c == nil {
+		return
+	}
+	if c.ctx != nil {
+		C.ollama_xgrammar_dynamic_compiler_free(c.ctx)
+		c.ctx = nil
+	}
+}
+
+type Matcher struct {
+	ctx       *C.ollama_xgrammar_matcher
+	vocabSize int
+	stops     []int32
+	// terminated: a stop token was accepted, ending the grammar; the matcher
+	// no longer constrains sampling.
+	terminated bool
+}
+
+// Terminated reports whether an accepted stop token ended the grammar; a
+// terminated matcher no longer constrains sampling.
+func (m *Matcher) Terminated() bool {
+	if m == nil {
+		return true
+	}
+	return m.terminated
+}
+
+// Fill writes the packed allowed-token bitmask for the next position into
+// row — bit id%32 of word id/32 is set when token id is allowed — and
+// reports whether it constrains sampling. A false return means every token
+// is allowed or the grammar has terminated; row contents are meaningful
+// only on true. row must hold ceil(vocabSize/32) words and may be one row
+// of a larger batch buffer.
+func (m *Matcher) Fill(row []int32) (bool, error) {
+	if m == nil {
+		return false, errors.New("grammar matcher is closed")
+	}
+	if m.ctx == nil {
+		return false, errors.New("grammar matcher is closed")
+	}
+	if want := (m.vocabSize + 31) / 32; len(row) != want {
+		return false, fmt.Errorf("mask row holds %d words; the vocabulary needs %d", len(row), want)
+	}
+	if m.terminated {
+		return false, nil
+	}
+
+	var needsApply C.int
+	var cError *C.char
+	if C.ollama_xgrammar_dynamic_matcher_fill(m.ctx, (*C.int32_t)(unsafe.Pointer(&row[0])), C.size_t(len(row)), &needsApply, &cError) != 0 {
+		return false, nativeError("build token mask", cError)
+	}
+	return needsApply != 0, nil
+}
+
+func (m *Matcher) Accept(tokenID int32) error {
+	if m == nil {
+		return errors.New("grammar matcher is closed")
+	}
+	if m.ctx == nil {
+		return errors.New("grammar matcher is closed")
+	}
+	var accepted C.int
+	var cError *C.char
+	if C.ollama_xgrammar_dynamic_matcher_accept(m.ctx, C.int32_t(tokenID), &accepted, &cError) != 0 {
+		return nativeError("accept token", cError)
+	}
+	if accepted == 0 {
+		return fmt.Errorf("grammar rejected sampled token %d", tokenID)
+	}
+	if slices.Contains(m.stops, tokenID) {
+		m.terminated = true
+	}
+	return nil
+}
+
+func (m *Matcher) Close() {
+	if m == nil {
+		return
+	}
+	if m.ctx != nil {
+		C.ollama_xgrammar_dynamic_matcher_free(m.ctx)
+		m.ctx = nil
+	}
+}
+
+func nativeError(op string, cMessage *C.char) error {
+	message := "unknown error"
+	if cMessage != nil {
+		defer C.free(unsafe.Pointer(cMessage))
+		message = C.GoString(cMessage)
+	}
+	if message == "" {
+		message = "unknown error"
+	}
+	return fmt.Errorf("%s: %s", op, message)
+}

x/mlxrunner/xgrammar/xgrammar_test.go

@@ -0,0 +1,236 @@
+package xgrammar_test
+
+import (
+	"errors"
+	"io/fs"
+	"path/filepath"
+	"slices"
+	"strings"
+	"testing"
+
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+	"github.com/ollama/ollama/x/mlxrunner/xgrammar"
+)
+
+const (
+	testEOS       int32 = 32                // Stop token in the second packed-mask word.
+	testVocabSize       = 40                // Spans two packed-mask words, with the second only partly used.
+	testPaddedID  int32 = testVocabSize - 1 // Empty padding at the vocabulary boundary must stay disallowed.
+)
+
+func testVocabulary() []string {
+	pieces := []string{
+		"{", "}", `"`, ":", ",", "a", "n", "s", "w", "e", "r", "o", "k", " ",
+		"1", "2", "[", "]", "true", "false", "ok", "answer", "\x00", "\xff",
+	}
+	for len(pieces) < testVocabSize {
+		pieces = append(pieces, "")
+	}
+	pieces[testEOS] = "<eos>"
+	return pieces
+}
+
+func testLibraryDir(t testing.TB) string {
+	t.Helper()
+	path, err := mlx.LoadedLibraryPath()
+	if err != nil {
+		t.Skipf("native MLX payload is not built: %v", err)
+	}
+	return filepath.Dir(path)
+}
+
+func testGrammarCompiler(t testing.TB) *xgrammar.Compiler {
+	t.Helper()
+	compiler, err := xgrammar.New(testLibraryDir(t), testVocabulary(), testVocabSize, []int32{testEOS}, 8, 128<<20)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			t.Skipf("native xgrammar payload is not built: %v", err)
+		}
+		t.Fatal(err)
+	}
+	t.Cleanup(compiler.Close)
+	return compiler
+}
+
+func FuzzJSONSchemaMatcher(f *testing.F) {
+	compiler := testGrammarCompiler(f)
+	for _, seed := range []struct {
+		schema string
+		tokens []byte
+	}{
+		{schema: `{}`, tokens: []byte{0, 1, byte(testEOS)}},                    // Accept {}, then EOS.
+		{schema: `{"type":"string"}`, tokens: []byte{2, 20, 2, byte(testEOS)}}, // Accept "ok", then EOS.
+		{schema: `{"$defs":{"node":{"type":"array","items":{"$ref":"#/$defs/node"}}},"$ref":"#/$defs/node"}`},
+		{schema: `{"enum":["a","b","c"]}`},
+		{schema: `{"type":`},
+		{schema: string([]byte{'{', '"', 'x', '"', ':', '"', 0xff, '"', '}'})},
+	} {
+		f.Add(seed.schema, seed.tokens)
+	}
+
+	f.Fuzz(func(t *testing.T, schema string, tokens []byte) {
+		if len(schema) > 1<<20 || len(tokens) > 256 {
+			return
+		}
+		matcher, err := compiler.Compile(xgrammar.JSONSchema, schema)
+		if err != nil {
+			return
+		}
+		defer matcher.Close()
+
+		row := make([]int32, (testVocabSize+31)/32)
+		for _, token := range tokens {
+			id := int32(token) % testVocabSize
+			if _, err := matcher.Fill(row); err != nil {
+				return
+			}
+			if err := matcher.Accept(id); err != nil {
+				return
+			}
+		}
+	})
+}
+
+func allowed(mask []int32, id int32) bool {
+	return uint32(mask[id/32])&(uint32(1)<<uint(id%32)) != 0
+}
+
+// fillMask fills a fresh mask row and reports whether it constrains.
+func fillMask(t *testing.T, matcher *xgrammar.Matcher) ([]int32, bool) {
+	t.Helper()
+	row := make([]int32, (testVocabSize+31)/32)
+	constrained, err := matcher.Fill(row)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return row, constrained
+}
+
+func acceptPieces(t *testing.T, matcher *xgrammar.Matcher, pieces ...string) {
+	t.Helper()
+	vocab := testVocabulary()
+	for _, piece := range pieces {
+		id := int32(slices.Index(vocab, piece))
+		if id < 0 {
+			t.Fatalf("test vocabulary does not contain %q", piece)
+		}
+		mask, constrained := fillMask(t, matcher)
+		if constrained && !allowed(mask, id) {
+			t.Fatalf("token %d is not allowed by mask %032b", id, uint32(mask[id/32]))
+		}
+		if err := matcher.Accept(id); err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestJSONSchemaMatcher(t *testing.T) {
+	compiler := testGrammarCompiler(t)
+	schema := `{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`
+	matcher, err := compiler.Compile(xgrammar.JSONSchema, schema)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer matcher.Close()
+
+	mask, constrained := fillMask(t, matcher)
+	if !constrained {
+		t.Fatal("the schema start does not constrain sampling")
+	}
+	if allowed(mask, testEOS) {
+		t.Fatal("EOS is allowed before the schema is complete")
+	}
+	if allowed(mask, testPaddedID) {
+		t.Fatal("a padded vocabulary ID is allowed")
+	}
+
+	acceptPieces(t, matcher, "{", `"`, "answer", `"`, ":", `"`, "ok", `"`, "}")
+	mask, constrained = fillMask(t, matcher)
+	if !constrained || !allowed(mask, testEOS) {
+		t.Fatal("EOS is not allowed after the schema is complete")
+	}
+	acceptPieces(t, matcher, "<eos>")
+}
+
+// A terminated matcher no longer constrains: Fill reports no constraint, so
+// a decoder can sample past the grammar's end.
+func TestFillAfterTermination(t *testing.T) {
+	compiler := testGrammarCompiler(t)
+	matcher, err := compiler.Compile(xgrammar.JSONSchema, "{}")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer matcher.Close()
+	acceptPieces(t, matcher, "{", "}", "<eos>")
+
+	if _, constrained := fillMask(t, matcher); constrained {
+		t.Fatal("terminated matcher still constrains sampling")
+	}
+}
+
+func TestInvalidGrammarReturnsError(t *testing.T) {
+	compiler := testGrammarCompiler(t)
+	for _, tt := range []struct {
+		name   string
+		source string
+	}{
+		{name: "empty", source: ""},
+		{name: "malformed", source: `{"type":`},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			matcher, err := compiler.Compile(xgrammar.JSONSchema, tt.source)
+			if matcher != nil {
+				matcher.Close()
+			}
+			if err == nil {
+				t.Fatal("Compile unexpectedly succeeded")
+			}
+			if !strings.Contains(err.Error(), "compile grammar") {
+				t.Fatalf("Compile error = %v", err)
+			}
+		})
+	}
+}
+
+func TestCompilerValidationAndClosedState(t *testing.T) {
+	loaded := testGrammarCompiler(t)
+	for _, tt := range []struct {
+		name      string
+		pieces    []string
+		vocabSize int
+		stops     []int32
+	}{
+		{name: "zero vocabulary", vocabSize: 0, stops: []int32{0}},
+		{name: "too many pieces", pieces: []string{"a", "b"}, vocabSize: 1, stops: []int32{0}},
+		{name: "no stop tokens", pieces: []string{"a"}, vocabSize: 1},
+		{name: "stop outside vocabulary", pieces: []string{"a"}, vocabSize: 1, stops: []int32{1}},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			compiler, err := xgrammar.New(testLibraryDir(t), tt.pieces, tt.vocabSize, tt.stops, 8, 128<<20)
+			if compiler != nil {
+				compiler.Close()
+			}
+			if err == nil {
+				t.Fatal("New unexpectedly succeeded")
+			}
+		})
+	}
+
+	matcher, err := loaded.Compile(xgrammar.JSONSchema, "{}")
+	if err != nil {
+		t.Fatal(err)
+	}
+	matcher.Close()
+	matcher.Close()
+	if _, err := matcher.Fill(make([]int32, (testVocabSize+31)/32)); err == nil {
+		t.Fatal("Fill on closed matcher unexpectedly succeeded")
+	}
+	if err := matcher.Accept(0); err == nil {
+		t.Fatal("Accept on closed matcher unexpectedly succeeded")
+	}
+	loaded.Close()
+	loaded.Close()
+	if matcher, err := loaded.Compile(xgrammar.JSONSchema, "{}"); err == nil || matcher != nil {
+		t.Fatalf("Compile on closed compiler = %v, %v; want error", matcher, err)
+	}
+}