diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index b537d2a13a..bbd4dc5c5f 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -27,6 +27,7 @@ from vllm.entrypoints.chat_utils import ( parse_chat_messages_async, ) from vllm.exceptions import VLLMValidationError +from vllm.renderers import TokenizeParams from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, _build_mixed_prompt_embeds, @@ -575,3 +576,49 @@ async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn): assert mask.count(False) == LEN_SYS + LEN_USR assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_SYS], t_sys) assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_USR], t_usr) + + +def test_truncation_keeps_the_mixed_mask_aligned_with_the_prompt(): + """`prompt_is_token_ids` must be truncated with `prompt_token_ids`. + + The mask is positional: the engine pairs it element-wise with the prompt + tokens to decide which positions read a real token id and which read a row + of `prompt_embeds`. Truncating the prompt without truncating the mask + leaves the two describing different positions, so embed placeholders stop + being recognised as such. + """ + hidden_size = 8 + placeholder_token_id = 0 + num_head, embed_len, num_tail = 10, 5, 5 + + token_ids = ( + list(range(1, num_head + 1)) + + [placeholder_token_id] * embed_len + + list(range(1, num_tail + 1)) + ) + tensors = [torch.randn(embed_len, hidden_size)] + positions = [(num_head, embed_len)] + embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions) + + prompt = { + "prompt_token_ids": token_ids, + "prompt_embeds": embeds, + "prompt_is_token_ids": mask, + } + keep = 8 + tok_params = TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=keep, + truncation_side="left", + ) + + result = tok_params.apply_post_tokenization(None, prompt) + + # Keeping the last 8 of 20 positions starts inside the embed span, so the + # first 3 surviving positions are embed rows and the rest are real tokens. + # Comparing contents rather than lengths is what pins the alignment: a + # mask reduced from the wrong side has the right length and still + # describes the wrong positions. + assert result["prompt_is_token_ids"] == [False] * 3 + [True] * num_tail + assert result["prompt_token_ids"] == token_ids[-keep:] + assert torch.equal(result["prompt_embeds"], embeds[-keep:]) diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 048f450ac1..3ecd8f4828 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -420,8 +420,8 @@ class TokenizeParams: """The slice truncation applies to a sequence of `length` tokens. `None` means no truncation. Anything parallel to the prompt tokens - (see `prompt_token_offsets`) must be reduced with this same slice to - stay aligned with them. + (see `prompt_token_offsets` and `prompt_is_token_ids`) must be reduced + with this same slice to stay aligned with them. """ max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: @@ -516,5 +516,10 @@ class TokenizeParams: truncation = self._truncation_slice(tokenizer, len(offsets)) if truncation is not None: prompt["prompt_token_offsets"] = offsets[truncation] # type: ignore[typeddict-unknown-key] + is_token_ids = cast(EmbedsPrompt, prompt).get("prompt_is_token_ids") + if is_token_ids is not None: + truncation = self._truncation_slice(tokenizer, len(is_token_ids)) + if truncation is not None: + prompt["prompt_is_token_ids"] = is_token_ids[truncation] # type: ignore[typeddict-unknown-key] return prompt