DOCSEARCHbyAlgolia
Submit Docs

MLX Swift LM

repository·main·Indexed 20 days ago

https://github.com/ml-explore/mlx-swift-lm

A Swift package for building applications with Large Language Models (LLMs) and Vision Language Models (VLMs) using the MLX framework. It provides model loading, fine-tuning capabilities, and Hugging Face integrations. The library includes MLXEmbedders for generating embeddings, MLXGuidedGeneration for JSON Schema constrained output, and MLXFoundationModels, which acts as an adapter for Apple's FoundationModels.LanguageModel protocol to enable tool calling, reasoning, and streaming.

Tokens
56.1K
Snippets
134
Records
186
Agent score
72%

What's inside mlx-swift-lm

  3. Overview of KV Cache Types

    main

    The KV (Key-Value) cache stores attention key and value tensors from previous tokens to enable efficient autoregressive generation. Choosing the right cache type involves trading off memory usage, context length, and performance.

    TypeUse CaseMemoryMax Context
    KVCacheSimpleDefault, unboundedGrows with contextUnlimited
    RotatingKVCacheLong contextsFixedmaxKVSize
    QuantizedKVCacheMemory-constrained4-8x lessUnlimited
    ChunkedKVCacheLarge prompt processingControlledChunked
    MambaCacheMamba/SSM modelsFixed stateN/A

  4. Handle optional modules and parameters in Swift

    main

    Some models contain modules or parameters that are only instantiated based on specific configuration settings (e.g., tie_word_embeddings).

    In Swift, you must use Optional types and the @ModuleInfo property wrapper to handle these. If a module is defined but not instantiated, it must be marked as optional so that the parameter loader does not fail when looking for its keys.

    Important: If the module is not created, ensure your callAsFunction logic handles the nil case (e.g., by falling back to a different operation like asLinear) to avoid runtime crashes.

    @ModuleInfo(key: "lm_head") var lmHead: Linear?

public init(_ args: Qwen2Configuration) {
    if !args.tieWordEmbeddings {
        _lmHead.wrappedValue = Linear(args.hiddenSize, args.vocabularySize, bias: false)
    }
}

public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray {
    var out = ...
    if let lmHead {
        out = lmHead(out)
    } else {
        out = model.embedTokens.asLinear(out)
    }
    return out
}

  5. Identify the authoritative source for model compatibility

    main

    To verify if a specific model architecture is supported, do not rely on static lists. Instead, consult the model factories and registries, which are the authoritative sources of truth:

    • Language Models: Check LLMModelFactory and LLMRegistry.
    • Vision-Language Models (VLMs): Check VLMModelFactory, VLMRegistry, and processor registries.
    • Embedding Models: Check MLXEmbedders model factories and registries.

    Note: ModelConfiguration preconfigured values are convenience entries for known-good models, but they do not represent the entire compatibility surface. You can load compatible architectures with matching weights using an explicit configuration.

  6. How ChatSession works for multi-turn conversations

    main

    The ChatSession is the recommended high-level API for chat interfaces. It maintains the conversation state (history) and manages the KV cache for you. You can initialize it with custom instructions or existing message history. To reset the conversation, call await session.clear().

    let session = ChatSession(
    modelContainer,
    instructions: "You are a helpful assistant",
    generateParameters: GenerateParameters(maxTokens: 500, temperature: 0.7)
)

let r1 = try await session.respond(to: "What is 2+2?")
let r2 = try await session.respond(to: "And if you multiply that by 3?")

await session.clear()

  7. How to estimate KV cache and attention workspace

    main

    Inference tickets should account for both persistent KV cache and transient prefill workspace.

    Dense full-attention KV cache

    For standard layers, calculate bytes as:

    1. elements per token per layer = 2 * kvHeads * headDim
    2. layer elements = tokens * elements per token per layer
    3. layer bytes = layer elements * bytesPerElement
    4. total KV bytes = layer bytes * numAttentionLayers

    Note: bytesPerElement is 2 for FP16/BF16, 1 for INT8, and 0.5 for INT4.

    Hybrid / MoE models with SSM

    For models like Qwen3-Next that alternate full-attention layers with linear/SSM layers, sum the full-attention KV cache (using the math above) with the SSM cache sizes for the linear layers.

    Prefill attention workspace (transient)

    Prefill allocates large temporary buffers. To estimate the peak transient workspace, sum the bytes for the following tensors (multiplied by bytesPerElement):

    • Q = B * H * L * D
    • K = B * Hkv * L * D
    • V = B * Hkv * L * D
    • Scores = B * H * L * L
    • Output = B * H * L * D
    • (Optional) Gating tensor: B * L * (H * D)

    Where: B=Batch, H=Heads, L=Prefill Chunk Size, D=Head Dim, Hkv=KV Heads.

    Practical Implementation Pattern

    • Single Ticket Pattern: Most callers create one ticket for the entire generate() call. Budget this ticket for the peak usage: weights + KV cache + prefill workspace.
    • Separate Reservation Pattern: If you have already created a separate reservation ticket for weights, the inference ticket should only cover KV cache + prefill workspace.

  8. Coordinate concurrent inference with Wired Memory

    main

    When running multiple inference tasks concurrently, use WiredMemoryPolicies and tickets to manage memory pressure. You can create a policy (e.g., WiredSumPolicy) and request a ticket with an estimated size. Pass this ticket to modelContainer.generate(..., wiredMemoryTicket:) to ensure the system coordinates memory usage according to the policy.

    let policy = WiredSumPolicy()
let ticket = policy.ticket(size: estimatedBytes, kind: .active)

let userInput = UserInput(prompt: "Summarize this text")
let lmInput = try await modelContainer.prepare(input: userInput)

let stream = try await modelContainer.generate(
    input: lmInput,
    parameters: GenerateParameters(),
    wiredMemoryTicket: ticket
)

  9. Use Active vs Reservation tickets

    main

    When creating a ticket via a policy, specify the kind to distinguish between different types of memory usage:

    • .active: Use this for memory that contributes to the limit while inference is actively running (e.g., KV cache and transient workspace).
    • .reservation: Use this to track long-lived budgets, such as model weights, without keeping the memory limit elevated when no active inference is occurring.
    let reservation = policy.ticket(size: weightBytes, kind: .reservation)
let inference = policy.ticket(size: kvAndWorkspaceBytes, kind: .active)

  10. Implement semantic search and RAG workflows

    main

    The Embedders library enables several common NLP workflows:

    Semantic Search

    Encode a query and a set of documents, then compute similarity using a matrix multiplication (matmul) and sort the results.

    RAG (Retrieval-Augmented Generation)

    1. Index: Generate and store embeddings for documents in a vector database.
    2. Retrieve: Encode the user query and search the vector DB for the top-K relevant documents.
    3. Generate: Construct a prompt containing the retrieved context and pass it to an LLM.

    Similarity Scoring

    Since embeddings are typically L2 normalized during pooling, cosine similarity can be computed simply as the sum of the element-wise product of two vectors.

    Clustering

    Generate batch embeddings to be used as input for algorithms like k-means or DBSCAN.

    // Similarity Scoring Example
let emb1 = await embed(container, "The cat sat on the mat")
let emb2 = await embed(container, "A cat was sitting on a rug")

// Cosine similarity (embeddings already normalized)
let similarity = sum(emb1 * emb2).item(Float.self)
print("Similarity: \(similarity)")

  11. Exclude pre-computed arrays from parameter loading

    main

    If you need to store an MLXArray that is pre-computed (like positionIds) but should not be treated as a loadable parameter, prefix the property name with an underscore (_).

    By naming the property with a leading underscore, the MLX parameter loader will ignore it, preventing validation failures that occur when the loader expects a key in the weights file that doesn't exist.

    fileprivate class VisionEmbeddings: Module, UnaryLayer {
    let positions: Int
    // Prefixing with underscore prevents this from being treated as a loadable parameter
    let _positionIds: MLXArray

    public init(_ config: PaliGemmaConfiguration.VisionConfiguration) {
        let d = config.imageSize / config.patchSize
        self.positions = d * d
        self._positionIds = MLXArray(0 ..< positions)[.newAxis, 0...]
    }
    // ...
}

  12. Map Python model structures to Swift patterns

    main

    When porting, follow these conventional naming patterns to map Python mlx-lm components to Swift MLX-Swift-LM components:

    Python patternSwift patternExample (Qwen2)
    @dataclass class ModelArgsstruct {ModelName}Configuration: Codable, SendableQwen2Configuration
    class Attention(nn.Module)class {ModelName}Attention: ModuleQwen2Attention
    class MLP(nn.Module)class {ModelName}MLP: Module, UnaryLayerQwen2MLP
    class TransformerBlock(nn.Module)class {ModelName}TransformerBlock: ModuleQwen2TransformerBlock
    class Model(nn.Module)class {ModelName}: Module, LLMModel, KVCacheDimensionProviderQwen2Model