Skip to content

fix(mllama): image inference crash on CUDA, eager-only paged flag, vision head config - #2337

Open
subin9 wants to merge 3 commits into
EricLBuehler:masterfrom
subin9:fix/mllama-vision
Open

fix(mllama): image inference crash on CUDA, eager-only paged flag, vision head config#2337
subin9 wants to merge 3 commits into
EricLBuehler:masterfrom
subin9:fix/mllama-vision

Conversation

@subin9

@subin9 subin9 commented Jul 11, 2026

Copy link
Copy Markdown

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 when has_changed_prompt is false, so has_images() stays true after prefill. A decode input has no <|image|> token, so convert_sparse_cross_attention_mask_to_dense produced a cross-attention mask with a zero-sized max_num_images dimension, e.g. [1, 1, 0, 4] (0 elements). prepare_cross_attention_mask then calls .to_dtype(F32) on that tensor; launching a CUDA kernel over 0 elements returns CUDA_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, matching Qwen2VLImageProcessor (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_attention advertised support the model can't do

VLlamaLoader::supports_paged_attention returned true, but MLlamaTextModel requires eager attention, so a default launch aborted with Error: Expected eager attention implementation rather than auto-disabling PagedAttention. Return false, matching the Qwen2-VL / Qwen2.5-VL loaders (which have the same eager requirement and correctly fall back).

3. Vision attention_heads config key ignored

MLlamaVisionConfig::num_attention_heads had #[serde(default = "d_attn_heads")] (=16) but no alias for HF's actual key attention_heads, so the vision head count from config.json was 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 identical MLlamaModel forward 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 return 200 with no CUDA errors. The stock meta-llama/Llama-3.2-11B-Vision-Instruct reproduces (1) and (2) identically.

subin9 added 3 commits July 11, 2026 10:41
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.
@github-actions

Copy link
Copy Markdown
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
subin9 force-pushed the fix/mllama-vision branch from a733e55 to 4bf5e08 Compare July 11, 2026 01:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Llama-3.2-11B-Vision (mllama): errors with PagedAttention by default, and image inference crashes on CUDA with --paged-attn off

1 participant