[Bugfix][Frontend] Truncate prompt_token_offsets with the prompt (#54407)
Signed-off-by: hotragn <hotragn.pettugani_2024@woxsen.edu.in>
Hotragn Pettugani committed
Aug 31, 2026 at 02:13 UTC
8fd9eb85d5a3d8f36cacb2b5d570a7a51c8200d8
3 files changed
+106
-8
tests/entrypoints/scale_out/render/test_render.py
+36
@@ -347,6 +347,42 @@ async def test_chat_render_default_no_token_offsets(client):
347
assert data["token_offsets"] is None
348
349
350
+@pytest.mark.asyncio
351
+async def test_completion_render_truncated_token_offsets(client):
352
+ """Truncation must shorten token_offsets together with token_ids.
353
+
354
+ An explicit truncation_side turns off tokenizer-level truncation, so the
355
+ tokenizer returns offsets for the whole prompt and they are reduced
356
+ afterwards -- separately from token_ids. GenerateRequest documents that the
357
+ two lists have equal length.
358
+ """
359
+ prompt = "The quick brown fox jumps over the lazy dog."
360
+ keep = 4
361
+ response = await client.post(
362
+ "/v1/completions/render",
363
+ json={
364
+ "model": MODEL_NAME,
365
+ "prompt": prompt,
366
+ "return_token_offsets": True,
367
+ "truncate_prompt_tokens": keep,
368
+ "truncation_side": "left",
369
+ },
370
+ )
371
+
372
+ assert response.status_code == 200
373
+ data = response.json()
374
+ token_ids = data[0]["token_ids"]
375
+ offsets = data[0]["token_offsets"]
376
+
377
+ assert len(token_ids) == keep
378
+ assert len(offsets) == len(token_ids)
379
+
380
+ # Equal length is not enough: truncating from the left keeps the *last*
381
+ # tokens, so the surviving offsets must cover the end of the prompt.
382
+ assert offsets[-1][1] == len(prompt)
383
+ assert offsets[0][0] > 0
384
+
385
+
386
@pytest.mark.asyncio
387
async def test_completion_render_multiple_prompts_token_offsets(client):
388
"""Each prompt in a batch gets its own offsets aligned with its tokens."""
tests/renderers/test_token_offsets.py
+42
@@ -197,3 +197,45 @@ class TestProcessTokensForwardsOffsets:
197
engine_input = renderer._process_tokens(tokens_prompt)
198
199
assert "prompt_token_offsets" not in engine_input
200
+
201
+
202
+class TestTruncationKeepsOffsetsAligned:
203
+ """``prompt_token_offsets`` runs parallel to ``prompt_token_ids``, and both
204
+ ``TokensPrompt`` and the render API's ``GenerateRequest`` document that the
205
+ two have equal length. Only ``prompt_token_ids`` was being truncated.
206
+ """
207
+
208
+ @pytest.mark.parametrize("side", ["left", "right"])
209
+ def test_explicit_truncation_side_truncates_offsets(self, fast_tokenizer, side):
210
+ renderer = _make_base_renderer_with(fast_tokenizer)
211
+ text = "The quick brown fox jumps over the lazy dog."
212
+ keep = 4
213
+
214
+ untruncated = renderer._tokenize_prompt(
215
+ {"prompt": text},
216
+ TokenizeParams(max_total_tokens=1024, return_token_offsets=True),
217
+ )
218
+ full_offsets = untruncated["prompt_token_offsets"]
219
+ assert len(full_offsets) > keep
220
+
221
+ # An explicit truncation_side disables tokenizer-level truncation (see
222
+ # get_encode_kwargs), so the tokenizer returns the whole sequence and
223
+ # truncation happens in apply_post_tokenization.
224
+ params = TokenizeParams(
225
+ max_total_tokens=1024,
226
+ return_token_offsets=True,
227
+ truncate_prompt_tokens=keep,
228
+ truncation_side=side,
229
+ )
230
+ result = params.apply_post_tokenization(
231
+ fast_tokenizer, renderer._tokenize_prompt({"prompt": text}, params)
232
+ )
233
+
234
+ offsets = result["prompt_token_offsets"]
235
+ assert len(result["prompt_token_ids"]) == keep
236
+ assert len(offsets) == keep
237
+
238
+ # Equal length is not enough: the surviving offsets must be the ones
239
+ # belonging to the surviving tokens.
240
+ expected = full_offsets[-keep:] if side == "left" else full_offsets[:keep]
241
+ assert offsets == expected
vllm/renderers/params.py
+28
-8
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass, field
4
-from typing import TYPE_CHECKING, Any, Literal, TypeVar
4
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
5
6
from vllm.exceptions import VLLMValidationError
7
from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
@@ -414,24 +414,39 @@ class TokenizeParams:
414
415
return tokens + [tokenizer.pad_token_id] * (pad_length - len(tokens))
416
417
- def _token_truncation(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
418
- """Apply truncation to prompt tokens if necessary."""
417
+ def _truncation_slice(
418
+ self, tokenizer: TokenizerLike | None, length: int
419
+ ) -> slice | None:
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.
425
+ """
426
max_length = self.truncate_prompt_tokens
427
if max_length is not None and max_length < 0:
428
max_length = self.max_input_tokens
429
423
- if max_length is None or max_length >= len(tokens):
424
- return tokens
430
+ if max_length is None or max_length >= length:
431
+ return None
432
if max_length == 0:
426
- return tokens[:0]
433
+ return slice(0, 0)
434
435
side = self.truncation_side or (
436
tokenizer.truncation_side if tokenizer is not None else None
437
)
438
if side == "left":
432
- return tokens[-max_length:]
439
+ return slice(-max_length, None)
440
+
441
+ return slice(0, max_length)
442
+
443
+ def _token_truncation(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
444
+ """Apply truncation to prompt tokens if necessary."""
445
+ truncation = self._truncation_slice(tokenizer, len(tokens))
446
+ if truncation is None:
447
+ return tokens
448
434
- return tokens[:max_length]
449
+ return tokens[truncation]
450
451
def _token_len_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S:
452
"""Apply length checks to prompt tokens if necessary."""
@@ -496,5 +511,10 @@ class TokenizeParams:
511
tokenizer,
512
prompt["prompt_embeds"], # type: ignore[typeddict-item]
513
)
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
520
return prompt