[BUG] Qwen 3.5 VLM model crashes on subsequent requests with manual KV cache prefix matching

[neon52]

Describe the bug
When using VLMModelFactory to load a multimodal model (specifically Qwen 3.5), reusing a ModelContainer for multiple sequential requests fails if manual KV cache prefix matching is employed.

While this pattern (finding a matching KV cache "slot" and initializing a TokenIterator with remainingTokens) works reliably for text-only models, VLM models appear to maintain internal state that is not cleared or updated correctly between perform blocks unless the entire ModelContainer is reloaded via VLMModelFactory.shared.loadContainer().

To Reproduce

1. Load a Qwen 3.5 VLM model (for example: https://huggingface.co/mlx-community/Qwen3.5-9B-MLX-4bit) using VLMModelFactory.shared.loadContainer.
2. Perform the first inference request using any KV cache (it will always succeed since it's a fresh ModelContainer)
3. For the 2nd request load a KV cache that is different than what was used for the 1st request and attempt to initialize a TokenIterator for the remaining tokens.

On step 3 the app will crash with a broadcast_shapes mismatch error. Note that I can workaround the issue by completely reloading the model before step 3, but this somewhat defeats the purpose of KV cache reuse.

Example code snippet

// Current workaround: uncommenting the lines below for VLM models only makes everything work
//let freshModelContainer = try await VLMModelFactory.shared.loadContainer(configuration: modelContainer.configuration)
//modelContainer = freshModelContainer

// Find the best matching KV cache or create a new one
let preparedUserInput = try await modelContainer.prepare(input: userInput)
let preparedTokens = preparedUserInput.text.tokens.asArray(Int.self)
let (bestCache, remainingTokens) = findBestCache(preparedTokens)

// Perform inference using the loaded KV cache
await modelContainer.perform { context in
    let tokenIterator = try TokenIterator(
        input: LMInput(tokens: MLXArray(remainingTokens).reshaped([1, -1])), 
        model: context.model,
        cache: bestCache
    )
    generateTokenTask(
        promptTokenCount: preparedTokens.count, // use the full prompt size here instead of remainingTokens.count so the stats are correct
        modelConfiguration: context.configuration,
        tokenizer: context.tokenizer,
        iterator: tokenIterator
    )
}

I should reiterate that this same code works fine for multiple requests using all regular LLM models that I have tried, so the issue appears to be unique to VLM models. It can also work reliably for VLM models the workaround in the code snippet above, but that makes things much slower.

Expected behavior
VLMModel should allow for TokenIterator initialization from an existing KVCache state, provided the prefix matches, without requiring a full container reload. Internal vision-related states (like image grid offsets or positional indicators) should be reset or exposed so they can be managed alongside the KV cache.

Desktop (please complete the following information):

• OS Version: MacOS 15
• Device: M4 MacBook Pro
• mlx-swift-lm Version: latest main from 3 days ago (https://github.com/ml-explore/mlx-swift-lm/tree/3a7503df30e2cd100be933322dcf9b25d6459803)

[davidkoski]

This looks a bit like #149 -- the followup queries might be using the wrong shape which would result in a broadcast error (likely).

But the text of the description makes it sound like there is state held in the model itself. Reloading the weights (or just instantiating a new model and loading the same weights) sounds like it fixes the problem. Is that the case? This would be a far more serious error -- the model must be stateless outside the passed in KVCache.

[neon52]

I was able to reduce this even further. I found out that the crash occurs even if you create a complete new KV cache each time, so you can remove all the prefix matching stuff. Here's an example that should allow you to reproduce it on your end:

import Foundation
import MLX
import MLXLLM
import MLXVLM
@preconcurrency import MLXLMCommon

final class UncheckedSendableBox<T>: @unchecked Sendable {
    private var value: T?
    
    init(_ value: T) {
        self.value = value
    }
    
    func consume() -> T {
        guard let value else {
            fatalError("UncheckedSendableBox: value already consumed")
        }
        self.value = nil
        return value
    }
}

func runInference(origModelContainer: ModelContainer, history: [Chat.Message]) async throws {

    var modelContainer = origModelContainer

    // NOTE: uncommenting the lines below makes it work
    //let freshModelContainer = try await VLMModelFactory.shared.loadContainer(configuration: modelContainer.configuration)
    //modelContainer = freshModelContainer

    let generateParameters = GenerateParameters()
    let additionalContext: [String: any Sendable] = ["enable_thinking": false] // disable Qwen thinking
    let userInput = UserInput(chat: history, additionalContext: additionalContext)

    // apply chat template and retrieve tokens
    let preparedUserInput = try await modelContainer.prepare(input: userInput)
    let preparedTokens = preparedUserInput.text.tokens.asArray(Int.self)

    // load the best matching cache (in this case we just create a new cache each time)
    let bestCache = await modelContainer.perform { context in
        return context.model.newCache(parameters: generateParameters)
    }

    // separate the tokens into 2 phases (for example, if we want to save the KV cache after processing system prompt and tools only)
    let numTokensPhase1 = preparedTokens.count - 10 // simple hard coded split for this example, but in real life it could be chosen more intelligently
    let numTokensPhase2 = preparedTokens.count - numTokensPhase1
    if (numTokensPhase1 < 1) || (numTokensPhase2 < 1) {
        fatalError("both phases should have at least 1 token")
    }
    let tokensArray1 = Array(preparedTokens[0..<numTokensPhase1])
    let tokensArray2 = Array(preparedTokens[numTokensPhase1...])

    // phase 1 (prefill only)
    if (numTokensPhase1 > 0) {
        let cacheBox1 = UncheckedSendableBox(bestCache)
        await modelContainer.perform { context in
            let cache = cacheBox1.consume()
            print("mark1") // it will crash after this printout on the 2nd inference request
            let model = context.model
            _ = model(MLXArray(tokensArray1).expandedDimensions(axis: 0), cache: cache)
            print("mark2")
            for c in cache {
                for tensor in c.state {
                    tensor.eval() // force mlx to finish the computation before saving
                }
            }
            print("mark3")
        }
    }

    // save the KV cache "slot" (skipped for this example)

    // phase 2 (remaining prefill and all output generation)
    let cacheBox2 = UncheckedSendableBox(bestCache) // wrap non-sendable cache in a SendableBox so it can cross the isolation boundary
    let (stream, _) = try await modelContainer.perform { context in
        let cache = cacheBox2.consume()
        print("mark4")
        let remainingInput: LMInput
        if context.model is VLMModel {
            remainingInput = LMInput(tokens: MLXArray(tokensArray2).expandedDimensions(axis: 0))
        } else {
            remainingInput = LMInput(tokens: MLXArray(tokensArray2))
        }
        let tokenIterator = try TokenIterator( // it will crash here with broadcast_shapes error on the 2nd round
            input: remainingInput,
            model: context.model,
            cache: cache,
            parameters: generateParameters
        )
        print("mark5")
        return generateTokenTask(
            promptTokenCount: preparedTokens.count, // use the full prompt size here instead of tokensArray2.count so the stats are correct
            modelConfiguration: context.configuration,
            tokenizer: context.tokenizer,
            iterator: tokenIterator
        )
    }
    print("mark6")
    for await generation in stream {
        switch generation {
        case .token(let tokenID):
            print("output token ID: \(tokenID)")
        case .info(let info):
            print("generated \(info.generationTokenCount) total tokens")
        }
    }
}

// load the model
let url = URL(fileURLWithPath: "/path/to/model/Qwen3.5-9B-4bit") // downloaded from: https://huggingface.co/mlx-community/Qwen3.5-9B-MLX-4bit
let modelConfiguration = ModelConfiguration(directory: url)
//let modelContainer = try await LLMModelFactory.shared.loadContainer(configuration: modelConfiguration) // NOTE: loading as LLM instead of VLM will work
let modelContainer = try await VLMModelFactory.shared.loadContainer(configuration: modelConfiguration)
print("model loaded")

// define message history
let history1: [Chat.Message] = [
    .user("hi")
]
let history2: [Chat.Message] = [
    .user("hi"),
    .assistant("Hello! How can I help you today?"),
    .user("hello"),
]

// perform inference

print("starting inference round 1")
try await runInference(origModelContainer: modelContainer, history: history1)

print("starting inference round 2")
try await runInference(origModelContainer: modelContainer, history: history2)

print("finished")

When I run that code I get this output:

model loaded
starting inference round 1
mark1
mark2
mark3
mark4
mark5
mark6
output token ID: 9419
output token ID: 0
output token ID: 2500
output token ID: 628
output token ID: 353
output token ID: 1438
output token ID: 488
output token ID: 3242
output token ID: 30
generated 9 total tokens
starting inference round 2
mark1
MLX/ErrorHandler.swift:343: Fatal error: [broadcast_shapes] Shapes (1,16,23,64) and (1,1,10,64) cannot be broadcast. at /path/to/project/.build/checkouts/mlx-swift/Source/Cmlx/mlx-c/mlx/c/ops.cpp:2193

The 2 different "NOTE" comments in the code show ways that you can get it working by either reloading the model each time, or by loading as a LLM model instead of VLM. I did test the Qwen3-VL-4B model and surprisingly that worked, so this might be limited to Qwen 3.5 specifically, although it's possible other vision models I didn't test are also affected.

Now that I had such a simple example I was able to have Opus look through the repo and determine the root cause. Here's the analysis in case it helps:

------BEGIN LLM OUTPUT------

The crash is caused by stale precomputedPositionIds state persisting on the Qwen35Language.LanguageModel instance between inference rounds.

Here's the detailed mechanism:

1. Round 1, Phase 2: The TokenIterator calls Qwen35.prepare(), which calls languageModel.resetPositionState() (clearing state), then calls the language model. Since ropeDeltas == nil, getRopeIndex is called with the 10 phase-2 tokens. This sets precomputedPositionIds to shape [3, 1, 10] on the language model instance.

2. Round 1 completes: Token generation finishes, but precomputedPositionIds remains [3, 1, 10] and ropeDeltas remains zeros([1]) on the languageModel instance (these are never cleared).

3. Round 2, Phase 1: The user calls model(MLXArray(tokensArray1).reshaped([1, -1]), cache: cache) directly (via Qwen35.callAsFunction(_ inputs:, cache:)), bypassing prepare() and thus never calling resetPositionState().

4. Inside languageModel.callAsFunction():

   • cacheOffset = 0 (fresh cache)
   • The condition (cache != nil && cache?[model.faIdx] != nil && cacheOffset == 0) is true
   • Since precomputedPositionIds is non-nil (stale [3, 1, 10]), it is used directly:

     positionIds = precomputedPositionIds[0..., 0..., 0 ..< seqLength]

   • This slices a 10-position array for 23 tokens, producing position IDs for only 10 positions (MLX clips the range)

5. In the Attention layer: The RoPE cos/sin are computed for 10 positions [1, 1, 10, 64], but queries/keys have 23 positions [1, 16, 23, 64]. Broadcasting (1, 16, 23, 64) with (1, 1, 10, 64) fails because 23 ≠ 10.

Why it works with LLMModelFactory: The LLM version of Qwen35 (Qwen35TextModel) does NOT have precomputedPositionIds or ropeDeltas state. Its attention layer computes RoPE directly from cache.offset, which is always correct regardless of call history.

Why creating a fresh model container fixes it: A new model instance has precomputedPositionIds = nil and ropeDeltas = nil, so the stale state doesn't exist.

The fix: The Qwen35.callAsFunction(_ inputs: MLXArray, cache: [any KVCache]?) -> MLXArray method should call languageModel.resetPositionState() before calling the language model, or the user should explicitly call resetPositionState() when reusing the model with a fresh cache. Alternatively, the precomputedPositionIds caching logic should validate that the stored positions are compatible with the current input before reusing them (e.g., by checking that cacheOffset + seqLength <= precomputedPositionIds.dim(2)).

------END LLM OUTPUT------

I had only been testing dense models, but I went ahead and tried the sample code above with Qwen3.5-35B-A3B to compare against #149 but it behaved the same way. VLM load would crash, LLM load would work, and completely reloading the model each time made both work. Note that I don't get the "SmallVector out of range" error they mentioned in that issue, I always get the "broadcast_shapes" crash. That makes me think the 2 issues are separate.

[davidkoski]

Ah, I see. This is forbidden to have per-evaluation state in the model:

        func resetPositionState() {
            precomputedPositionIds = nil
            ropeDeltas = nil
        }

There is some mechanism to put things like that in LMOutput, but that probably isn't appropriate either.

I don't see these on the mlx-vlm side so need to investigate further, but the fix is clear: this needs to be removed.

[davidkoski]

Thank you for tracking this down!

[davidkoski]

Oh, I take that back, here it is:

        if pixel_values is not None:
            self._rope_deltas = None
            self._position_ids = None

It looks like it has been there all along and this is a faithful port, but still it won't work.

[davidkoski]

A note about how this state might work. The model evaluation produces LMOutput, which can include state:

public struct LMOutput {
    /// optional ``State`` to carry forward into the next step
    public let state: State?

    public struct State {
        public let crossAttentionStates: MLXArray?
    }

    public init(logits: MLXArray, state: LMOutput.State? = nil)
}

Additional properties could be added to state. Then this is hooked up here:

public protocol LanguageModel: Module {

    func callAsFunction(_ input: LMInput.Text, cache: [KVCache]?, state: LMOutput.State?)
        -> LMOutput

That can take the state back in and compute the values if they are missing. This is just evaluation state, not "permanent" state like the KVCache.

I think this is probably the right approach. It looks like this "stateful model" is a common pattern in mlx-vlm -- I think it is a bug there and certainly a BIG bug in mlx-swift-lm -- none of this is thread safe, much less correct when evaluating multiple responses. I think we need to add some docs to discuss this and the care that is needed.

[davidkoski]

Other instances:

GlmOcrTextModel:

        var _positionIds: MLXArray?
        var _ropeDeltas: MLXArray?

Qwen3VL (surprised this one hasn't show up earlier based on the name -- it will require weights in the model):

        private var ropeDeltas: MLXArray? = nil

and Qwen35 as noted above.

[davidkoski]

GPTOSS

        let sinksActive =
            cachedSinksActive
            ?? {
                let active = (sinks * sinks).max().item(Float.self) > 0
                cachedSinksActive = active
                return active
            }()

slightly different but still an illegal mutation that we should fix.

[alankessler]

Took a look at the code paths here. The fix seems straightforward in concept — move the mutable fields into LMOutput.State — but there's a wrinkle with how the VLM models hook into the protocol.

Qwen35, Qwen3VL, and GlmOcr all override the simplified callAsFunction(_ inputs: MLXArray, cache:) which doesn't have a state parameter. They'd need to switch to the state-aware variant (callAsFunction(_ input: LMInput.Text, cache:, state:)) to receive and return the position data.

GPTOSS is simpler — cachedSinksActive is just a lazy constant derived from weights, could compute it in init().

Want me to take a crack at the implementation, or do you want to handle this one given it touches the VLM core?