[Render][Speculator] Add return_loss_mask to render endpoint for training data generation (vllm-project#46846)

Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

Author: 2 people

tests/entrypoints/serve/render/test_render.py

@@ -14,7 +14,7 @@
 
 @pytest.fixture(scope="module")
 def server():
-    args: list[str] = []
+    args: list[str] = ["--trust-request-chat-template"]
 
     with RemoteLaunchRenderServer(MODEL_NAME, args) as remote_server:
         yield remote_server
@@ -369,3 +369,127 @@ async def test_completion_render_multiple_prompts_token_offsets(client):
         assert len(offsets) == len(item["token_ids"])
         for start, end in offsets:
             assert 0 <= start <= end <= len(prompt)
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_render_assistant_tokens_mask_default(client):
+    """Without return_assistant_tokens_mask, assistant_tokens_mask should be null."""
+    response = await client.post(
+        "/v1/chat/completions/render",
+        json={
+            "model": MODEL_NAME,
+            "messages": [
+                {"role": "user", "content": "Hello"},
+                {"role": "assistant", "content": "Hi!"},
+                {"role": "user", "content": "How are you?"},
+            ],
+        },
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data.get("assistant_tokens_mask") is None
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_render_assistant_tokens_mask_false(client):
+    """Explicitly setting return_assistant_tokens_mask=false gives null."""
+    response = await client.post(
+        "/v1/chat/completions/render",
+        json={
+            "model": MODEL_NAME,
+            "messages": [
+                {"role": "user", "content": "Hello"},
+            ],
+            "return_assistant_tokens_mask": False,
+        },
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data.get("assistant_tokens_mask") is None
+
+
+@pytest.mark.asyncio
+async def test_chat_render_assistant_tokens_mask_null_without_gen_tags(
+    client,
+):
+    """The tiny test model lacks ``{% generation %}`` tags, so the mask is null."""
+    response = await client.post(
+        "/v1/chat/completions/render",
+        json={
+            "model": MODEL_NAME,
+            "messages": [
+                {"role": "user", "content": "Hello"},
+                {"role": "assistant", "content": "Hi!"},
+            ],
+            "return_assistant_tokens_mask": True,
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.json().get("assistant_tokens_mask") is None
+
+
+# A minimal chat template with {% generation %} tags so we can test that
+# the mask correctly marks assistant tokens.
+_TEMPLATE_WITH_GENERATION = (
+    "{% for m in messages %}"
+    "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n"
+    "{% elif m['role'] == 'assistant' %}"
+    "{% generation %}Assistant: {{ m['content'] }}\n{% endgeneration %}"
+    "{% endif %}"
+    "{% endfor %}"
+)
+
+
+@pytest.mark.asyncio
+async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags(
+    client,
+):
+    """With a ``{% generation %}``-enabled template, the mask marks assistant
+    tokens and the masked tokens decode to the assistant content."""
+    response = await client.post(
+        "/v1/chat/completions/render",
+        json={
+            "model": MODEL_NAME,
+            "messages": [
+                {"role": "user", "content": "Hello"},
+                {"role": "assistant", "content": "Hi!"},
+                {"role": "user", "content": "Bye"},
+            ],
+            "chat_template": _TEMPLATE_WITH_GENERATION,
+            "return_assistant_tokens_mask": True,
+        },
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+
+    mask = data["assistant_tokens_mask"]
+    token_ids = data["token_ids"]
+    assert mask is not None
+    assert isinstance(mask, list)
+    assert len(mask) == len(token_ids)
+    assert all(v in (0, 1) for v in mask)
+    assert sum(mask) > 0, "mask should mark at least one assistant token"
+
+    # Detokenize masked (assistant) and unmasked (non-assistant) tokens
+    # separately to verify the mask is correct, not just non-empty.
+    masked_ids = [t for t, m in zip(token_ids, mask, strict=True) if m]
+    unmasked_ids = [t for t, m in zip(token_ids, mask, strict=True) if not m]
+
+    detok = await client.post(
+        "/detokenize",
+        json={"model": MODEL_NAME, "tokens": masked_ids},
+    )
+    assert detok.status_code == 200
+    assert "Hi!" in detok.json()["prompt"]
+
+    detok_rest = await client.post(
+        "/detokenize",
+        json={"model": MODEL_NAME, "tokens": unmasked_ids},
+    )
+    assert detok_rest.status_code == 200
+    assert "Hi!" not in detok_rest.json()["prompt"]
+    assert "Bye" in detok_rest.json()["prompt"]

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -407,6 +407,18 @@ class ChatCompletionRequest(OpenAIBaseModel):
         ),
     )
 
+    return_assistant_tokens_mask: bool = Field(
+        default=False,
+        description=(
+            "If true, the /render response will include an "
+            "``assistant_tokens_mask`` field — a per-token list of 0/1 "
+            "values indicating which tokens were assistant-generated. "
+            "Requires the chat template to use ``{% generation %}`` "
+            "tags.  When the template does not support it, "
+            "``assistant_tokens_mask`` will be ``null``."
+        ),
+    )
+
     cache_salt: str | None = Field(
         default=None,
         description=(
@@ -520,6 +532,7 @@ def build_chat_params(
                 extra_kwargs,
             ),
             media_io_kwargs=self.media_io_kwargs,
+            return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask),
         )
 
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:

vllm/entrypoints/serve/disagg/protocol.py

@@ -75,6 +75,14 @@ class GenerateRequest(BaseModel):
     token_ids: list[int] = Field(min_length=1)
     """The token ids to generate text from."""
 
+    assistant_tokens_mask: list[int] | None = None
+    """Per-token mask (1 = assistant-generated, 0 = not).
+
+    Only populated when the render request sets ``return_assistant_tokens_mask=True``
+    and the chat template supports ``{% generation %}``.
+    ``None`` when the mask was not requested or could not be computed.
+    """
+
     @field_validator("token_ids")
     @classmethod
     def validate_token_ids(cls, v: list[int]) -> list[int]:

vllm/entrypoints/serve/render/serving.py

@@ -127,11 +127,33 @@ async def render_chat_request(
         )
         params = request.to_sampling_params(max_tokens, self.default_sampling_params)
 
+        assistant_tokens_mask: list[int] | None = engine_input.get(  # type: ignore[assignment]
+            "assistant_tokens_mask"
+        )
+        if assistant_tokens_mask is not None and len(assistant_tokens_mask) != len(
+            token_ids
+        ):
+            logger.warning(
+                "assistant_tokens_mask length (%d) != token_ids length (%d); "
+                "this can happen with multimodal inputs where "
+                "placeholder expansion changes the token count. "
+                "The mask may be positionally misaligned.",
+                len(assistant_tokens_mask),
+                len(token_ids),
+            )
+            if len(assistant_tokens_mask) < len(token_ids):
+                assistant_tokens_mask.extend(
+                    [0] * (len(token_ids) - len(assistant_tokens_mask))
+                )
+            else:
+                assistant_tokens_mask = assistant_tokens_mask[: len(token_ids)]
+
         request_id = f"chatcmpl-{random_uuid()}"
 
         return GenerateRequest(
             request_id=request_id,
             token_ids=token_ids,
+            assistant_tokens_mask=assistant_tokens_mask,
             features=self._extract_mm_features(engine_input),
             sampling_params=params,
             model=request.model,

vllm/inputs/engine.py

@@ -42,6 +42,11 @@ class TokensInput(_InputOptions):
     """Char-level (start, end) offsets per token, propagated from the
     renderer's TokensPrompt when offsets were computed."""
 
+    assistant_tokens_mask: NotRequired[list[int] | None]
+    """Per-token 0/1 mask marking assistant-generated tokens.
+    Populated when ``return_assistant_tokens_mask=True`` is set on the
+    render request and the chat template supports ``{% generation %}``."""
+
 
 def tokens_input(
     prompt_token_ids: list[int],
@@ -151,6 +156,11 @@ class MultiModalInput(_InputOptions):
     `prompt_token_ids`.
     """
 
+    assistant_tokens_mask: NotRequired[list[int] | None]
+    """Per-token 0/1 mask marking assistant-generated tokens.
+    Populated when ``return_assistant_tokens_mask=True`` is set on the
+    render request and the chat template supports ``{% generation %}``."""
+
 
 def mm_input(
     prompt_token_ids: list[int],