fix(mllama): image inference crash on CUDA, eager-only paged flag, vision head config - #2337
Open
subin9 wants to merge 3 commits into
Open
fix(mllama): image inference crash on CUDA, eager-only paged flag, vision head config#2337subin9 wants to merge 3 commits into
subin9 wants to merge 3 commits into
Conversation
VLlamaLoader::supports_paged_attention returned true, but MLlamaTextModel requires eager attention, so default launches aborted with "Expected eager attention implementation" instead of auto-disabling PagedAttention. Match the Qwen2-VL loaders and return false.
On decode steps the image block still ran (take_images clones rather than empties when has_changed_prompt is false, so has_images() stays true after prefill). The decode input has no <|image|> token, so the cross-attention mask collapsed to a zero-image tensor e.g. [1, 1, 0, 4]; to_dtype on that 0-element tensor aborts the CUDA stream with CUDA_ERROR_INVALID_VALUE (a harmless no-op on CPU, which is why text and CPU worked). Gate the image block on is_prompt, matching Qwen2VLImageProcessor. Cross-attention is skipped during decode anyway.
MLlamaVisionConfig::num_attention_heads had no serde alias for HF's `attention_heads` key, so the vision head count was ignored and always defaulted to 16. The stock 11B has 16 vision heads so it was masked; other mllama checkpoints loaded with wrong dims / a weight shape mismatch.
Code Metrics Report━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Language Files Lines Code Comments Blanks ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ C Header 23 4454 3116 790 548 CSS 3 282 252 6 24 CUDA 119 23681 19217 1704 2760 Dockerfile 1 35 19 9 7 HTML 2 27 27 0 0 JavaScript 3 577 562 12 3 Jinja2 7 694 656 5 33 JSON 27 15864 15861 0 3 Makefile 1 18 16 0 2 MDX 36 5901 0 4327 1574 Metal Shading Lan| 37 14287 11284 1136 1867 PowerShell 1 656 570 31 55 Python 151 12284 10191 484 1609 Shell 3 1062 843 117 102 Plain Text 53 10687 0 9209 1478 TOML 28 1368 1189 39 140 TypeScript 11 1641 1404 66 171 YAML 3 25 23 2 0 ───────────────────────────────────────────────────────────────────────────────── Jupyter Notebooks 3 122 83 23 16 |- Markdown 1 60 30 22 8 |- Python 1 122 113 1 8 (Total) 304 226 46 32 ───────────────────────────────────────────────────────────────────────────────── Markdown 273 11726 0 8736 2990 |- BASH 23 295 217 46 32 |- Dockerfile 2 5 5 0 0 |- JSON 6 289 289 0 0 |- PowerShell 1 1 1 0 0 |- Python 135 7239 6021 310 908 |- Rust 62 3820 2838 388 594 |- TOML 7 77 65 0 12 (Total) 23452 9436 9480 4536 ───────────────────────────────────────────────────────────────────────────────── Rust 674 299702 267081 5872 26749 |- Markdown 413 9761 452 8126 1183 (Total) 309463 267533 13998 27932 ───────────────────────────────────────────────────────────────────────────────── Svelte 19 1969 1826 51 92 |- CSS 1 4 4 0 0 |- JavaScript 19 921 767 25 129 (Total) 2894 2597 76 221 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Total 1478 429656 345022 41537 43097 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ |
subin9
force-pushed
the
fix/mllama-vision
branch
from
July 11, 2026 01:51
a733e55 to
4bf5e08
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes three issues that together make
Llama-3.2-11B-Vision(mllama /vllama) unusable for image inputs. Closes #2335.1. Image inference crashes on CUDA (main fix)
On decode steps the mllama inputs processor still ran its image block, because
Sequence::take_images()clones (rather than empties) the images whenhas_changed_promptis false, sohas_images()stays true after prefill. A decode input has no<|image|>token, soconvert_sparse_cross_attention_mask_to_denseproduced a cross-attention mask with a zero-sizedmax_num_imagesdimension, e.g.[1, 1, 0, 4](0 elements).prepare_cross_attention_maskthen calls.to_dtype(F32)on that tensor; launching a CUDA kernel over 0 elements returnsCUDA_ERROR_INVALID_VALUE. On CPU it is a harmless no-op, which is why text generation (and CPU) worked.Fix: gate the image block on
is_prompt, matchingQwen2VLImageProcessor(let mut pixel_values = if is_prompt { pixel_values } else { None };). Cross-attention is skipped during decode anyway (mllama/text.rs:if cross_attn_states.is_none() { continue }), so passing image inputs on decode was both wasteful and, on CUDA, fatal.2.
supports_paged_attentionadvertised support the model can't doVLlamaLoader::supports_paged_attentionreturnedtrue, butMLlamaTextModelrequires eager attention, so a default launch aborted withError: Expected eager attention implementationrather than auto-disabling PagedAttention. Returnfalse, matching the Qwen2-VL / Qwen2.5-VL loaders (which have the same eager requirement and correctly fall back).3. Vision
attention_headsconfig key ignoredMLlamaVisionConfig::num_attention_headshad#[serde(default = "d_attn_heads")](=16) but no alias for HF's actual keyattention_heads, so the vision head count fromconfig.jsonwas ignored and always 16. The stock 11B has 16 vision heads so it was masked; any mllama checkpoint with a different vision head count loaded with wrong dimensions or a weight shape mismatch. Added#[serde(alias = "attention_heads")].Reproduction / verification
Reproduced and verified all three on an RTX 4070 Ti SUPER (CUDA 12.9) using a tiny same-architecture model,
yujiepan/llama-3.2-vision-tiny-random(~22 MB), which hits the identicalMLlamaModelforward path without the 20 GB download. Before: default launch errors;--paged-attn off+ first image request ->CUDA_ERROR_INVALID_VALUE. After: model loads by default and small/large/repeat image requests plus text all return200with no CUDA errors. The stockmeta-llama/Llama-3.2-11B-Vision-Instructreproduces (1) and (2) identically.