[Bugfix][Frontend] Truncate prompt_is_token_ids with the prompt (#54509)
Signed-off-by: hotragn <hotragn.pettugani_2024@woxsen.edu.in>
Hotragn Pettugani committed
Aug 31, 2026 at 05:29 UTC
28bf75c9a951b9871d2646df1c7828ee5a66e953
2 files changed
+54
-2
tests/renderers/test_chat_utils_prompt_embeds.py
+47
@@ -27,6 +27,7 @@ from vllm.entrypoints.chat_utils import (
27
parse_chat_messages_async,
28
)
29
from vllm.exceptions import VLLMValidationError
30
+from vllm.renderers import TokenizeParams
31
from vllm.renderers.hf import (
32
_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR,
33
_build_mixed_prompt_embeds,
@@ -575,3 +576,49 @@ async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn):
576
assert mask.count(False) == LEN_SYS + LEN_USR
577
assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_SYS], t_sys)
578
assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_USR], t_usr)
579
+
580
+
581
+def test_truncation_keeps_the_mixed_mask_aligned_with_the_prompt():
582
+ """`prompt_is_token_ids` must be truncated with `prompt_token_ids`.
583
+
584
+ The mask is positional: the engine pairs it element-wise with the prompt
585
+ tokens to decide which positions read a real token id and which read a row
586
+ of `prompt_embeds`. Truncating the prompt without truncating the mask
587
+ leaves the two describing different positions, so embed placeholders stop
588
+ being recognised as such.
589
+ """
590
+ hidden_size = 8
591
+ placeholder_token_id = 0
592
+ num_head, embed_len, num_tail = 10, 5, 5
593
+
594
+ token_ids = (
595
+ list(range(1, num_head + 1))
596
+ + [placeholder_token_id] * embed_len
597
+ + list(range(1, num_tail + 1))
598
+ )
599
+ tensors = [torch.randn(embed_len, hidden_size)]
600
+ positions = [(num_head, embed_len)]
601
+ embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions)
602
+
603
+ prompt = {
604
+ "prompt_token_ids": token_ids,
605
+ "prompt_embeds": embeds,
606
+ "prompt_is_token_ids": mask,
607
+ }
608
+ keep = 8
609
+ tok_params = TokenizeParams(
610
+ max_total_tokens=100,
611
+ truncate_prompt_tokens=keep,
612
+ truncation_side="left",
613
+ )
614
+
615
+ result = tok_params.apply_post_tokenization(None, prompt)
616
+
617
+ # Keeping the last 8 of 20 positions starts inside the embed span, so the
618
+ # first 3 surviving positions are embed rows and the rest are real tokens.
619
+ # Comparing contents rather than lengths is what pins the alignment: a
620
+ # mask reduced from the wrong side has the right length and still
621
+ # describes the wrong positions.
622
+ assert result["prompt_is_token_ids"] == [False] * 3 + [True] * num_tail
623
+ assert result["prompt_token_ids"] == token_ids[-keep:]
624
+ assert torch.equal(result["prompt_embeds"], embeds[-keep:])
vllm/renderers/params.py
+7
-2
@@ -420,8 +420,8 @@ class TokenizeParams:
420
"""The slice truncation applies to a sequence of `length` tokens.
421
422
`None` means no truncation. Anything parallel to the prompt tokens
423
- (see `prompt_token_offsets`) must be reduced with this same slice to
424
- stay aligned with them.
423
+ (see `prompt_token_offsets` and `prompt_is_token_ids`) must be reduced
424
+ with this same slice to stay aligned with them.
425
"""
426
max_length = self.truncate_prompt_tokens
427
if max_length is not None and max_length < 0:
@@ -516,5 +516,10 @@ class TokenizeParams:
516
truncation = self._truncation_slice(tokenizer, len(offsets))
517
if truncation is not None:
518
prompt["prompt_token_offsets"] = offsets[truncation] # type: ignore[typeddict-unknown-key]
519
+ is_token_ids = cast(EmbedsPrompt, prompt).get("prompt_is_token_ids")
520
+ if is_token_ids is not None:
521
+ truncation = self._truncation_slice(tokenizer, len(is_token_ids))
522
+ if truncation is not None:
523
+ prompt["prompt_is_token_ids"] = is_token_ids[truncation] # type: ignore[typeddict-unknown-key]
524
525
return prompt