[Bugfix][Frontend] Truncate the assistant tokens mask with the prompt (#54539)
Signed-off-by: hotragn <hotragn.pettugani_2024@woxsen.edu.in>
Hotragn Pettugani committed
Sep 1, 2026 at 01:53 UTC
ff0c3cb03c5d52fe21ebcf98f3e45325b1689116
2 files changed
+67
-11
tests/entrypoints/scale_out/render/test_render.py
+46
@@ -531,6 +531,52 @@ async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags
531
assert "Bye" in detok_rest.json()["prompt"]
532
533
534
+@pytest.mark.asyncio
535
+async def test_chat_render_assistant_tokens_mask_follows_truncation(client):
536
+ """The assistant mask must be truncated with the prompt it describes.
537
+
538
+ `assistant_tokens_mask` is positional: entry i labels token i. Truncating
539
+ `token_ids` from the left without truncating the mask leaves the two
540
+ describing different positions, and the mask ends up marking whichever
541
+ tokens happen to sit at the old offsets.
542
+ """
543
+ messages = [
544
+ # Deliberately lopsided: a long leading user turn and a short trailing
545
+ # one, so keeping the head of the mask is distinguishable from keeping
546
+ # its tail.
547
+ {"role": "user", "content": "Hello hello hello hello hello hello"},
548
+ {"role": "assistant", "content": "Hi!"},
549
+ {"role": "user", "content": "Bye"},
550
+ ]
551
+ body = {
552
+ "model": MODEL_NAME,
553
+ "messages": messages,
554
+ "chat_template": _TEMPLATE_WITH_GENERATION,
555
+ "return_assistant_tokens_mask": True,
556
+ }
557
+
558
+ full = await client.post("/v1/chat/completions/render", json=body)
559
+ assert full.status_code == 200
560
+ full_token_ids = full.json()["token_ids"]
561
+ full_mask = full.json()["assistant_tokens_mask"]
562
+ assert sum(full_mask) > 0
563
+
564
+ keep = len(full_token_ids) - 4
565
+ # Precondition: with this prompt the head and tail slices of the mask
566
+ # really do differ, so the assertion below can tell them apart.
567
+ assert full_mask[-keep:] != full_mask[:keep]
568
+
569
+ truncated = await client.post(
570
+ "/v1/chat/completions/render",
571
+ json={**body, "truncate_prompt_tokens": keep, "truncation_side": "left"},
572
+ )
573
+ assert truncated.status_code == 200
574
+ data = truncated.json()
575
+
576
+ assert data["token_ids"] == full_token_ids[-keep:]
577
+ assert data["assistant_tokens_mask"] == full_mask[-keep:]
578
+
579
+
580
@pytest.mark.asyncio
581
async def test_messages_render_basic(client):
582
"""Test basic Anthropic Messages render endpoint."""
vllm/renderers/params.py
+21
-11
@@ -24,6 +24,18 @@ logger = init_logger(__name__)
24
25
_S = TypeVar("_S", list[int], "torch.Tensor")
26
27
+# Prompt keys whose entry i describes token i of the prompt, so they must be
28
+# reduced with the same slice as the prompt tokens themselves. Anything added
29
+# here is truncated by `TokenizeParams.apply_post_tokenization`.
30
+# `_assistant_tokens_mask` is renderer-internal: `HfRenderer.render_messages`
31
+# stashes it on the prompt so `_process_tokens` can move it onto the engine
32
+# input.
33
+_PARALLEL_TO_PROMPT_TOKENS = (
34
+ "prompt_token_offsets",
35
+ "prompt_is_token_ids",
36
+ "_assistant_tokens_mask",
37
+)
38
+
39
40
def merge_kwargs(
41
defaults: dict[str, Any] | None,
@@ -420,8 +432,8 @@ class TokenizeParams:
432
"""The slice truncation applies to a sequence of `length` tokens.
433
434
`None` means no truncation. Anything parallel to the prompt tokens
423
- (see `prompt_token_offsets` and `prompt_is_token_ids`) must be reduced
424
- with this same slice to stay aligned with them.
435
+ must be reduced with this same slice to stay aligned with them; list
436
+ such keys in `_PARALLEL_TO_PROMPT_TOKENS`.
437
"""
438
max_length = self.truncate_prompt_tokens
439
if max_length is not None and max_length < 0:
@@ -511,15 +523,13 @@ class TokenizeParams:
523
tokenizer,
524
prompt["prompt_embeds"], # type: ignore[typeddict-item]
525
)
514
- offsets = cast(TokensPrompt, prompt).get("prompt_token_offsets")
515
- if offsets is not None:
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))
526
+ prompt_dict = cast(dict, prompt)
527
+ for key in _PARALLEL_TO_PROMPT_TOKENS:
528
+ parallel = prompt_dict.get(key)
529
+ if parallel is None:
530
+ continue
531
+ truncation = self._truncation_slice(tokenizer, len(parallel))
532
if truncation is not None:
523
- prompt["prompt_is_token_ids"] = is_token_ids[truncation] # type: ignore[typeddict-unknown-key]
533
+ prompt_dict[key] = parallel[truncation]
534
535
return prompt