Monolithic stateful dynamic prefill becomes nondeterministic above 16 tokens

[kylejfrost]

Summary

The macOS monolithic stateful LLM export path appears deterministic and HF-matching only when each prefill call updates the KV cache with query_len <= 16. For Qwen3 dynamic monolithic bundles, wider stateful prefill calls (17, 19, 64, or full batch prompt) produce process-to-process nondeterministic greedy tokens. The issue reproduces on a minimal public model (Qwen/Qwen3-0.6B) in fp16, so it is independent of 4-bit quantization.

This looks isolated to the stateful KV update/readback path (KVCache.update_and_fetch + dynamic mutable_slice_update + subsequent SDPA over the updated cache), not to tokenization, sampling, RoPE, or caix runtime chunking.

Affected Path

Relevant constants and code paths:

• coreai_models/export/_constants.py
  • QUANT_TRACE_QUERY_LEN = 16
  • QUANT_TRACE_OFFSET = 8
• coreai_models/models/macos/qwen3.py
  • attention uses SDPA(is_causal=True)
  • prefill/decode call cache.update_and_fetch(layer_idx, offset, key, value, seq_len=seq_len, query_len=query_len)
• coreai_models/primitives/macos/cache.py
  • KVCache.update_and_fetch uses mutable_slice_update to update K/V state for the current query range, then returns narrow(..., 0, seq_len) for attention.

Repro Setup

Model:

• Qwen/Qwen3-0.6B
• macOS monolithic dynamic bundle
• fp16, no compression
• max_context_length = 8192

Prompt set:

Name one primary color.
Write one factual sentence about local inference on Apple silicon.
Return a JSON object with keys "language" and "runtime" for Swift and Core AI.
Explain in two sentences why logits should stay on the final stage.
Write a Swift function signature for adding two Int values.
List three constraints for a safe distributed model test.
Answer with one word: what device family runs macOS?
Summarize pipeline parallelism in one short paragraph.

Reference HF fp16 greedy 8-token sequences:

576 4226 374 1112 30 6771 752 1744
4710 26567 702 11537 264 2205 44378 4712
576 4718 1265 614 264 897 315 330
7281 11 10339 304 1378 22870 3170 279
576 729 1265 470 279 2629 315 279
576 4226 1265 387 304 264 3175 11652
362 13 8162 362 13 8162 425 13
3555 374 279 6672 1948 15279 2142 323

Test method:

• Load the bundle in fresh processes / fresh engines.
• Run greedy generation for 8 tokens per prompt.
• Repeat 3 times and compare generated token sequences.
• Compare passing modes against the HF fp16 token oracle above.

Observed Results

Existing Qwen3-0.6B fp16 monolithic bundle

• Default caix mitigation (<=16 chunks): deterministic and HF-matching.
• Explicit chunk16: deterministic and HF-matching.
• chunk17: nondeterministic.
  • Example failure: prompt 2 token 0 changed from 576 to 1096.
• chunk64: nondeterministic.
  • Example failure: prompt 2 token 0 changed from 576 to 659.
• Full batch prefill: nondeterministic.
  • Example failure: prompt 2 token 0 changed from 4710 to 2055.

Trace-width experiments

I tested whether the failure was just the export reference query width (QUANT_TRACE_QUERY_LEN=16) by re-exporting Qwen3-0.6B fp16 with larger/equal trace widths.

Trace width 64

• default / chunk16: deterministic and HF-matching.
• chunk64: still nondeterministic.
  • Example failure: prompt 2 token 4 changed from 3405 to 1196.
• Full batch prefill: still nondeterministic.
  • Example failure: prompt 2 token 0 changed from 4710 to 576.

Trace width 19

This matches the known failing prompt length.

• default / chunk16: deterministic and HF-matching.
• chunk19: still nondeterministic.
  • Example failure: prompt 2 token 0 changed from 4710 to 2055.
• chunk20: still nondeterministic.
  • Example failure: prompt 2 token 1 changed from 565 to 334.
• Full batch prefill: still nondeterministic.
  • Example failure: prompt 2 token 0 changed from 4710 to 220.

Explicit causal mask experiment

I also tested whether the issue was caused by implicit SDPA(is_causal=True) geometry specialization. A Qwen3-0.6B fp16 export using an explicit in-graph causal attention mask still failed above 16:

• default / chunk16: deterministic and HF-matching.
• chunk17: nondeterministic.
  • Example failure: prompt 2 token 1 changed from 1196 to 1156.
• chunk64: nondeterministic.
  • Example failure: prompt 2 token 1 changed from 4416 to 32313.
• Full batch prefill: nondeterministic.
  • Example failure: prompt 2 token 0 changed from 4710 to 576.

Expected Behavior

For a dynamic stateful monolithic graph, a prefill call with query_len > 16 should be deterministic across fresh processes and should match the same greedy sequence produced by either:

1. HF fp16 reference, or
2. Repeated safe <=16 stateful prefill chunks.

Actual Behavior

The same bundle is stable and HF-matching when prefill is split into <=16 token calls, but wider stateful updates produce process-varying greedy results. The first failure appears exactly when a single stateful prefill update crosses 16 tokens.

Why this appears isolated to stateful update/fetch

Ruled out:

• Tokenization: same prompt IDs and same HF oracle.
• Sampling: greedy only; failures occur at token IDs returned by Core AI.
• Quantization: reproduces on fp16, no compression.
• Trace query width: trace19 and trace64 exports still fail above 16.
• Implicit causal mask geometry: explicit causal mask export still fails above 16.
• Runtime chunk policy: chunk16 is deterministic and HF-matching on the same model/bundle family.

Remaining common axis:

• A single Core AI stateful graph call with query_len > 16 performs a wider mutable_slice_update into K/V state and immediately reads the updated cache slice back into SDPA.

Impact

This forces production runtimes to chunk monolithic stateful prefill at <=16 for correctness. That avoids nondeterminism but costs substantial TTFT/prefill throughput. On the same machine/model class, caix decode is already at MLX parity, while prefill remains slower because the safe path must use many host-dispatched <=16 updates.

Ask

Is this a known limitation of Core AI dynamic mutable state / mutable_slice_update when the update width exceeds the trace-time query width, or a bug in the macOS monolithic stateful export pattern?

If there is a supported way to express a wider deterministic stateful prefill update, or a recommended replacement for this update_and_fetch pattern, guidance would unblock a correctness-safe fast prefill path.

[stikves]

Thanks for your report,

I've tried replicating it with both llm-runner and also running the model directly in Python. However I was unable to replicate the issue, and got consistent logits every time.

Can you describe your machine setup?

Hardware?
M1/M2/..., Base/Pro/Max/Ultra, RAM, GPU Cores?

Software?
macOS 27.0 Beta N? Xcode version?

That would be helpful to pin down the issue.