[CB] [Major] Add multimodal support to continous batching by remi-or · Pull Request #46376 · huggingface/transformers

remi-or · 2026-06-03T10:27:22Z

This PR adds support for image and text models to continuous batching.

Tests run, will add image and audio perf numbers later.

This is more a beta versions, we are missing a bunch of features: CUDA graphs for encoder, multiple modalities at once, etc. But better start from here and build little by little IMO.

HuggingFaceDocBuilderDev · 2026-06-03T10:38:51Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-06-04T12:37:21Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46376&sha=3da698

zucchini-nlp · 2026-06-18T07:47:28Z

You might be affected by #46405. Hopefully useful for CB since the image embeddings will be of same sizes after it's merged, with each Nth image emb obtainable as features[N]`

molbap

Left some discussion around the structure (and naming) of the cache as I feel it'd benefit from a simplification as a simple request store. Other than that I tested the branch with Qwen2VL with a wide range of images sizes and it seems to work well, idling comes mostly from things independent from this PR that lie within either the processor (relatively cheap) or some convolution kernels that are either slow to run or big to allocate in the vision encoder.

Very excited, let's goo

molbap · 2026-06-23T10:45:46Z

+def find_num_kv_heads(config: PretrainedConfig) -> int | None:
+    """Finds the number of key-value heads for the given config."""
+    for attr in ["num_key_value_heads", "num_attention_heads"]:
+        if hasattr(config, attr):


won't that break?

I don't think so, this change is actually going to be live before this PR, albeit in a more readable form. What this does is it looks for num_key_value_heads if there is GQA and num_attention_heads if GQA is not there. If you see a reason it would break, lmk!

where is this change going to be?
And no sounds fine, I was looking at MLA/compressed latent but it is kind of an edge case for paged IMO

The more-prefill branch. And for MLA, it's easy: it isn't supported yet. It will be soon!

molbap · 2026-06-24T06:49:57Z

+        cache_size = max(16384, max_batch_tokens)
+        cache_shape = (cache_size, config.text_config.hidden_size)
+        self.cache = torch.empty(cache_shape, dtype=model_dtype, device=device)


so allocation will be max(16384, max_batch_tokens) × hidden × model_dtype no? might be worth underallocating

So I did the math on this, with bf16 and cache_size=16384 we will get:
VRAM_used = 16384 * hidden * sizeof(bf16) = hidden * 16 * 1024 * 2 bytes = hidden * 32 KiB
so with a big hidden_size like 8192 this is
VRAM_used = 8192 * 32 KiB = 8 * 32 MiB = 256 MiB
This is ok imo but if we use fp32 embeddings and double the hidden size this can go to 1Go, so I agree it is a lot.

The best way to solve this would be to have a "max_num_embeddings" attribute we can retrieve from the model: does that exist? I know there are dynamic resolution encoders but if there is truly no limit, then we need to add a hard stop for inputs "too big" and find a way to chunk these

IMO _get_num_multimodal_tokens returns what we need for ~40 models and happens before processing is ran. (input_ids == special_token_id).sum() happens during processing, but you can reuse it and it's enough. I don't think we need to add an attribute to the models like max_num_embeddings , and it's not really a model property

For the variable resolution it's not infinite ackshually, there's a max_pixels that translates into the maximum possible size to allocate, I computed a couple below:

model | tokens per image ================= paddleocr_vl | 2916 qwen2_vl | 1225 glm4v | 7,396 ernie4_5_vl_moe | 6084 video_llama_3 | 5041

However of course "images per request" is unbounded

I don't think _get_num_multimodal_tokens would be relevant here because we need a number before processing happens (this is at cache creation time). So 16384 would be a safe guess, but I can take it down to 8192 as well, seeing the values you posted above.

For reference, VLLM uses the size max(max_batch_tokens, max_mm_enbeddings_per_item) so having the function you used to compute the table above would be useful to size the cache. If you can paste it here, I can whip up something!

molbap · 2026-06-24T13:03:38Z

A general comment after some heatwaved thinking: the main issue I have with this is that it reuses the logic of a KV cache paged block allocator for blocks that are written once, read once during prefill, and then freed. so a simpler per-request store would be better no, wdyt? Embeddings just need to survive and be accessible between prefill chunks. So there'd be no need for allocate_blocks, also seems that can_store_mm_embeddings was unused?

So i'd vouch for this simplification if it makes sense for you

Actually, we kind of want to re-use the logic of the KV cache, because of a few reasons:

it makes the amount of VRAM used by the encoder cache known and bounded

when the encoder will have CUDA graphs, we will need static addresses, so might as well add them now

can_store_mm_embeddings is used to know if there is enough space to store the embeddings of the request before scheduling it
Does that make sense?

What I agree with is that we could have a tighter cache loop: right now, we allocate and free the embeddings in one block. While allocating the cache for the whole embeddings of one request seems like the best option (we don't want to starve mid-way) we could progressively de-allocate the cache as we go. Will look into it.

interesting.. Thought some more and I still want to discuss this cache design part, bear with me:

for the VRAM, doesn't a single running count of how many embedding tokens are alive set the limit naturally? we would know the number of tokens when receiving a request, and refuse or delay a new multimodal request when that total would exceed our max bound

this I am less familiar with but doesn't it touch more the encoder itself? your TODO about binning cuda graphs is that no? Assuming we've got cuda graphs in the encoder forward, then, wouldn't we just need to pass inputs_embedss which at that point is a static tensor? It happens IIUC before _get_forward_fn, so why is the embedding cache needed?

well we only need to know if there is enough space to schedule the request if we don't do a check like in 1) (the dumb running count)

So I thought about 1. and I think if we go down that road, we end up inventing paged cache again.

Say we move to an integer count instead: whenever a new embedding arrives, we allocate a new tensor for it, as long as it does not drive the number of embeddings above the limit we set. Internally, torch has to check if there is enough contiguous reserved space for that embedding tensor. If we are lucky, there is, we use the already reserved space and the GPU sees no increase in VRAM usage. If we are unlucky, there is enough reserved space, but it is not contiguous: torch has to reserve more space to fit our new tensor, and we have fragmentation. That's because the int keeps track of the total VRAM used, but not how it's arranged, which leads to "holes" in the VRAM (ie. wasted space).
The issue is that we required the new embeddings tensor to be contiguous: if we drop that requirement, we can arrange it however we like in memory, and we can ensure the space occupied by the embeddings is always under the limit we set. So we cut our embeddings into parts, and we arrange those parts in a fixed sized tensor. Hence: paged cache.

Also, for 2. ideally we want to have the "write into cache" part included in the graph, so we need it to have a fixed address. And even later on, when I implement the new cache system, I want the storage of the embeddings cache and the KV cache to reside on the same underlying tensor.

OK the fragmentation point is pretty compelling. How often do you reckon that will happen? Can we measure it? I agree it's important but I wonder how much it is important and if we can have numbers/impressions on benchmarks, it would make the point stronger!

But ok, it means we need a fixed address bounded tensor allocated once, and then the page table to avoid fragmentation.

For point 2 I am not sure it belongs in this PR's scope then, and could be done later. For now the graph does not touch the embeddings cache, so since the PR is already hefty I'd defer it to another one later on. Besides, just having a fixed buffer with a static address is enough, no? ( Imean we don't need the page table?)

Ah and I forgot to answer about point 3 but the "can store mm embeddings" is indeed needed for anything bounded, but it's an integer check, so it's also needed (and I think, simple) for the "counter" route

So curious about benchmarks/some measurements there as fragmentation can be a killer, but it's a very short-lived portion of the inference, so I want to quantify a bit before merging a rather non-trivial new API 😁

I agree it is not trivial but it is very similar to the PagedCache itself so it's not a new API either. And either way, we are going to make this a static tensor down the road for CUDA graphs, so it would be weird to not make it a static tensor from the beginning. As for the measurements, here they are:

eager per-request tensors static cache + encoder in graph pool Δ

Decode throughput 40.7 tok/s 40.4 tok/s equal (±1% run noise)

Peak reserved VRAM 75.69 GiB 74.86 GiB −0.83 GiB

Trapped memory (inactive-split peak) 1.41 GiB 1.20 GiB −15%

cudaMalloc calls over the run 471 135 −71%

On a single H100 with a workload with varying image sizes and request length. So there is a small measurable impact, which is good to know actually because we prefer static tensors for down-the-road CUDA graphs

molbap · 2026-06-24T13:08:56Z

+    supported_modalities = set(input_modalities).intersection({"text", "image", "audio"})
+    if len(supported_modalities) not in {1, 2} or "text" not in supported_modalities:
+        raise ValueError(f"This model supports {input_modalities = } but CB only supports text+image or text+audio.")
+


a request with video will have multimodal_inputs so they should be dropped, else it might raise with something weird like get_image_features(**video_kwargs)

Not sure I understand this comments: since a model with the video modality cannot be initialized with CB, there will never be a request w/ video data unless the user / server messes up

ah? but like qwen2_vl has get_video_features and video_token_id=151656, and check_modality_support(("image","video","text")) returns "image" so the model does init CB no? so a request could arrive with video no?

Oh yeah you are right, I forgot I changed this part to not error out, sorry. Yeah really good point. Is there a way to distinguish from video and image inputs? If not, I guess we can hard error for an unsupported modality like video.

Changed to only support 2 modalities for now -- I can change it back afterwards depending on the answer to the com above

molbap · 2026-06-24T13:17:27Z

        victims: list[RequestState] = []
        while demand > free_blocks and starved and num_active - len(victims) > 1:
            state, blocks_needed = starved.pop()
+            # If a request already consumed its MM inputs, since they cannot be retrieved, it cannot be offloaded


for me this is another argument towards a simple request store

IMO this is more a sign that the processor should be inside CB scope, that way we could hash the processors inputs and hash them + re-run the processor from those "small" CPU inputs. But then it's more of a down-the-road PR then.

molbap

Continued the cache discussion!

molbap · 2026-06-30T14:46:24Z

OK the fragmentation point is pretty compelling. How often do you reckon that will happen? Can we measure it? I agree it's important but I wonder how much it is important and if we can have numbers/impressions on benchmarks, it would make the point stronger!

But ok, it means we need a fixed address bounded tensor allocated once, and then the page table to avoid fragmentation.

For point 2 I am not sure it belongs in this PR's scope then, and could be done later. For now the graph does not touch the embeddings cache, so since the PR is already hefty I'd defer it to another one later on. Besides, just having a fixed buffer with a static address is enough, no? ( Imean we don't need the page table?)

Ah and I forgot to answer about point 3 but the "can store mm embeddings" is indeed needed for anything bounded, but it's an integer check, so it's also needed (and I think, simple) for the "counter" route

So curious about benchmarks/some measurements there as fragmentation can be a killer, but it's a very short-lived portion of the inference, so I want to quantify a bit before merging a rather non-trivial new API 😁

molbap · 2026-06-30T15:22:50Z

+        cache_size = max(16384, max_batch_tokens)
+        cache_shape = (cache_size, config.text_config.hidden_size)
+        self.cache = torch.empty(cache_shape, dtype=model_dtype, device=device)


IMO _get_num_multimodal_tokens returns what we need for ~40 models and happens before processing is ran. (input_ids == special_token_id).sum() happens during processing, but you can reuse it and it's enough. I don't think we need to add an attribute to the models like max_num_embeddings , and it's not really a model property

For the variable resolution it's not infinite ackshually, there's a max_pixels that translates into the maximum possible size to allocate, I computed a couple below:

model | tokens per image ================= paddleocr_vl | 2916 qwen2_vl | 1225 glm4v | 7,396 ernie4_5_vl_moe | 6084 video_llama_3 | 5041

However of course "images per request" is unbounded

github-actions · 2026-07-16T07:53:43Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29415017650:1
Result: failure | Jobs: 2 | Tests: 19 | Failures: 1 | Duration: 1m 52s

Code quality check failed: test jobs were skipped. Fix the code quality issues and push again to run tests.

ArthurZucker · 2026-07-23T06:38:01Z

+        encoder_kwargs = self.inputs_and_outputs.encoder_kwargs
+        if encoder_kwargs:


https://qiita.com/you22fy/items/3bfcd1b0fbe12ea2dda9

ArthurZucker · 2026-07-23T06:40:34Z

+# TODO: add hash-based indexing for multimodal inputs
+class EmbeddingsCache:
+    # One embedding is stored per row of the storage tensor. Rows are named this way to avoid confusion with the KV
+    # cache "blocks" (each of which spans block_size tokens).


can you add why we need a cache? -> prefix split i suppose? 😉 with ascii example

ArthurZucker · 2026-07-23T06:41:20Z

+        # Retrieve the actual token ID
+        token_id = None
+        for token_name in possible_token_names:
+            token_id = getattr(config, token_name, None)
+            if token_id is not None:
+                break


interesting, I would think the tokenizer is the better place to actually retrieve these

ArthurZucker · 2026-07-23T06:46:06Z

+            submask = rows_mask[past_length : past_length + query_length]
+            read_ids = submask.tolist()
+            missing_indices = query_length - len(read_ids)
+            # Check if any of the multimodal embeddings for this request are read in this batch
+            cache_read = (submask != -1).any().item()
+            # Check if all the multimodal embeddings for this request have been read
+            if past_length + query_length >= len(rows_mask):
+                to_free = True


i might be wrong, but I am not sure you have to go to list domain? torch ops should be better ?

only cast upon check passed?

ArthurZucker · 2026-07-23T06:47:02Z

+        # Store the multimodal embeddings in the cache
+        self.storage[allocated_rows] = mm_embedding
+
+    def release_cache_for_requests(self, requests: set[str]) -> None:


can this be async?, its not slow but wondering if you can just run this in bg ?

ArthurZucker · 2026-07-23T06:49:49Z

+                request_id = encoder_kw.pop(self.cache.embeddings_cache.REQUEST_ID_KEY)
+                # Run feature extractor and catch any errors (during feature extraction or storing in cache)
+                try:
+                    encoding_output = feature_extractor(**encoder_kw)


TODO but extractor should be run on device as well it can be much faster AFAIK

ArthurZucker · 2026-07-23T06:50:40Z

+        mm_embeddings = self.cache.embeddings_cache.storage[mm_embeddings_read_index]  # type: ignore
+        mm_embeddings = mm_embeddings.unsqueeze(0)  # shape [1, q_tokens, hidden_size]
+        mask = (mm_embeddings_read_index == -1).unsqueeze(-1)  # shape [q_tokens, 1]
+        inputs_embeds.copy_(torch.where(mask, inputs_embeds, mm_embeddings))


copy + torch.where vs masked_fill that we use usually?

ArthurZucker · 2026-07-23T06:51:02Z

+        inputs_embeds.copy_(torch.where(mask, inputs_embeds, mm_embeddings))
+
+    def _pop_or_get_input_ids(self, batch_data: PagedAttentionArgs) -> torch.Tensor:
+        """Retrieves the input ids from the batch data, popping it if the inputs_embeds are present."""


only used once not sure its needed

ArthurZucker

noice
just need to read the tests

ArthurZucker · 2026-07-23T06:51:48Z

+        """Fill the inputs_embeds tensor inside the batch_data dictionary."""
+        # Run the embedding layer to get all text tokens embeddings
+        inputs_embeds: torch.Tensor = batch_data["inputs_embeds"]  # shape [1, q_tokens, hidden_size]
+        embedding_module = model.get_input_embeddings()


i might be wrong, but this part is never compiled / never in the graph?(just checking)

zucchini-nlp · 2026-07-23T12:26:29Z

+        return None
+    if input_modalities == {"text", "image"}:
+        return "image"
+    if input_modalities == {"text", "audio"}:
+        return "audio"
+    if input_modalities == {"text", "image", "audio"}:
+        logger.warning(


a lot of VLMs support video and image under same backbone, we could check if input_modalities == {"text", "image", "video"} and raise warning that this model will work only with images instead of refusing to load it, no?

remi-or requested a review from molbap June 3, 2026 10:27

remi-or added this to Continuous batching Jun 9, 2026

remi-or moved this to In progress in Continuous batching Jun 9, 2026

remi-or mentioned this pull request Jun 9, 2026

[Continuous batching] Tracker for issues and feature requests #46516

Closed

remi-or force-pushed the cb-multimodal branch from 3da6981 to e6275de Compare June 18, 2026 07:09

remi-or force-pushed the cb-multimodal branch 3 times, most recently from a3797df to 24368e5 Compare June 24, 2026 12:43

molbap reviewed Jun 24, 2026

View reviewed changes

remi-or mentioned this pull request Jun 25, 2026

Fix continuous batching for composite / multimodal configs #46864

Closed

molbap reviewed Jun 30, 2026

View reviewed changes

remi-or force-pushed the cb-multimodal branch from 4223f3d to 1e385bb Compare July 9, 2026 02:41

remi-or added 4 commits July 14, 2026 03:02

Squash for the second time bc rebasing is hell

183f9d7

Fixes

8684392

Bake encoder run inside the mem pool

244ab50

Fix bug

8f449aa

remi-or force-pushed the cb-multimodal branch from 6661aec to 8f449aa Compare July 14, 2026 03:02

remi-or added 10 commits July 14, 2026 03:05

Get text cfg

65ba079

Small fixes

6f42cc7

Fix too many mm embeddings crash

f450b8d

Minor fixes

3bcc053

Shifting around args

8038f02

Fixes against dumb failures

8e3080d

Fixes final maybe who knows

de8e19c

add writing in try

6cc0bb1

FIx tests

06c85d1

Swap decorator order

926856f

remi-or self-assigned this Jul 16, 2026

remi-or marked this pull request as ready for review July 16, 2026 08:18

ArthurZucker reviewed Jul 23, 2026

View reviewed changes

zucchini-nlp reviewed Jul 23, 2026

View reviewed changes

remi-or mentioned this pull request Aug 3, 2026

Continuous batching (generate_batch) fails on models with heterogeneous per-layer num_key_value_heads/head_dim (e.g. Gemma4) #47721

Closed

4 tasks

	eager per-request tensors	static cache + encoder in graph pool	Δ
Decode throughput	40.7 tok/s	40.4 tok/s	equal (±1% run noise)
Peak reserved VRAM	75.69 GiB	74.86 GiB	−0.83 GiB
Trapped memory (inactive-split peak)	1.41 GiB	1.20 GiB	−15%
cudaMalloc calls over the run	471	135	−71%

		encoder_kwargs = self.inputs_and_outputs.encoder_kwargs
		if encoder_kwargs:

Uh oh!

Conversation

remi-or commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jun 3, 2026

Uh oh!

github-actions Bot commented Jun 4, 2026

Uh oh!

zucchini-nlp commented Jun 18, 2026

Uh oh!

molbap left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

remi-or Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

molbap left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 16, 2026

CI recap

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

remi-or commented Jun 3, 2026 •

edited

Loading

remi-or Jun 25, 2026 •

edited

Loading