Skip to content

Repository files navigation

MLX Structured

MLX Structured is a Swift library for structured output generation using constrained decoding. It's built on top of the XGrammar library, which provides efficient, flexible, and portable structured generation. You can learn more about the XGrammar algorithm in their technical report.

Installation

To use MLX Structured in your project, add the following to your Package.swift file:

dependencies: [
    .package(url: "https://github.com/petrukha-ivan/mlx-swift-structured", from: "0.2.0")
]

Don't forget to add the library as a dependency for your targets:

dependencies: [
    .product(name: "MLXStructured", package: "mlx-swift-structured")
]               

Usage

Grammar

Start by defining a Grammar. You can use JSON Schema to describe the desired output:

let schema = JSONSchema.object(
    description: "Person info",
    properties: [
        "name": .string(),
        "age": .integer()
    ], required: [
        "name",
        "age"
    ]
)

let grammar = try Grammar.schema(schema)

Starting with macOS 26 and iOS 26, you can use a @Generable type as a grammar source:

@Generable
struct PersonInfo: Codable {
    
    @Guide(description: "Person name")
    let name: String
    
    @Guide(description: "Person age")
    let age: Int
}

let grammar = try Grammar.generable(PersonInfo.self)

You can also use a regex:

let regex = #"^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$"# // Simple email regex
let grammar = Grammar.regex(regex)

Or define your own grammar rules with EBNF syntax:

let ebnf = #"root ::= ("YES" | "NO")"# // Answer only "YES" or "NO"
let grammar = Grammar.ebnf(ebnf)

Complex Grammar

You can define rich, composable grammar rules via a grammar builder. This enables you to describe structured output formats precisely:

let grammar = try Grammar {
    SequenceFormat {
        ConstTextFormat(text: "Hello!")
        OrFormat {
            JSONSchemaFormat(...)
            RegexFormat(...)
        }
    }
}

This can be used in different ways. Here is an example of a constrained Qwen3 tool-calling format:

let grammar = try Grammar {
    SequenceFormat {
        if forceThinking {
            TagFormat(begin: "<think>", end: "</think>") {
                AnyTextFormat()
            }
        }
        TriggeredTagsFormat(triggers: ["<tool_call>"], options: [.atLeastOne, .stopAfterFirst]) {
            for tool in tools {
                TagFormat(begin: "<tool_call>\n{\"name\": \"\(tool.name)\", \"arguments\": ", end: "}\n</tool_call>") {
                    JSONSchemaFormat(schema: tool.parameters)
                }
            }
        }
    }
}

Generation

To use a defined grammar during text generation, use the convenient generate method. These overloads are fully compatible with MLXLM generation APIs, with the grammar passed as an additional argument:

let stream = try await generate(
    input: input, 
    context: context, 
    grammar: grammar
)

for await generation in stream {
    switch generation {
    case .chunk(let chunk):
        print(chunk, terminator: "")
    case .toolCall(let toolCall):
        // Handle tool call
    case .info(let info):
        // Handle completion info
    }
}

You can also decode constrained JSON output into a Decodable type:

let model = try await generate(
    input: input, 
    context: context, 
    schema: schema,
    generating: PersonInfo.self
)

With a Generable type, you can generate a validated value directly:

let model = try await generate(
    input: input, 
    context: context,
    generating: PersonInfo.self
)

You can also stream partial Generable updates, which return PartiallyGenerated content for your type:

let stream = try await generate(
    input: input, 
    context: context, 
    partially: PersonInfo.self
)

for try await content in stream {
    print("Partially generated:", content)
}

To enable jump forwarding, pass the .jumpForwarding option to any of generate function:

let stream = try await generate(
    input: input,
    options: .jumpForwarding,
    context: context,
    grammar: grammar
)

You can find more usage examples in the MLXStructuredCLI target and in the unit tests.

Experiments

Performance

Constrained decoding has effectively zero-overhead, so you can produce valid output without sacrificing generation speed. Moreover, in some cases, jump forwarding can make generation even faster by skipping model calls when the output is already determined by the current grammar state. For example, the model can skip predictable JSON syntax such as keys, brackets, and whitespace while generating tokens only for field values.

Model Plain (tokens/s) Constrained (tokens/s) Jump forwarding (tokens/s)
Qwen3-0.6B-4bit 394.1 394.3 604.4
Qwen3-4B-4bit 98.9 98.3 173.1
Qwen3-8B-4bit 58.1 58.1 100.9
Qwen3-14B-4bit 32.8 32.4 51.0
Qwen3-32B-4bit 14.4 14.5 23.6

These benchmark results were collected on an Apple M3 Max using the MLXStructuredCLI benchmark command and the movie record extraction example from the next section. The speedup from jump forwarding depends on the grammar and the ratio of static to dynamic output. In the benchmark example, nearly half of the output is predefined syntax, so jump forwarding can approach almost a 2x speedup.

Accuracy

For example, given a task to extract a movie record from text and output it in JSON format, the prompt is:

Instruction: Extract movie record from the text, output in JSON format according to schema: \(schema)
Text: The Dark Knight (2008) is a superhero crime film directed by Christopher Nolan. Starring Christian Bale, Heath Ledger, and Michael Caine.

And the grammar definition looks like this:

let grammar = try Grammar.schema(.object(
    description: "Movie record",
    properties: [
        "title": .string(),
        "year": .integer(minimum: 1900, maximum: 2026),
        "genres": .array(items: .string(), maxItems: 3),
        "director": .string(),
        "actors": .array(items: .string(), maxItems: 5)
    ], required: [
        "title",
        "year",
        "genres",
        "director",
        "actors"
    ]
))

For large proprietary models like ChatGPT, this is not a problem. With the right prompt, they can successfully generate valid JSON even without constrained decoding. However, with smaller models like Gemma3 270M (especially when quantized to 4-bit), the output almost always contains invalid JSON, even if the schema is provided in the prompt.

[
  "title": "The Dark Knight",
  "actors": [
    "Christian Bale",
    "Heath Ledger",
    "Michael Caine"
  ],
  "genre": "crime",
  "director": "Christopher Nolan",
  "actors": [
    "Christian Bale",
    "Heath Ledger",
    "Michael Caine"
  ],
  "description": "The Dark Knight is a superhero crime film directed by Christopher Nolan. Starring Christian Bale, Heath Ledger, Michael Caine."
]

This output has several issues:

  • Root starts with [ instead of {
  • Incorrect key and type for genres field
  • Missing required year field
  • Duplicated actors field
  • Extra description field

Here is the output using constrained decoding:

{
  "title": "The Dark Knight",
  "year": 2008,
  "genres": [
    "superhero",
    "crime"
  ],
  "director": "Christopher Nolan",
  "actors": [
    "Christian Bale",
    "Heath Ledger",
    "Michael Caine"
  ]
}

The output is fully valid JSON that exactly matches the provided schema. This shows that, with the right approach, even small models like Gemma3 270M 4-bit (which is just 150 MB) can produce correct structured output.

Troubleshooting

This library is still in an early stage of development. While it is already functional, it may have unexpected issues or even crash your program. If you encounter a problem, please create an issue or open a pull request. Contributions are welcome!

About

Structured output generation in Swift

Topics

Resources

Stars

76 stars

Watchers

5 watching

Forks

Releases

Packages

Contributors

Languages