Chunked Prefill VLM #3188
Chunked Prefill VLM #3188
Conversation
|
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. |
docs/source/reference/launcher.md
Outdated
| @@ -248,9 +248,18 @@ Options: | |||
| -p, --port <PORT> | |||
| The port to listen on | |||
|
|
|||
| [env: PORT=] | |||
| [env: PORT=80] | |||
There was a problem hiding this comment.
Did you change something ?
| @@ -116,11 +116,10 @@ def __init__(self, prefix: str, config, weights, layer_id): | |||
| ) | |||
| self.num_heads = config.num_attention_heads | |||
| self.hidden_size = config.hidden_size | |||
| if hasattr(config, "head_dim"): | |||
| if hasattr(config, "head_dim") and config.head_dim is not None: | |||
There was a problem hiding this comment.
| if hasattr(config, "head_dim") and config.head_dim is not None: | |
| if getattr(config, "head_dim", None) is not None: |
Nit
| @@ -81,27 +112,13 @@ def forward( | |||
| image_sizes: Optional[torch.Tensor] = None, | |||
| adapter_data: Optional[torch.Tensor] = None, | |||
| image_grid_thw: Optional[torch.LongTensor] = None, | |||
| inputs_embeds: Optional[torch.Tensor] = None, | |||
There was a problem hiding this comment.
Can we not have both, or either input_ids and input_embeds ?
Seems like it's an antipattern to accept both (since we don't know which one is valid.
It's totally fine to handle the embeddings before this step and only accept input_embeds imho
| self, | ||
| input_ids: torch.Tensor, | ||
| vision_embeds: torch.Tensor = None, | ||
| **kwargs, |
There was a problem hiding this comment.
Can we remove those kwargs everywhere they are not needed ?
| vision_embeds: torch.Tensor = None, | ||
| pixel_values: torch.FloatTensor = None, | ||
| pixel_attention_mask: torch.BoolTensor = None, | ||
| **kwargs, |
| ][0] | ||
| length = tokens.numel() | ||
|
|
||
| assert ( |
There was a problem hiding this comment.
Why do we care, I thought that chunking inside an image would be fine (in the LLM part where we do not care about the results anyway, since we're going to replace more of the input_embeds with the image embedding anyway)
|
|
||
| self.pixel_values = [] | ||
|
|
||
| for i, ( |
There was a problem hiding this comment.
A small assert that all these have the same lengths wouldn't hurt imho
| graph = torch.cuda.CUDAGraph() | ||
| self.cuda_graphs[bs] = { | ||
| "input_ids": input_ids, | ||
| "inputs_embeds": inputs_embeds, |
There was a problem hiding this comment.
There's a big issue to adding that here.
With 1M context window, we're now reserving 1M * ~4k (hidden_dim) * 2 (f16) = 8GB of VRAM for every graph we save. This seems prohibitive.
Instead of doing that, couldn't we send pixel_embeds (which are just the image tokens, which are capped at max_batch_prefill_tokens) and some pixel_embed_positions (a sequence_length mask of 0, 1 which is simply 1M * 8bit so 1MB worth of storage ?). And within the modeling code, we simply have:
input_embeds[pixel_embed_positions] = pixel_embeds ?
| "inputs_embeds": inputs_embeds, | ||
| "position_ids": position_ids, | ||
| "kv_cache": self.kv_cache, | ||
| "block_tables": block_tables, |
There was a problem hiding this comment.
Also duplicating this function in here, has a maintenance cost. I was happy we managed to avoid it until now (for sanity's sake).
I'm wondering if we could factor some stuff out, decompose this in a cleaner way.
This should definitely way for a subsequent PR in the current form (the current PR is already humonguous let's not make it even bigger), I'm just musing at how we could clean this up.
| batch.input_ids, vision_embeds=vision_embeds | ||
| ) | ||
|
|
||
| batch.inputs_embeds = inputs_embeds |
There was a problem hiding this comment.
The function name is not correct, get_xx should always return something.
If you're not returning something but instead acting on some object, the name can be set_xxx.
Also let's keep consistency and use input_embeds everywhere instead of input_embeddings.
What does this PR do?
Models not supported:
Fixes # (issue)
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.