mtmd: fix miscounting n_tokens

[ngxson]

Overview

Supersede #24634

The logic was a bit messy, the new logic is much easier to read:

• If n_temporal_merge == 1, output tokens is nx * ny * n_batch
• If n_temporal_merge > 1, output is ceil-divided by n_temporal_merge

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: no

[gabe-l-hart]

Not sure how, but it seems that this broke support for Granite 4 vision:

b9658

./bin/llama-mtmd-cli -m ~/models/ibm-granite/granite-vision-4.1-4b/granite-4B-vision-4.1-BF16.gguf --mmproj ~/models/ibm-granite/granite-vision-4.1-4b/mmproj-granite-vision-4b-4.1-BF16.gguf --image ~/Pictures/cat.jpg -p "Describe this image" --jinja

0.00.042.118 I common_init_result: fitting params to device memory ...
0.00.042.123 I common_init_result: (for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)
0.01.215.973 I common_init_from_params: warming up the model with an empty run - please wait ... (--no-warmup to disable)
0.01.281.045 I mtmd_cli_context: chat template example:
<|start_of_role|>system<|end_of_role|>You are a helpful assistant<|end_of_text|>
<|start_of_role|>user<|end_of_role|>        Hello<|end_of_text|>
<|start_of_role|>assistant<|end_of_role|>        Hi there<|end_of_text|>
<|start_of_role|>user<|end_of_role|>        How are you?<|end_of_text|>
<|start_of_role|>assistant<|end_of_role|>
0.01.598.427 I main: loading model: /Users/ghart/models/ibm-granite/granite-vision-4.1-4b/granite-4B-vision-4.1-BF16.gguf
0.01.598.435 W WARN: This is an experimental CLI for testing multimodal capability.
0.01.598.436 W       For normal use cases, please use the standard llama-cli

The image depicts a young kitten with a light cream or white coat, with striking blue eyes. The kitten is perched on what appears to be a brown or tan-colored fabric, possibly a piece of furniture. The kitten's expression is one of curiosity or attentiveness, as it gazes towards the camera. The kitten's fur is soft and fluffy, and its ears are upright, indicating alertness. The background is dark and out of focus, which helps to emphasize the kitten as the main subject of the image.


0.04.912.279 W ~llama_context:       MTL0 compute buffer size of 381.5886 MiB, does not match expectation of 226.0117 MiB

b9659

./bin/llama-mtmd-cli -m ~/models/ibm-granite/granite-vision-4.1-4b/granite-4B-vision-4.1-BF16.gguf --mmproj ~/models/ibm-granite/granite-vision-4.1-4b/mmproj-granite-vision-4b-4.1-BF16.gguf --image ~/Pictures/cat.jpg -p "Describe this image" --jinja

0.00.050.012 I common_init_result: fitting params to device memory ...
0.00.050.017 I common_init_result: (for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)
0.01.265.275 I common_init_from_params: warming up the model with an empty run - please wait ... (--no-warmup to disable)
0.01.340.758 I mtmd_cli_context: chat template example:
<|start_of_role|>system<|end_of_role|>You are a helpful assistant<|end_of_text|>
<|start_of_role|>user<|end_of_role|>        Hello<|end_of_text|>
<|start_of_role|>assistant<|end_of_role|>        Hi there<|end_of_text|>
<|start_of_role|>user<|end_of_role|>        How are you?<|end_of_text|>
<|start_of_role|>assistant<|end_of_role|>
0.02.551.383 I main: loading model: /Users/ghart/models/ibm-granite/granite-vision-4.1-4b/granite-4B-vision-4.1-BF16.gguf
0.02.551.389 W WARN: This is an experimental CLI for testing multimodal capability.
0.02.551.389 W       For normal use cases, please use the standard llama-cli




0.05.905.526 W ~llama_context:       MTL0 compute buffer size of 384.5916 MiB, does not match expectation of 226.0117 MiB

I'll look into a fix!

[gabe-l-hart]

@ngxson it looks like the removal of this short-circuit is what broke Granite 4 Vision. In the previous path, this triggers because n_temporal_merge == 1 which falls down into return nx * ny (724 * 1). Now without this, it falls down to return nx * ny * nz (724 * 1 * 7).

[gabe-l-hart]

I think the issue is that for llava-next models, nz is the number of total tiles versus the bug in the upstream issue (#24634) where it's the number of images in a multi-image batch.

[ngxson]

@gabe-l-hart can you check if this PR resolves the problems? #24732

for ref, the problem was no about token counting, but it was because the llava-uhd try to construct a batch, that conflicts with #24384

[gabe-l-hart]

Yep, that does it! Much better than my fix #24733