[Deprecation] Remove assistant_token_mask support (#57520)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 19, 2026 at 00:17 UTC
56cac80b145ed1e9814bba519735ad4c71c530f5
7 files changed
+20
-334
tests/entrypoints/scale_out/render/test_render.py
-170
@@ -594,176 +594,6 @@ async def test_completion_render_multiple_prompts_token_offsets(client):
594
assert 0 <= start <= end <= len(prompt)
595
596
597
-@pytest.mark.asyncio
598
-async def test_chat_completion_render_assistant_tokens_mask_default(client):
599
- """Without return_assistant_tokens_mask, assistant_tokens_mask should be null."""
600
- response = await client.post(
601
- "/v1/chat/completions/render",
602
- json={
603
- "model": MODEL_NAME,
604
- "messages": [
605
- {"role": "user", "content": "Hello"},
606
- {"role": "assistant", "content": "Hi!"},
607
- {"role": "user", "content": "How are you?"},
608
- ],
609
- },
610
- )
611
-
612
- assert response.status_code == 200
613
- data = response.json()
614
- assert data.get("assistant_tokens_mask") is None
615
-
616
-
617
-@pytest.mark.asyncio
618
-async def test_chat_completion_render_assistant_tokens_mask_false(client):
619
- """Explicitly setting return_assistant_tokens_mask=false gives null."""
620
- response = await client.post(
621
- "/v1/chat/completions/render",
622
- json={
623
- "model": MODEL_NAME,
624
- "messages": [
625
- {"role": "user", "content": "Hello"},
626
- ],
627
- "return_assistant_tokens_mask": False,
628
- },
629
- )
630
-
631
- assert response.status_code == 200
632
- data = response.json()
633
- assert data.get("assistant_tokens_mask") is None
634
-
635
-
636
-@pytest.mark.asyncio
637
-async def test_chat_render_assistant_tokens_mask_null_without_gen_tags(
638
- client,
639
-):
640
- """The tiny test model lacks ``{% generation %}`` tags, so the mask is null."""
641
- response = await client.post(
642
- "/v1/chat/completions/render",
643
- json={
644
- "model": MODEL_NAME,
645
- "messages": [
646
- {"role": "user", "content": "Hello"},
647
- {"role": "assistant", "content": "Hi!"},
648
- ],
649
- "return_assistant_tokens_mask": True,
650
- },
651
- )
652
-
653
- assert response.status_code == 200
654
- assert response.json().get("assistant_tokens_mask") is None
655
-
656
-
657
-# A minimal chat template with {% generation %} tags so we can test that
658
-# the mask correctly marks assistant tokens.
659
-_TEMPLATE_WITH_GENERATION = (
660
- "{% for m in messages %}"
661
- "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n"
662
- "{% elif m['role'] == 'assistant' %}"
663
- "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}"
664
- "{% endif %}"
665
- "{% endfor %}"
666
-)
667
-
668
-
669
-@pytest.mark.asyncio
670
-async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags(
671
- client,
672
-):
673
- """With a ``{% generation %}``-enabled template, the mask marks assistant
674
- tokens and the masked tokens decode to the assistant content."""
675
- response = await client.post(
676
- "/v1/chat/completions/render",
677
- json={
678
- "model": MODEL_NAME,
679
- "messages": [
680
- {"role": "user", "content": "Hello"},
681
- {"role": "assistant", "content": "Hi!"},
682
- {"role": "user", "content": "Bye"},
683
- ],
684
- "chat_template": _TEMPLATE_WITH_GENERATION,
685
- "return_assistant_tokens_mask": True,
686
- },
687
- )
688
-
689
- assert response.status_code == 200
690
- data = response.json()
691
-
692
- mask = data["assistant_tokens_mask"]
693
- token_ids = data["token_ids"]
694
- assert mask is not None
695
- assert isinstance(mask, list)
696
- assert len(mask) == len(token_ids)
697
- assert all(v in (0, 1) for v in mask)
698
- assert sum(mask) > 0, "mask should mark at least one assistant token"
699
-
700
- # Detokenize masked (assistant) and unmasked (non-assistant) tokens
701
- # separately to verify the mask is correct, not just non-empty.
702
- masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m]
703
- unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m]
704
-
705
- detok = await client.post(
706
- "/detokenize",
707
- json={"model": MODEL_NAME, "tokens": masked_ids},
708
- )
709
- assert detok.status_code == 200
710
- assert "Hi!" in detok.json()["prompt"]
711
-
712
- detok_rest = await client.post(
713
- "/detokenize",
714
- json={"model": MODEL_NAME, "tokens": unmasked_ids},
715
- )
716
- assert detok_rest.status_code == 200
717
- assert "Hi!" not in detok_rest.json()["prompt"]
718
- assert "Bye" in detok_rest.json()["prompt"]
719
-
720
-
721
-@pytest.mark.asyncio
722
-async def test_chat_render_assistant_tokens_mask_follows_truncation(client):
723
- """The assistant mask must be truncated with the prompt it describes.
724
-
725
- `assistant_tokens_mask` is positional: entry i labels token i. Truncating
726
- `token_ids` from the left without truncating the mask leaves the two
727
- describing different positions, and the mask ends up marking whichever
728
- tokens happen to sit at the old offsets.
729
- """
730
- messages = [
731
- # Deliberately lopsided: a long leading user turn and a short trailing
732
- # one, so keeping the head of the mask is distinguishable from keeping
733
- # its tail.
734
- {"role": "user", "content": "Hello hello hello hello hello hello"},
735
- {"role": "assistant", "content": "Hi!"},
736
- {"role": "user", "content": "Bye"},
737
- ]
738
- body = {
739
- "model": MODEL_NAME,
740
- "messages": messages,
741
- "chat_template": _TEMPLATE_WITH_GENERATION,
742
- "return_assistant_tokens_mask": True,
743
- }
744
-
745
- full = await client.post("/v1/chat/completions/render", json=body)
746
- assert full.status_code == 200
747
- full_token_ids = full.json()["token_ids"]
748
- full_mask = full.json()["assistant_tokens_mask"]
749
- assert sum(full_mask) > 0
750
-
751
- keep = len(full_token_ids) - 4
752
- # Precondition: with this prompt the head and tail slices of the mask
753
- # really do differ, so the assertion below can tell them apart.
754
- assert full_mask[-keep:] != full_mask[:keep]
755
-
756
- truncated = await client.post(
757
- "/v1/chat/completions/render",
758
- json={**body, "truncate_prompt_tokens": keep, "truncation_side": "left"},
759
- )
760
- assert truncated.status_code == 200
761
- data = truncated.json()
762
-
763
- assert data["token_ids"] == full_token_ids[-keep:]
764
- assert data["assistant_tokens_mask"] == full_mask[-keep:]
765
-
766
-
597
@pytest.mark.asyncio
598
async def test_messages_render_basic(client):
599
"""Test basic Anthropic Messages render endpoint."""
vllm/entrypoints/openai/chat_completion/protocol.py
-13
@@ -453,18 +453,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
453
),
454
)
455
456
- return_assistant_tokens_mask: bool = Field(
457
- default=False,
458
- description=(
459
- "If true, the /render response will include an "
460
- "``assistant_tokens_mask`` field — a per-token list of 0/1 "
461
- "values indicating which tokens were assistant-generated. "
462
- "Requires the chat template to use ``{% generation %}`` "
463
- "tags. When the template does not support it, "
464
- "``assistant_tokens_mask`` will be ``null``."
465
- ),
466
- )
467
-
456
cache_salt: str | None = Field(
457
default=None,
458
min_length=1,
@@ -605,7 +593,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
593
extra_kwargs,
594
),
595
media_io_kwargs=self.media_io_kwargs,
608
- return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask),
596
# No-tools requests default to tool_choice="none" at the API
597
# layer. Collapse that default before rendering, so K3 emits a
598
# model-visible tool-choice instruction only for requests with a
vllm/entrypoints/scale_out/render/serving.py
-22
@@ -126,33 +126,11 @@ class ServingRender(BaseServing):
126
)
127
params = request.to_sampling_params(max_tokens, self.default_sampling_params)
128
129
- assistant_tokens_mask: list[int] | None = engine_input.get( # type: ignore[assignment]
130
- "assistant_tokens_mask"
131
- )
132
- if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len(
133
- token_ids
134
- ):
135
- logger.warning(
136
- "assistant_tokens_mask length (%d) != token_ids length (%d); "
137
- "this can happen with multimodal inputs where "
138
- "placeholder expansion changes the token count. "
139
- "The mask may be positionally misaligned.",
140
- len(assistant_tokens_mask),
141
- len(token_ids),
142
- )
143
- if len(assistant_tokens_mask) < len(token_ids):
144
- assistant_tokens_mask.extend(
145
- [0] * (len(token_ids) - len(assistant_tokens_mask))
146
- )
147
- else:
148
- assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)]
149
-
129
request_id = f"chatcmpl-{random_uuid()}"
130
131
return GenerateRequest(
132
request_id=request_id,
133
token_ids=token_ids,
155
- assistant_tokens_mask=assistant_tokens_mask,
134
features=self._extract_mm_features(engine_input),
135
sampling_params=params,
136
model=request.model,
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
-8
@@ -150,14 +150,6 @@ class GenerateRequest(BaseModel):
150
token_ids: list[int] = Field(min_length=1)
151
"""The token ids to generate text from."""
152
153
- assistant_tokens_mask: list[int] | None = None
154
- """Per-token mask (1 = assistant-generated, 0 = not).
155
-
156
- Only populated when the render request sets ``return_assistant_tokens_mask=True``
157
- and the chat template supports ``{% generation %}``.
158
- ``None`` when the mask was not requested or could not be computed.
159
- """
160
-
153
@field_validator("token_ids")
154
@classmethod
155
def validate_token_ids(cls, v: list[int]) -> list[int]:
vllm/inputs/engine.py
-10
@@ -43,11 +43,6 @@ class TokensInput(_InputOptions):
43
"""Char-level (start, end) offsets per token, propagated from the
44
renderer's TokensPrompt when offsets were computed."""
45
46
- assistant_tokens_mask: NotRequired[list[int] | None]
47
- """Per-token 0/1 mask marking assistant-generated tokens.
48
- Populated when ``return_assistant_tokens_mask=True`` is set on the
49
- render request and the chat template supports ``{% generation %}``."""
50
-
46
47
def tokens_input(
48
prompt_token_ids: list[int],
@@ -155,11 +150,6 @@ class MultiModalInput(_InputOptions):
150
`prompt_token_ids`.
151
"""
152
158
- assistant_tokens_mask: NotRequired[list[int] | None]
159
- """Per-token 0/1 mask marking assistant-generated tokens.
160
- Populated when ``return_assistant_tokens_mask=True`` is set on the
161
- render request and the chat template supports ``{% generation %}``."""
162
-
153
154
def mm_input(
155
prompt_token_ids: list[int],
vllm/renderers/hf.py
+20
-103
@@ -7,7 +7,7 @@ import inspect
7
import itertools
8
import weakref
9
from collections import defaultdict, deque
10
-from collections.abc import Mapping, Sequence
10
+from collections.abc import Sequence
11
from functools import lru_cache
12
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
13
@@ -758,7 +758,6 @@ def safe_apply_chat_template(
758
tools: list[dict[str, Any]] | None = ...,
759
chat_template: str | None = ...,
760
tokenize: Literal[True] = ...,
761
- return_assistant_tokens_mask: Literal[False] = ...,
761
**kwargs,
762
) -> list[int]: ...
763
@overload
@@ -770,20 +769,8 @@ def safe_apply_chat_template(
769
tools: list[dict[str, Any]] | None = ...,
770
chat_template: str | None = ...,
771
tokenize: Literal[False] = ...,
773
- return_assistant_tokens_mask: Literal[False] = ...,
772
**kwargs,
773
) -> str: ...
776
-@overload
777
-def safe_apply_chat_template(
778
- model_config: ModelConfig,
779
- tokenizer: HfTokenizer,
780
- conversation: list[ConversationMessage],
781
- *,
782
- tools: list[dict[str, Any]] | None = ...,
783
- chat_template: str | None = ...,
784
- return_assistant_tokens_mask: Literal[True],
785
- **kwargs,
786
-) -> tuple[list[int], list[int] | None]: ...
774
def safe_apply_chat_template(
775
model_config: ModelConfig,
776
tokenizer: HfTokenizer,
@@ -792,9 +779,8 @@ def safe_apply_chat_template(
779
tools: list[dict[str, Any]] | None = None,
780
chat_template: str | None = None,
781
tokenize: bool = True,
795
- return_assistant_tokens_mask: bool = False,
782
**kwargs,
797
-) -> str | list[int] | tuple[list[int], list[int] | None]:
783
+) -> str | list[int]:
784
chat_template = resolve_chat_template(
785
tokenizer,
786
chat_template=chat_template,
@@ -822,38 +808,6 @@ def safe_apply_chat_template(
808
chat_template_kwargs=kwargs,
809
)
810
825
- # assistant_tokens_mask requires tokenized output — force tokenize=True.
826
- if return_assistant_tokens_mask:
827
- tokenize = True
828
-
829
- # When return_assistant_tokens_mask is requested and the template supports it,
830
- # request assistant_tokens_mask via return_dict.
831
- # Check for the actual Jinja tag, not just the word "generation"
832
- # (which also appears in add_generation_prompt).
833
- if return_assistant_tokens_mask and "{% generation %}" in chat_template:
834
- resolved_kwargs["return_assistant_tokens_mask"] = True
835
- resolved_kwargs["return_dict"] = True
836
- resolved_kwargs.pop("tokenize", None)
837
- try:
838
- result = tokenizer.apply_chat_template(
839
- conversation=conversation, # type: ignore[arg-type]
840
- tools=tools, # type: ignore[arg-type]
841
- chat_template=chat_template,
842
- tokenize=True,
843
- **resolved_kwargs,
844
- )
845
- except (TypeError, ValueError) as exc:
846
- logger.warning(
847
- "apply_chat_template failed for assistant_tokens_mask: %s", exc
848
- )
849
- else:
850
- if isinstance(result, Mapping):
851
- token_ids = list(result.get("input_ids", []))
852
- mask_raw = result.get("assistant_masks")
853
- mask = list(mask_raw) if mask_raw is not None else None
854
- return token_ids, mask
855
- return list(result), None
856
-
811
# transformers v5 changed the default of `return_dict` to True, which
812
# makes `apply_chat_template(tokenize=True)` return a `BatchEncoding`
813
# instead of `list[int]`. Force `return_dict=False` so downstream code
@@ -879,9 +833,6 @@ def safe_apply_chat_template(
833
logger.warning("Chat template rejected the request: %s", e)
834
raise VLLMValidationError(_template_error_reason(e)) from e
835
882
- if return_assistant_tokens_mask:
883
- assert isinstance(plain, list), f"Expected list[int], got {type(plain)}"
884
- return plain, None
836
return plain
837
838
@@ -1066,22 +1017,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1017
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
1018
chat_template_kwargs["tokenize"] = True
1019
1069
- assistant_tokens_mask: list[int] | None = None
1070
- if params.return_assistant_tokens_mask:
1071
- prompt_raw, assistant_tokens_mask = safe_apply_chat_template(
1072
- model_config,
1073
- tokenizer,
1074
- conversation,
1075
- return_assistant_tokens_mask=True,
1076
- **chat_template_kwargs,
1077
- )
1078
- else:
1079
- prompt_raw = safe_apply_chat_template(
1080
- model_config,
1081
- tokenizer,
1082
- conversation,
1083
- **chat_template_kwargs,
1084
- )
1020
+ prompt_raw = safe_apply_chat_template(
1021
+ model_config,
1022
+ tokenizer,
1023
+ conversation,
1024
+ **chat_template_kwargs,
1025
+ )
1026
1027
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
1028
# model which uses unified vision chunks for both images and videos.
@@ -1107,9 +1048,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1048
1049
prompt = parse_dec_only_prompt(prompt_raw)
1050
1110
- if assistant_tokens_mask is not None:
1111
- cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
1112
-
1051
# When `prompt_embeds` is mixed with other modality data,
1052
# `_process_tokens` runs `_process_multimodal` first (expanding
1053
# `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then
@@ -1183,30 +1121,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1121
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
1122
chat_template_kwargs["tokenize"] = True
1123
1186
- assistant_tokens_mask: list[int] | None = None
1187
- if params.return_assistant_tokens_mask:
1188
- result_with_mask = cast(
1189
- tuple[list[int], list[int] | None],
1190
- await make_async(
1191
- safe_apply_chat_template,
1192
- executor=self._executor,
1193
- )(
1194
- model_config,
1195
- tokenizer,
1196
- conversation,
1197
- return_assistant_tokens_mask=True, # type: ignore[arg-type]
1198
- **chat_template_kwargs,
1199
- ),
1200
- )
1201
- prompt_raw: str | list[int] = result_with_mask[0]
1202
- assistant_tokens_mask = result_with_mask[1]
1203
- else:
1204
- prompt_raw = await self._apply_chat_template_async(
1205
- model_config,
1206
- tokenizer,
1207
- conversation,
1208
- **chat_template_kwargs,
1209
- )
1124
+ prompt_raw = await self._apply_chat_template_async(
1125
+ model_config,
1126
+ tokenizer,
1127
+ conversation,
1128
+ **chat_template_kwargs,
1129
+ )
1130
1131
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
1132
# model which uses unified vision chunks for both images and videos.
@@ -1232,9 +1152,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1152
1153
prompt = parse_dec_only_prompt(prompt_raw)
1154
1235
- if assistant_tokens_mask is not None:
1236
- cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
1237
-
1155
# See `render_messages` for the rationale.
1156
if prompt_embeds_tensors and mm_data:
1157
assert prompt_embeds_placeholder_token_id is not None
@@ -1276,7 +1193,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1193
processor records all placeholder offsets in the final (post-expansion)
1194
coordinate space, no offset shifting needed afterwards.
1195
"""
1279
- assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
1196
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1197
if prompt_embeds_info is not None:
1198
tensors, placeholder_token_id = prompt_embeds_info
@@ -1284,7 +1200,9 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1200
cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
1201
list(prompt["prompt_token_ids"]), mm_updates
1202
)
1203
+
1204
engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
1205
+
1206
if prompt_embeds_info is not None:
1207
tensors, _ = prompt_embeds_info
1208
self._apply_prompt_embeds_to_engine_input(
@@ -1292,8 +1210,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1210
tensors,
1211
mm_updates,
1212
)
1295
- if assistant_tokens_mask is not None:
1296
- engine_input["assistant_tokens_mask"] = assistant_tokens_mask
1213
+
1214
return engine_input
1215
1216
@override
@@ -1304,7 +1221,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1221
skip_mm_cache: bool = False,
1222
) -> TokensInput | MultiModalInput:
1223
"""Async equivalent of `_process_tokens`."""
1307
- assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
1224
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1225
if prompt_embeds_info is not None:
1226
tensors, placeholder_token_id = prompt_embeds_info
@@ -1312,9 +1228,11 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1228
cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
1229
list(prompt["prompt_token_ids"]), mm_updates
1230
)
1231
+
1232
engine_input = await super()._process_tokens_async(
1233
prompt, skip_mm_cache=skip_mm_cache
1234
)
1235
+
1236
if prompt_embeds_info is not None:
1237
tensors, _ = prompt_embeds_info
1238
self._apply_prompt_embeds_to_engine_input(
@@ -1322,8 +1240,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1240
tensors,
1241
mm_updates,
1242
)
1325
- if assistant_tokens_mask is not None:
1326
- engine_input["assistant_tokens_mask"] = assistant_tokens_mask
1243
+
1244
return engine_input
1245
1246
@staticmethod
vllm/renderers/params.py
-8
@@ -27,13 +27,9 @@ _S = TypeVar("_S", list[int], "torch.Tensor")
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.
30
_PARALLEL_TO_PROMPT_TOKENS = (
31
"prompt_token_offsets",
32
"prompt_is_token_ids",
36
- "_assistant_tokens_mask",
33
)
34
35
@@ -99,9 +95,6 @@ class ChatParams:
95
mm_processor_kwargs: dict[str, Any] | None = None
96
"""The kwargs to pass to the multi-modal processor."""
97
102
- return_assistant_tokens_mask: bool = False
103
- """Request a per-token assistant mask from apply_chat_template."""
104
-
98
tool_choice: Any | None = None
99
"""Request-level tool choice for renderers that need API metadata."""
100
@@ -136,7 +129,6 @@ class ChatParams:
129
default_mm_processor_kwargs,
130
self.mm_processor_kwargs,
131
),
139
- return_assistant_tokens_mask=self.return_assistant_tokens_mask,
132
tool_choice=self.tool_choice,
133
response_format=self.response_format,
134
)