SIGN IN SIGN UP

feat(vlm): chunk the media-expanded prefill so a long VL chat stops hitting a wall

The media forward pushed the WHOLE media-expanded prompt through the language
model in ONE command buffer. Measured on an M5 Max: a 28,483-token prompt
returned kIOGPUCommandBufferCallbackErrorOutOfMemory with 89GB free -- the
failure is one enormous allocation, not exhausted memory. Because a chat
client re-sends the image on every turn, the prompt only grows, so the first
turn to cross the line kills the conversation and it never recovers.

The vision tower genuinely needs the whole image. The language model does not
need the whole sequence. Every VL wrapper in this tree already exposes that
seam -- qwen3_5 even pre-computes its mRoPE position_ids under the comment
"Pre-calculate position_ids for chunked prefill", and its language model
carries a "chunked prefill compatibility" mask check. The plumbing was built
for this and simply was not used.

_media_forward now merges the embeddings once and walks the language model in
chunks, carrying the cache forward and releasing allocator transients between
chunks.

CAPABILITY DETECTION, NOT A FAMILY LIST. The language model must NAME an
embeddings parameter -- `inputs_embeds`, or `input_embeddings` which is what
minimax spells it. **kwargs does NOT count: several wrappers here accept
`position_ids` only in the sense that it vanishes into **kwargs unread, so
counting that as support would build a chunked prefill on a model that
ignores the per-chunk positions and returns confident garbage. Anything
undetected, unsupported, or raising falls straight back to the one-shot call,
which is preserved verbatim.

Sizing is deliberately BIG. A smaller chunk does not reduce weight streaming,
it multiplies it: the chunk bounds only the terms that scale with it, while
the weights are re-read in full every chunk. dots3 restreams ~85GB of expert
weights per chunk. So the floor is 4096, and prompts under 8192 tokens keep
the one-shot path outright -- it reads the weights once and was never the
shape that failed.

Chunk boundaries snap away from the middle of a media placeholder run. Once
embeddings are merged this is harmless for every family here today, because
they all build masks from the cache offset. It stops being harmless the
moment a family builds a mask from whole-sequence image geometry -- gemma4's
config already declares use_bidirectional_attention="vision" even though the
MLX language model does not implement it, and qwen3_vl's deepstack injection
is keyed to visual rows in the current window. Snapping costs nothing.
`no_chunked_prefill` on a wrapper is honoured as a kill switch, and
VMLX_DISABLE_MEDIA_CHUNKED_PREFILL=1 forces one-shot.

Per-chunk extras that are not derivable from the cache offset are carried:
position_ids slices (qwen mRoPE), per_layer_inputs (gemma4, passed whole
because its language model self-slices by offset), image_mask slices (zaya's
per-token LoRA gate). `mask` is forced to None -- the engine's padding mask is
(B, seq) and must never reach a chunked call; muse actively overrides a bare
mask for its sliding layers.

19 new tests, including the one that matters: a model whose signature only
swallows **kwargs is NOT treated as chunkable.
J
Jinho Jang committed
831a305288d0bead9f0d9faec8bb8b3b784d39c3
Parent: 276574c