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