[Render][Speculator] Add return_loss_mask to render endpoint for training data generation (#46846)
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Ranran committed
Jun 28, 2026 at 02:07 UTC
09841ae705ce73967b3303cdda5c7046d7710f5f
7 files changed
+288
-20
tests/entrypoints/serve/render/test_render.py
+125
-1
@@ -14,7 +14,7 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
14
15
@pytest.fixture(scope="module")
16
def server():
17
- args: list[str] = []
17
+ args: list[str] = ["--trust-request-chat-template"]
18
19
with RemoteLaunchRenderServer(MODEL_NAME, args) as remote_server:
20
yield remote_server
@@ -369,3 +369,127 @@ async def test_completion_render_multiple_prompts_token_offsets(client):
369
assert len(offsets) == len(item["token_ids"])
370
for start, end in offsets:
371
assert 0 <= start <= end <= len(prompt)
372
+
373
+
374
+@pytest.mark.asyncio
375
+async def test_chat_completion_render_assistant_tokens_mask_default(client):
376
+ """Without return_assistant_tokens_mask, assistant_tokens_mask should be null."""
377
+ response = await client.post(
378
+ "/v1/chat/completions/render",
379
+ json={
380
+ "model": MODEL_NAME,
381
+ "messages": [
382
+ {"role": "user", "content": "Hello"},
383
+ {"role": "assistant", "content": "Hi!"},
384
+ {"role": "user", "content": "How are you?"},
385
+ ],
386
+ },
387
+ )
388
+
389
+ assert response.status_code == 200
390
+ data = response.json()
391
+ assert data.get("assistant_tokens_mask") is None
392
+
393
+
394
+@pytest.mark.asyncio
395
+async def test_chat_completion_render_assistant_tokens_mask_false(client):
396
+ """Explicitly setting return_assistant_tokens_mask=false gives null."""
397
+ response = await client.post(
398
+ "/v1/chat/completions/render",
399
+ json={
400
+ "model": MODEL_NAME,
401
+ "messages": [
402
+ {"role": "user", "content": "Hello"},
403
+ ],
404
+ "return_assistant_tokens_mask": False,
405
+ },
406
+ )
407
+
408
+ assert response.status_code == 200
409
+ data = response.json()
410
+ assert data.get("assistant_tokens_mask") is None
411
+
412
+
413
+@pytest.mark.asyncio
414
+async def test_chat_render_assistant_tokens_mask_null_without_gen_tags(
415
+ client,
416
+):
417
+ """The tiny test model lacks ``{% generation %}`` tags, so the mask is null."""
418
+ response = await client.post(
419
+ "/v1/chat/completions/render",
420
+ json={
421
+ "model": MODEL_NAME,
422
+ "messages": [
423
+ {"role": "user", "content": "Hello"},
424
+ {"role": "assistant", "content": "Hi!"},
425
+ ],
426
+ "return_assistant_tokens_mask": True,
427
+ },
428
+ )
429
+
430
+ assert response.status_code == 200
431
+ assert response.json().get("assistant_tokens_mask") is None
432
+
433
+
434
+# A minimal chat template with {% generation %} tags so we can test that
435
+# the mask correctly marks assistant tokens.
436
+_TEMPLATE_WITH_GENERATION = (
437
+ "{% for m in messages %}"
438
+ "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n"
439
+ "{% elif m['role'] == 'assistant' %}"
440
+ "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}"
441
+ "{% endif %}"
442
+ "{% endfor %}"
443
+)
444
+
445
+
446
+@pytest.mark.asyncio
447
+async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags(
448
+ client,
449
+):
450
+ """With a ``{% generation %}``-enabled template, the mask marks assistant
451
+ tokens and the masked tokens decode to the assistant content."""
452
+ response = await client.post(
453
+ "/v1/chat/completions/render",
454
+ json={
455
+ "model": MODEL_NAME,
456
+ "messages": [
457
+ {"role": "user", "content": "Hello"},
458
+ {"role": "assistant", "content": "Hi!"},
459
+ {"role": "user", "content": "Bye"},
460
+ ],
461
+ "chat_template": _TEMPLATE_WITH_GENERATION,
462
+ "return_assistant_tokens_mask": True,
463
+ },
464
+ )
465
+
466
+ assert response.status_code == 200
467
+ data = response.json()
468
+
469
+ mask = data["assistant_tokens_mask"]
470
+ token_ids = data["token_ids"]
471
+ assert mask is not None
472
+ assert isinstance(mask, list)
473
+ assert len(mask) == len(token_ids)
474
+ assert all(v in (0, 1) for v in mask)
475
+ assert sum(mask) > 0, "mask should mark at least one assistant token"
476
+
477
+ # Detokenize masked (assistant) and unmasked (non-assistant) tokens
478
+ # separately to verify the mask is correct, not just non-empty.
479
+ masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m]
480
+ unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m]
481
+
482
+ detok = await client.post(
483
+ "/detokenize",
484
+ json={"model": MODEL_NAME, "tokens": masked_ids},
485
+ )
486
+ assert detok.status_code == 200
487
+ assert "Hi!" in detok.json()["prompt"]
488
+
489
+ detok_rest = await client.post(
490
+ "/detokenize",
491
+ json={"model": MODEL_NAME, "tokens": unmasked_ids},
492
+ )
493
+ assert detok_rest.status_code == 200
494
+ assert "Hi!" not in detok_rest.json()["prompt"]
495
+ assert "Bye" in detok_rest.json()["prompt"]
vllm/entrypoints/openai/chat_completion/protocol.py
+13
@@ -407,6 +407,18 @@ class ChatCompletionRequest(OpenAIBaseModel):
407
),
408
)
409
410
+ return_assistant_tokens_mask: bool = Field(
411
+ default=False,
412
+ description=(
413
+ "If true, the /render response will include an "
414
+ "``assistant_tokens_mask`` field — a per-token list of 0/1 "
415
+ "values indicating which tokens were assistant-generated. "
416
+ "Requires the chat template to use ``{% generation %}`` "
417
+ "tags. When the template does not support it, "
418
+ "``assistant_tokens_mask`` will be ``null``."
419
+ ),
420
+ )
421
+
422
cache_salt: str | None = Field(
423
default=None,
424
description=(
@@ -520,6 +532,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
532
extra_kwargs,
533
),
534
media_io_kwargs=self.media_io_kwargs,
535
+ return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask),
536
)
537
538
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
vllm/entrypoints/serve/disagg/protocol.py
+8
@@ -75,6 +75,14 @@ class GenerateRequest(BaseModel):
75
token_ids: list[int] = Field(min_length=1)
76
"""The token ids to generate text from."""
77
78
+ assistant_tokens_mask: list[int] | None = None
79
+ """Per-token mask (1 = assistant-generated, 0 = not).
80
+
81
+ Only populated when the render request sets ``return_assistant_tokens_mask=True``
82
+ and the chat template supports ``{% generation %}``.
83
+ ``None`` when the mask was not requested or could not be computed.
84
+ """
85
+
86
@field_validator("token_ids")
87
@classmethod
88
def validate_token_ids(cls, v: list[int]) -> list[int]:
vllm/entrypoints/serve/render/serving.py
+22
@@ -127,11 +127,33 @@ class ServingRender(BaseServing):
127
)
128
params = request.to_sampling_params(max_tokens, self.default_sampling_params)
129
130
+ assistant_tokens_mask: list[int] | None = engine_input.get( # type: ignore[assignment]
131
+ "assistant_tokens_mask"
132
+ )
133
+ if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len(
134
+ token_ids
135
+ ):
136
+ logger.warning(
137
+ "assistant_tokens_mask length (%d) != token_ids length (%d); "
138
+ "this can happen with multimodal inputs where "
139
+ "placeholder expansion changes the token count. "
140
+ "The mask may be positionally misaligned.",
141
+ len(assistant_tokens_mask),
142
+ len(token_ids),
143
+ )
144
+ if len(assistant_tokens_mask) < len(token_ids):
145
+ assistant_tokens_mask.extend(
146
+ [0] * (len(token_ids) - len(assistant_tokens_mask))
147
+ )
148
+ else:
149
+ assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)]
150
+
151
request_id = f"chatcmpl-{random_uuid()}"
152
153
return GenerateRequest(
154
request_id=request_id,
155
token_ids=token_ids,
156
+ assistant_tokens_mask=assistant_tokens_mask,
157
features=self._extract_mm_features(engine_input),
158
sampling_params=params,
159
model=request.model,
vllm/inputs/engine.py
+10
@@ -42,6 +42,11 @@ class TokensInput(_InputOptions):
42
"""Char-level (start, end) offsets per token, propagated from the
43
renderer's TokensPrompt when offsets were computed."""
44
45
+ assistant_tokens_mask: NotRequired[list[int] | None]
46
+ """Per-token 0/1 mask marking assistant-generated tokens.
47
+ Populated when ``return_assistant_tokens_mask=True`` is set on the
48
+ render request and the chat template supports ``{% generation %}``."""
49
+
50
51
def tokens_input(
52
prompt_token_ids: list[int],
@@ -151,6 +156,11 @@ class MultiModalInput(_InputOptions):
156
`prompt_token_ids`.
157
"""
158
159
+ assistant_tokens_mask: NotRequired[list[int] | None]
160
+ """Per-token 0/1 mask marking assistant-generated tokens.
161
+ Populated when ``return_assistant_tokens_mask=True`` is set on the
162
+ render request and the chat template supports ``{% generation %}``."""
163
+
164
165
def mm_input(
166
prompt_token_ids: list[int],
vllm/renderers/hf.py
+106
-19
@@ -7,7 +7,7 @@ import inspect
7
import itertools
8
import weakref
9
from collections import defaultdict, deque
10
-from collections.abc import Sequence
10
+from collections.abc import Mapping, Sequence
11
from functools import lru_cache
12
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
13
@@ -678,6 +678,7 @@ def safe_apply_chat_template(
678
tools: list[dict[str, Any]] | None = ...,
679
chat_template: str | None = ...,
680
tokenize: Literal[True] = ...,
681
+ return_assistant_tokens_mask: Literal[False] = ...,
682
**kwargs,
683
) -> list[int]: ...
684
@overload
@@ -689,8 +690,20 @@ def safe_apply_chat_template(
690
tools: list[dict[str, Any]] | None = ...,
691
chat_template: str | None = ...,
692
tokenize: Literal[False] = ...,
693
+ return_assistant_tokens_mask: Literal[False] = ...,
694
**kwargs,
695
) -> str: ...
696
+@overload
697
+def safe_apply_chat_template(
698
+ model_config: ModelConfig,
699
+ tokenizer: HfTokenizer,
700
+ conversation: list[ConversationMessage],
701
+ *,
702
+ tools: list[dict[str, Any]] | None = ...,
703
+ chat_template: str | None = ...,
704
+ return_assistant_tokens_mask: Literal[True],
705
+ **kwargs,
706
+) -> tuple[list[int], list[int] | None]: ...
707
def safe_apply_chat_template(
708
model_config: ModelConfig,
709
tokenizer: HfTokenizer,
@@ -699,8 +712,9 @@ def safe_apply_chat_template(
712
tools: list[dict[str, Any]] | None = None,
713
chat_template: str | None = None,
714
tokenize: bool = True,
715
+ return_assistant_tokens_mask: bool = False,
716
**kwargs,
703
-) -> str | list[int]:
717
+) -> str | list[int] | tuple[list[int], list[int] | None]:
718
chat_template = resolve_chat_template(
719
tokenizer,
720
chat_template=chat_template,
@@ -728,6 +742,38 @@ def safe_apply_chat_template(
742
chat_template_kwargs=kwargs,
743
)
744
745
+ # assistant_tokens_mask requires tokenized output — force tokenize=True.
746
+ if return_assistant_tokens_mask:
747
+ tokenize = True
748
+
749
+ # When return_assistant_tokens_mask is requested and the template supports it,
750
+ # request assistant_tokens_mask via return_dict.
751
+ # Check for the actual Jinja tag, not just the word "generation"
752
+ # (which also appears in add_generation_prompt).
753
+ if return_assistant_tokens_mask and "{% generation %}" in chat_template:
754
+ resolved_kwargs["return_assistant_tokens_mask"] = True
755
+ resolved_kwargs["return_dict"] = True
756
+ resolved_kwargs.pop("tokenize", None)
757
+ try:
758
+ result = tokenizer.apply_chat_template(
759
+ conversation=conversation, # type: ignore[arg-type]
760
+ tools=tools, # type: ignore[arg-type]
761
+ chat_template=chat_template,
762
+ tokenize=True,
763
+ **resolved_kwargs,
764
+ )
765
+ except (TypeError, ValueError) as exc:
766
+ logger.warning(
767
+ "apply_chat_template failed for assistant_tokens_mask: %s", exc
768
+ )
769
+ else:
770
+ if isinstance(result, Mapping):
771
+ token_ids = list(result.get("input_ids", []))
772
+ mask_raw = result.get("assistant_masks")
773
+ mask = list(mask_raw) if mask_raw is not None else None
774
+ return token_ids, mask
775
+ return list(result), None
776
+
777
# transformers v5 changed the default of `return_dict` to True, which
778
# makes `apply_chat_template(tokenize=True)` return a `BatchEncoding`
779
# instead of `list[int]`. Force `return_dict=False` so downstream code
@@ -737,23 +783,24 @@ def safe_apply_chat_template(
783
resolved_kwargs["return_dict"] = False
784
785
try:
740
- return tokenizer.apply_chat_template(
786
+ plain = tokenizer.apply_chat_template(
787
conversation=conversation, # type: ignore[arg-type]
788
tools=tools, # type: ignore[arg-type]
789
chat_template=chat_template,
790
tokenize=tokenize,
791
**resolved_kwargs,
792
)
747
- # External library exceptions can sometimes occur despite the framework's
748
- # internal exception management capabilities.
793
except Exception as e:
750
- # Log and report any library-related exceptions for further
751
- # investigation.
794
logger.exception(
795
"An error occurred in `transformers` while applying chat template"
796
)
797
raise ValueError(str(e)) from e
798
799
+ if return_assistant_tokens_mask:
800
+ assert isinstance(plain, list), f"Expected list[int], got {type(plain)}"
801
+ return plain, None
802
+ return plain
803
+
804
805
def rebuild_mm_uuids_from_mm_data(
806
mm_uuids: MultiModalUUIDDict,
@@ -934,12 +981,22 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
981
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
982
chat_template_kwargs["tokenize"] = True
983
937
- prompt_raw = safe_apply_chat_template(
938
- model_config,
939
- tokenizer,
940
- conversation,
941
- **chat_template_kwargs,
942
- )
984
+ assistant_tokens_mask: list[int] | None = None
985
+ if params.return_assistant_tokens_mask:
986
+ prompt_raw, assistant_tokens_mask = safe_apply_chat_template(
987
+ model_config,
988
+ tokenizer,
989
+ conversation,
990
+ return_assistant_tokens_mask=True,
991
+ **chat_template_kwargs,
992
+ )
993
+ else:
994
+ prompt_raw = safe_apply_chat_template(
995
+ model_config,
996
+ tokenizer,
997
+ conversation,
998
+ **chat_template_kwargs,
999
+ )
1000
1001
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
1002
# model which uses unified vision chunks for both images and videos.
@@ -965,6 +1022,9 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1022
1023
prompt = parse_dec_only_prompt(prompt_raw)
1024
1025
+ if assistant_tokens_mask is not None:
1026
+ cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
1027
+
1028
# When `prompt_embeds` is mixed with other modality data,
1029
# `_process_tokens` runs `_process_multimodal` first (expanding
1030
# `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then
@@ -1038,12 +1098,30 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1098
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
1099
chat_template_kwargs["tokenize"] = True
1100
1041
- prompt_raw = await self._apply_chat_template_async(
1042
- model_config,
1043
- tokenizer,
1044
- conversation,
1045
- **chat_template_kwargs,
1046
- )
1101
+ assistant_tokens_mask: list[int] | None = None
1102
+ if params.return_assistant_tokens_mask:
1103
+ result_with_mask = cast(
1104
+ tuple[list[int], list[int] | None],
1105
+ await make_async(
1106
+ safe_apply_chat_template,
1107
+ executor=self._executor,
1108
+ )(
1109
+ model_config,
1110
+ tokenizer,
1111
+ conversation,
1112
+ return_assistant_tokens_mask=True, # type: ignore[arg-type]
1113
+ **chat_template_kwargs,
1114
+ ),
1115
+ )
1116
+ prompt_raw: str | list[int] = result_with_mask[0]
1117
+ assistant_tokens_mask = result_with_mask[1]
1118
+ else:
1119
+ prompt_raw = await self._apply_chat_template_async(
1120
+ model_config,
1121
+ tokenizer,
1122
+ conversation,
1123
+ **chat_template_kwargs,
1124
+ )
1125
1126
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
1127
# model which uses unified vision chunks for both images and videos.
@@ -1067,6 +1145,9 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1145
1146
prompt = parse_dec_only_prompt(prompt_raw)
1147
1148
+ if assistant_tokens_mask is not None:
1149
+ cast(dict, prompt)["_assistant_tokens_mask"] = assistant_tokens_mask
1150
+
1151
# See `render_messages` for the rationale.
1152
if prompt_embeds_tensors and mm_data:
1153
assert prompt_embeds_placeholder_token_id is not None
@@ -1108,6 +1189,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1189
processor records all placeholder offsets in the final (post-expansion)
1190
coordinate space, no offset shifting needed afterwards.
1191
"""
1192
+ assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
1193
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1194
if prompt_embeds_info is not None:
1195
tensors, placeholder_token_id = prompt_embeds_info
@@ -1123,6 +1205,8 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1205
tensors,
1206
mm_updates,
1207
)
1208
+ if assistant_tokens_mask is not None:
1209
+ engine_input["assistant_tokens_mask"] = assistant_tokens_mask
1210
return engine_input
1211
1212
@override
@@ -1133,6 +1217,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1217
skip_mm_cache: bool = False,
1218
) -> TokensInput | MultiModalInput:
1219
"""Async equivalent of `_process_tokens`."""
1220
+ assistant_tokens_mask = cast(dict, prompt).pop("_assistant_tokens_mask", None)
1221
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1222
if prompt_embeds_info is not None:
1223
tensors, placeholder_token_id = prompt_embeds_info
@@ -1150,6 +1235,8 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
1235
tensors,
1236
mm_updates,
1237
)
1238
+ if assistant_tokens_mask is not None:
1239
+ engine_input["assistant_tokens_mask"] = assistant_tokens_mask
1240
return engine_input
1241
1242
@staticmethod
vllm/renderers/params.py
+4
@@ -87,6 +87,9 @@ class ChatParams:
87
mm_processor_kwargs: dict[str, Any] | None = None
88
"""The kwargs to pass to the multi-modal processor."""
89
90
+ return_assistant_tokens_mask: bool = False
91
+ """Request a per-token assistant mask from apply_chat_template."""
92
+
93
def with_defaults(
94
self,
95
default_chat_template_kwargs: dict[str, Any] | None = None,
@@ -115,6 +118,7 @@ class ChatParams:
118
default_mm_processor_kwargs,
119
self.mm_processor_kwargs,
120
),
121
+ return_assistant_tokens_mask=self.return_assistant_tokens_mask,
122
)
123
124
def get_apply_chat_template_kwargs(self) -> dict[str, Any]: