Skip to content

[BUG] KV prefix cache replays a previous image's features for token-identical vision requests (different images → identical answers) #2272

Description

@emilio-melroseinc

Describe the bug

When two /v1/chat/completions vision requests have identical text content (same prompt, same image placeholder token layout) but different images, the second request can silently reuse the KV prefix-cache entry built for the first image. The model then answers about the previous image — or, in a related degraded state, behaves as if no image was attached at all ("I cannot describe the image because no image was provided").

This is severe for programmatic captioning workloads, where every request uses the same prompt template: after the first request, every image in a batch receives the same caption. The failure is silent — responses are fluent and plausible — so downstream systems ingest wrong data with no error signal.

A freshly launched instance answers correctly, which rules out the encoders/models themselves; state accumulated in the cache layer is what corrupts subsequent requests.

We believe the hole is in src/exo/worker/engines/mlx/cache.py::_validate_media_match (see Analysis below). Related but distinct from #2151/#2152 (cache key omitting dimensions): here the images have different pixel bytes entirely, yet the KV entry is still reused because media-region validation is skipped or misses.

To Reproduce

Minimal, deterministic, no local files needed:

  1. Launch a fresh instance of any vision model (reproduced with mlx-community/gemma-4-e4b-it-8bit, mlx-community/Qwen3-VL-4B-Instruct-4bit, and mlx-community/Qwen3.5-122B-A10B-8bit, in both Tensor·RDMA and Pipeline·Ring modes).
  2. Send a solid red 64×64 PNG with the prompt "One word: what color?" → answers "Red" (correct).
  3. Immediately send a solid blue 64×64 PNG with the same prompt → answers "Red" (stale replay of request 1).
  4. Delete + relaunch the instance, send the blue square first → answers "Blue" (correct), confirming per-instance state, not model behavior.
import base64, json, urllib.request, zlib
from struct import pack

def png(rgb):  # minimal 64x64 solid-color RGB PNG
    w = h = 64
    raw = b"".join(b"\x00" + bytes(rgb) * w for _ in range(h))
    def chunk(t, d):
        c = t + d
        return pack(">I", len(d)) + c + pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
    return (b"\x89PNG\r\n\x1a\n"
            + chunk(b"IHDR", pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
            + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))

def ask(rgb, base="http://<exo-host>:52415", model="mlx-community/gemma-4-e4b-it-8bit"):
    url = "data:image/png;base64," + base64.b64encode(png(rgb)).decode()
    body = {"model": model, "messages": [{"role": "user", "content": [
        {"type": "image_url", "image_url": {"url": url}},
        {"type": "text", "text": "One word: what color?"}]}], "max_tokens": 200}
    req = urllib.request.Request(base + "/v1/chat/completions",
        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    r = json.loads(urllib.request.urlopen(req, timeout=180).read())
    return r["choices"][0]["message"]["content"].strip()

print("red  ->", ask((255, 0, 0)))   # "Red"  (correct, fresh instance)
print("blue ->", ask((0, 0, 255)))   # "Red"  (BUG: replay of request 1)

Real-photo variant: two different photographs sent with the identical prompt "Describe this shot in one sentence." return byte-identical (or trivially reworded) descriptions of whichever image was processed first. In one of our runs, five different film keyframes all received the same caption. In a further degraded state, requests return "no image was provided" for valid image payloads until the instance is relaunched.

Expected behavior

Each request's answer reflects the image actually attached to that request. If a cached prefix cannot be safely validated against the incoming request's images, the cache entry should be bypassed for the image region, never reused.

Analysis

In src/exo/worker/engines/mlx/cache.py, get_kv_cache() matches on token IDs via get_prefix_length(). Image placeholder tokens are identical across requests, so two vision requests with the same prompt template token-match through the entire image region. _validate_media_match() is the only guard, and it has silent-pass paths:

query_r = query_by_start.get(cached_r.start_pos)
if query_r is None:
    continue                      # ← cached image region with no positional twin is NOT validated
  • Validation only fires when a query media region starts at exactly the cached region's start_pos. Any positional skew (template/think-token fixups, differing n_tokens_per_image, empty query_regions because _find_media_regions found no placeholder runs) means the mismatched image is never compared, and the stale KV — which encodes the previous image's features — is reused as an "exact" match.
  • if not cached_regions: return match_length similarly trusts any entry cached without region metadata.

Suggested direction: treat unvalidatable overlap as a mismatch — if the token-level match extends into any cached media region that cannot be positively hash-matched against a query region at the same span, truncate the match to that region's start (fail closed instead of open).

Workaround for API clients

Prepend a unique nonce as a text content part before the image in each request. The token prefix then diverges ahead of the image region, so the cache can never extend into it. This fully resolved the issue for our captioning pipeline (verified across hundreds of requests).

{"role": "user", "content": [
  {"type": "text", "text": "[request 3f9a1c2e]"},
  {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}},
  {"type": "text", "text": "Describe this shot in one sentence."}]}

Environment

  • exo: main-branch build, ~June 2026 (post-v1.0.71; model registry includes Kimi K2.6/K2.7 cards)
  • Cluster: 4× Mac Studio (M3 Ultra, 256 GB), macOS 26.5.1 (25F80), Thunderbolt RDMA
  • Reproduced on: Tensor·MLX RDMA (multi-node) and Pipeline·MLX Ring (single-node) instances
  • Models: gemma-4-e4b-it-8bit, Qwen3-VL-4B-Instruct-4bit, Qwen3.5-122B-A10B-8bit (all mlx-community)
  • Client: plain /v1/chat/completions with OpenAI-style image_url data URLs (curl / Python urllib)

Related

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions