@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 950ee4c2e4

[API] Add token offsets to render endpoints (/v1/.../render) (#44226)

Signed-off-by: HyunKyun Moon <mhg5303@gmail.com>

Hyunkyun Moon committed Jun 26, 2026 at 21:02 UTC 950ee4c2e48fd462a03cebb3283ac9fdcb27b2e4
13 files changed +516 -20
tests/entrypoints/openai/test_render_token_offsets.py new
+80
@@ -0,0 +1,80 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the token-offsets request/response protocol wiring:
4 +the request flag flowing into ``TokenizeParams`` and the ``GenerateRequest``
5 +serialization boundary. End-to-end behavior is covered by
6 +``tests/entrypoints/serve/render/test_render.py``; plain Pydantic field
7 +storage is not retested here.
8 +"""
9 +
10 +from unittest.mock import Mock
11 +
12 +from vllm.config import ModelConfig
13 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
14 +from vllm.entrypoints.openai.completion.protocol import CompletionRequest
15 +from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
16 +from vllm.sampling_params import SamplingParams
17 +
18 +
19 +def _model_config() -> Mock:
20 + model_config = Mock(spec=ModelConfig)
21 + model_config.max_model_len = 128
22 + return model_config
23 +
24 +
25 +def test_completion_flag_forwarded_to_tok_params():
26 + """build_tok_params must forward return_token_offsets, defaulting to
27 + False (zero behavioral change for existing callers) and coercing JSON
28 + null to False via the bool() guard."""
29 + cfg = _model_config()
30 +
31 + default = CompletionRequest(model="m", prompt="hi")
32 + assert default.build_tok_params(cfg).return_token_offsets is False
33 +
34 + on = CompletionRequest(model="m", prompt="hi", return_token_offsets=True)
35 + assert on.build_tok_params(cfg).return_token_offsets is True
36 +
37 + null = CompletionRequest(model="m", prompt="hi", return_token_offsets=None)
38 + assert null.build_tok_params(cfg).return_token_offsets is False
39 +
40 +
41 +def test_chat_flag_forwarded_to_tok_params():
42 + """Chat build_tok_params has its own (max_completion_tokens) branch, so
43 + its return_token_offsets forwarding is verified independently."""
44 + cfg = _model_config()
45 + messages = [{"role": "user", "content": "hi"}]
46 +
47 + default = ChatCompletionRequest(model="m", messages=messages)
48 + assert default.build_tok_params(cfg).return_token_offsets is False
49 +
50 + on = ChatCompletionRequest(model="m", messages=messages, return_token_offsets=True)
51 + assert on.build_tok_params(cfg).return_token_offsets is True
52 +
53 + null = ChatCompletionRequest(
54 + model="m", messages=messages, return_token_offsets=None
55 + )
56 + assert null.build_tok_params(cfg).return_token_offsets is False
57 +
58 +
59 +def test_generate_request_token_offsets_default_none():
60 + """Defaults to None so existing /v1/.../render responses are unchanged."""
61 + req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=SamplingParams())
62 + assert req.token_offsets is None
63 +
64 +
65 +def test_generate_request_token_offsets_survive_json_round_trip():
66 + """GenerateRequest crosses the disagg serialization boundary; the
67 + tuple[int, int] offsets must survive model_dump and re-validate."""
68 + req = GenerateRequest(
69 + token_ids=[10, 20],
70 + sampling_params=SamplingParams(),
71 + token_offsets=[(0, 1), (1, 3)],
72 + )
73 + dumped = req.model_dump()
74 + assert dumped["token_offsets"] == [(0, 1), (1, 3)]
75 + # Re-validate from the dumped dict (sampling_params doesn't round-trip
76 + # cleanly via dump, so re-inject a fresh instance).
77 + again = GenerateRequest.model_validate(
78 + {**dumped, "sampling_params": SamplingParams()}
79 + )
80 + assert again.token_offsets == [(0, 1), (1, 3)]
tests/entrypoints/serve/render/test_render.py
+106
@@ -263,3 +263,109 @@ async def test_chat_completion_render_with_sampling_params(client):
263
264 # Check that internal fields are not present
265 assert "_all_stop_token_ids" not in sampling_params
266 +
267 +
268 +@pytest.mark.asyncio
269 +async def test_completion_render_emits_token_offsets(client):
270 + """With return_token_offsets, /v1/completions/render returns per-token
271 + (start, end) char offsets aligned with token_ids."""
272 + prompt = "Hello, world."
273 + response = await client.post(
274 + "/v1/completions/render",
275 + json={
276 + "model": MODEL_NAME,
277 + "prompt": prompt,
278 + "return_token_offsets": True,
279 + },
280 + )
281 +
282 + assert response.status_code == 200
283 + data = response.json()
284 + assert isinstance(data, list)
285 + offsets = data[0]["token_offsets"]
286 + assert offsets is not None
287 + assert len(offsets) == len(data[0]["token_ids"])
288 + for start, end in offsets:
289 + assert isinstance(start, int) and isinstance(end, int)
290 + assert 0 <= start <= end <= len(prompt)
291 +
292 +
293 +@pytest.mark.asyncio
294 +async def test_completion_render_default_no_token_offsets(client):
295 + """Without the flag, token_offsets must be null (existing responses
296 + unchanged)."""
297 + response = await client.post(
298 + "/v1/completions/render",
299 + json={
300 + "model": MODEL_NAME,
301 + "prompt": "Hello, world.",
302 + },
303 + )
304 +
305 + assert response.status_code == 200
306 + data = response.json()
307 + assert data[0]["token_offsets"] is None
308 +
309 +
310 +@pytest.mark.asyncio
311 +async def test_chat_render_emits_token_offsets(client):
312 + """With return_token_offsets, /v1/chat/completions/render returns
313 + per-token offsets relative to the templated prompt string."""
314 + response = await client.post(
315 + "/v1/chat/completions/render",
316 + json={
317 + "model": MODEL_NAME,
318 + "messages": [{"role": "user", "content": "Hello, world."}],
319 + "return_token_offsets": True,
320 + },
321 + )
322 +
323 + assert response.status_code == 200
324 + data = response.json()
325 + assert isinstance(data, dict)
326 + offsets = data["token_offsets"]
327 + assert offsets is not None
328 + assert len(offsets) == len(data["token_ids"])
329 + for start, end in offsets:
330 + assert isinstance(start, int) and isinstance(end, int)
331 + assert 0 <= start <= end
332 +
333 +
334 +@pytest.mark.asyncio
335 +async def test_chat_render_default_no_token_offsets(client):
336 + """Without the flag, chat render token_offsets must be null."""
337 + response = await client.post(
338 + "/v1/chat/completions/render",
339 + json={
340 + "model": MODEL_NAME,
341 + "messages": [{"role": "user", "content": "Hello, world."}],
342 + },
343 + )
344 +
345 + assert response.status_code == 200
346 + data = response.json()
347 + assert data["token_offsets"] is None
348 +
349 +
350 +@pytest.mark.asyncio
351 +async def test_completion_render_multiple_prompts_token_offsets(client):
352 + """Each prompt in a batch gets its own offsets aligned with its tokens."""
353 + prompts = ["Hello, world.", "Goodbye, world."]
354 + response = await client.post(
355 + "/v1/completions/render",
356 + json={
357 + "model": MODEL_NAME,
358 + "prompt": prompts,
359 + "return_token_offsets": True,
360 + },
361 + )
362 +
363 + assert response.status_code == 200
364 + data = response.json()
365 + assert len(data) == len(prompts)
366 + for item, prompt in zip(data, prompts):
367 + offsets = item["token_offsets"]
368 + assert offsets is not None
369 + assert len(offsets) == len(item["token_ids"])
370 + for start, end in offsets:
371 + assert 0 <= start <= end <= len(prompt)
tests/renderers/test_completions.py
+5
@@ -79,6 +79,11 @@ class DummyTokenizer:
79
80 return list(range(in_length))
81
82 + def __call__(self, text: str, **kwargs):
83 + # BaseRenderer._tokenize_prompt calls the tokenizer via __call__ (to
84 + # unify the output type), so mirror a real tokenizer's BatchEncoding.
85 + return {"input_ids": self.encode(text, **kwargs)}
86 +
87
88 def _build_renderer(
89 model_config: MockModelConfig,
tests/renderers/test_token_offsets.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for renderer-level token-offset behavior.
4 +
5 +These exercise ``_tokenize_prompt`` (offset extraction + capability/MM
6 +gating) and the ``_tokenize_prompt -> _process_tokens -> TokensInput``
7 +forwarding chain. Endpoint-level coverage lives in
8 +``tests/entrypoints/serve/render/test_render.py``.
9 +"""
10 +
11 +import pytest
12 +
13 +from vllm.renderers.params import TokenizeParams
14 +
15 +
16 +@pytest.fixture
17 +def fast_tokenizer():
18 + """gpt2 ships a Fast tokenizer; use it to test the offsets happy path."""
19 + from transformers import AutoTokenizer
20 +
21 + return AutoTokenizer.from_pretrained("openai-community/gpt2", use_fast=True)
22 +
23 +
24 +def _make_base_renderer_with(tokenizer):
25 + """Build a minimal BaseRenderer subclass that exposes the tokenizer so we
26 + can call ``_tokenize_prompt`` directly. BaseRenderer is abstract because of
27 + ``render_messages``; we just need a stub."""
28 + from vllm.renderers.base import BaseRenderer
29 +
30 + class _StubRenderer(BaseRenderer):
31 + def __init__(self, tok):
32 + # Bypass BaseRenderer.__init__ — we don't need a VllmConfig.
33 + from vllm.utils.async_utils import make_async
34 +
35 + self.tokenizer = tok
36 + self._executor = None
37 + # Mirror BaseRenderer.__init__: the async path offloads the sync
38 + # ``_tokenize_prompt`` to a thread pool.
39 + self._tokenize_prompt_async = make_async(self._tokenize_prompt)
40 + self.mm_processor = None
41 +
42 + def get_tokenizer(self):
43 + return self.tokenizer
44 +
45 + def _can_produce_offsets(self):
46 + # Mirror HfRenderer: offsets only for fast tokenizers.
47 + return self.tokenizer is not None and self.tokenizer.is_fast
48 +
49 + def render_messages(self, messages, params): # pragma: no cover
50 + raise NotImplementedError
51 +
52 + return _StubRenderer(tokenizer)
53 +
54 +
55 +class TestTokenizePromptOffsets:
56 + def test_fast_tokenizer_with_flag_returns_offsets(self, fast_tokenizer):
57 + renderer = _make_base_renderer_with(fast_tokenizer)
58 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
59 + prompt = {"prompt": "Hello, world."}
60 +
61 + result = renderer._tokenize_prompt(prompt, params)
62 +
63 + assert "prompt_token_ids" in result
64 + offsets = result["prompt_token_offsets"]
65 + assert offsets is not None
66 + # Length must match the token sequence, and each (start, end) is an
67 + # ordered pair within the source text.
68 + assert len(offsets) == len(result["prompt_token_ids"])
69 + text_len = len("Hello, world.")
70 + for s, e in offsets:
71 + assert isinstance(s, int) and isinstance(e, int)
72 + assert 0 <= s <= e <= text_len
73 +
74 + def test_base_renderer_without_override_yields_no_offsets(self, fast_tokenizer):
75 + """A renderer that does not override ``_can_produce_offsets`` never
76 + emits offsets, even with a fast tokenizer and the flag set. This locks
77 + in the base-default-False / subclass-override design."""
78 + from vllm.renderers.base import BaseRenderer
79 +
80 + class _BareRenderer(BaseRenderer):
81 + def __init__(self, tok):
82 + self.tokenizer = tok
83 + self._executor = None
84 + self.mm_processor = None
85 +
86 + def get_tokenizer(self):
87 + return self.tokenizer
88 +
89 + def render_messages(self, messages, params): # pragma: no cover
90 + raise NotImplementedError
91 +
92 + renderer = _BareRenderer(fast_tokenizer)
93 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
94 +
95 + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params)
96 +
97 + assert "prompt_token_offsets" not in result
98 +
99 + def test_default_flag_no_offsets(self, fast_tokenizer):
100 + renderer = _make_base_renderer_with(fast_tokenizer)
101 + params = TokenizeParams(max_total_tokens=None) # flag defaults False
102 +
103 + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params)
104 +
105 + # Field must be absent (not None) so TokensInput serialization stays
106 + # minimal for existing consumers.
107 + assert "prompt_token_offsets" not in result
108 +
109 + def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer):
110 + """Force is_fast=False to simulate a Slow tokenizer: the flag is set
111 + but offsets must not be returned because it cannot produce them."""
112 + from unittest.mock import PropertyMock, patch
113 +
114 + renderer = _make_base_renderer_with(fast_tokenizer)
115 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
116 +
117 + with patch.object(
118 + type(fast_tokenizer),
119 + "is_fast",
120 + new_callable=PropertyMock,
121 + return_value=False,
122 + ):
123 + result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params)
124 +
125 + assert "prompt_token_offsets" not in result
126 +
127 + @pytest.mark.parametrize("mm_key", ["multi_modal_data", "multi_modal_uuids"])
128 + def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key):
129 + """Offsets index the text prompt, which is meaningless once multimodal
130 + data is interleaved, so they are suppressed when MM inputs are present."""
131 + renderer = _make_base_renderer_with(fast_tokenizer)
132 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
133 + prompt = {"prompt": "Hello.", mm_key: {"image": ["x"]}}
134 +
135 + result = renderer._tokenize_prompt(prompt, params)
136 +
137 + assert "prompt_token_offsets" not in result
138 +
139 + @pytest.mark.asyncio
140 + async def test_tokenize_prompt_async_returns_offsets(self, fast_tokenizer):
141 + """The async path offloads the sync tokenizer; it must yield the same
142 + offsets as the sync path."""
143 + renderer = _make_base_renderer_with(fast_tokenizer)
144 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
145 +
146 + result = await renderer._tokenize_prompt_async(
147 + {"prompt": "Hello, world."}, params
148 + )
149 +
150 + offsets = result["prompt_token_offsets"]
151 + assert offsets is not None
152 + assert len(offsets) == len(result["prompt_token_ids"])
153 +
154 +
155 +class TestProcessTokensForwardsOffsets:
156 + """Tests that the ``_tokenize_prompt -> _process_tokens -> TokensInput``
157 + chain carries ``prompt_token_offsets`` through to the engine input.
158 + ``_process_tokens`` rebuilds the engine input from scratch, so it must
159 + copy the field explicitly. The sync and async variants are independent
160 + implementations, so both are checked.
161 + """
162 +
163 + def test_sync_forwards_offsets_to_engine_input(self, fast_tokenizer):
164 + renderer = _make_base_renderer_with(fast_tokenizer)
165 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
166 +
167 + tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params)
168 + # Sanity: offsets must reach the TokensPrompt, else this guards the
169 + # wrong layer.
170 + expected = tokens_prompt["prompt_token_offsets"]
171 +
172 + engine_input = renderer._process_tokens(tokens_prompt)
173 +
174 + assert engine_input["prompt_token_offsets"] == expected
175 +
176 + @pytest.mark.asyncio
177 + async def test_async_forwards_offsets_to_engine_input(self, fast_tokenizer):
178 + renderer = _make_base_renderer_with(fast_tokenizer)
179 + params = TokenizeParams(max_total_tokens=None, return_token_offsets=True)
180 +
181 + tokens_prompt = await renderer._tokenize_prompt_async(
182 + {"prompt": "Hello, world."}, params
183 + )
184 + expected = tokens_prompt["prompt_token_offsets"]
185 +
186 + engine_input = await renderer._process_tokens_async(tokens_prompt)
187 +
188 + assert engine_input["prompt_token_offsets"] == expected
189 +
190 + def test_no_offsets_forwarded_when_flag_off(self, fast_tokenizer):
191 + renderer = _make_base_renderer_with(fast_tokenizer)
192 + params = TokenizeParams(max_total_tokens=None) # flag defaults False
193 +
194 + tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params)
195 + assert "prompt_token_offsets" not in tokens_prompt
196 +
197 + engine_input = renderer._process_tokens(tokens_prompt)
198 +
199 + assert "prompt_token_offsets" not in engine_input
vllm/entrypoints/openai/chat_completion/protocol.py
+16
@@ -382,6 +382,21 @@ class ChatCompletionRequest(OpenAIBaseModel):
382 "need to map generated text back to input tokens."
383 ),
384 )
385 + return_token_offsets: bool | None = Field(
386 + default=False,
387 + description=(
388 + "If true, return char-level (start, end) offsets for each "
389 + "token relative to the tokenized source string in the "
390 + "`token_offsets` field of the rendered response. Only "
391 + "supported on the `/v1/completions/render` and "
392 + "`/v1/chat/completions/render` endpoints; ignored on regular "
393 + "generation endpoints. Honored only for Fast (Rust-backed) "
394 + "tokenizers; otherwise `token_offsets` is null. For chat "
395 + "requests, offsets are relative to the templated prompt "
396 + "string (after applying the chat template). Multimodal "
397 + "inputs and pre-tokenized inputs always yield null."
398 + ),
399 + )
400 return_prompt_text: bool | None = Field(
401 default=None,
402 description=(
@@ -524,6 +539,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
539 needs_detokenization=bool(self.echo and not self.return_token_ids),
540 max_total_tokens_param="max_model_len",
541 max_output_tokens_param=max_output_tokens_param,
542 + return_token_offsets=bool(self.return_token_offsets),
543 )
544
545 # Default sampling parameters for chat completion requests
vllm/entrypoints/openai/completion/protocol.py
+16
@@ -152,6 +152,21 @@ class CompletionRequest(OpenAIBaseModel):
152 "need to map generated text back to input tokens."
153 ),
154 )
155 + return_token_offsets: bool | None = Field(
156 + default=False,
157 + description=(
158 + "If true, return char-level (start, end) offsets for each "
159 + "token relative to the tokenized source string in the "
160 + "`token_offsets` field of the rendered response. Only "
161 + "supported on the `/v1/completions/render` and "
162 + "`/v1/chat/completions/render` endpoints; ignored on regular "
163 + "generation endpoints. Honored only for Fast (Rust-backed) "
164 + "tokenizers; otherwise `token_offsets` is null. For chat "
165 + "requests, offsets are relative to the templated prompt "
166 + "string (after applying the chat template). Multimodal "
167 + "inputs and pre-tokenized inputs always yield null."
168 + ),
169 + )
170
171 cache_salt: str | None = Field(
172 default=None,
@@ -209,6 +224,7 @@ class CompletionRequest(OpenAIBaseModel):
224 needs_detokenization=bool(self.echo and not self.return_token_ids),
225 max_total_tokens_param="max_model_len",
226 max_output_tokens_param="max_tokens",
227 + return_token_offsets=bool(self.return_token_offsets),
228 )
229
230 # Default sampling parameters for completion requests
vllm/entrypoints/serve/disagg/protocol.py
+7
@@ -82,6 +82,13 @@ class GenerateRequest(BaseModel):
82 raise ValueError("token_ids must not contain negative values")
83 return v
84
85 + token_offsets: list[tuple[int, int]] | None = None
86 + """Char-level (start, end) offsets per token, relative to the
87 + tokenized source string. Present only when the request set
88 + `return_token_offsets=True` and the renderer was able to compute
89 + them (Fast tokenizer, text input, no multimodal data). List length
90 + equals `token_ids` length when present. None otherwise."""
91 +
92 features: MultiModalFeatures | None = None
93 """Multimodal hashes and placeholder positions (populated for MM inputs)."""
94
vllm/entrypoints/serve/render/serving.py
+2
@@ -139,6 +139,7 @@ class ServingRender(BaseServing):
139 stream_options=(request.stream_options if request.stream else None),
140 cache_salt=request.cache_salt,
141 priority=request.priority,
142 + token_offsets=engine_input.get("prompt_token_offsets"),
143 )
144
145 async def render_completion_request(
@@ -194,6 +195,7 @@ class ServingRender(BaseServing):
195 stream_options=(request.stream_options if request.stream else None),
196 cache_salt=request.cache_salt,
197 priority=request.priority,
198 + token_offsets=engine_input.get("prompt_token_offsets"),
199 )
200 )
201
vllm/inputs/engine.py
+4
@@ -38,6 +38,10 @@ class TokensInput(_InputOptions):
38 prompt: NotRequired[str]
39 """The prompt text corresponding to the token IDs, if available."""
40
41 + prompt_token_offsets: NotRequired[list[tuple[int, int]] | None]
42 + """Char-level (start, end) offsets per token, propagated from the
43 + renderer's TokensPrompt when offsets were computed."""
44 +
45
46 def tokens_input(
47 prompt_token_ids: list[int],
vllm/inputs/llm.py
+6
@@ -115,6 +115,12 @@ class TokensPrompt(_PromptOptions):
115 token_type_ids: NotRequired[list[int]]
116 """A list of token type IDs to pass to the cross encoder model."""
117
118 + prompt_token_offsets: NotRequired[list[tuple[int, int]] | None]
119 + """Char-level (start, end) offsets per token, relative to the
120 + tokenized source string. Present only when offsets were requested
121 + AND a Fast (Rust-backed) tokenizer was used AND no multimodal data
122 + was present. The list length equals the length of `prompt_token_ids`."""
123 +
124
125 class EmbedsPrompt(_PromptOptions):
126 """Schema for a prompt provided via token embeddings."""
vllm/renderers/base.py
+65 -20
@@ -89,8 +89,13 @@ class BaseRenderer(ABC, Generic[_T]):
89 # to keep the asyncio event loop responsive under concurrent load.
90 self._mm_executor: Executor = self._executor
91
92 - # Offloading tokenizer encode & decode to thread pool.
93 - self._async_tokenizer_encode = make_async(self._encode, executor=self._executor)
92 + # Offload tokenization to the thread pool. The sync
93 + # ``_tokenize_prompt`` already encapsulates the unified ``__call__``
94 + # path and char-offset extraction, so the async variant is just it
95 + # offloaded (mirrors ``_process_multimodal_async`` below).
96 + self._tokenize_prompt_async = make_async(
97 + self._tokenize_prompt, executor=self._executor
98 + )
99 self._async_tokenizer_decode = make_async(self._decode, executor=self._executor)
100
101 self.mm_processor: BaseMultiModalProcessor | None = None
@@ -147,9 +152,6 @@ class BaseRenderer(ABC, Generic[_T]):
152 def _decode(self, *args, **kwargs):
153 return self.get_tokenizer().decode(*args, **kwargs)
154
150 - def _encode(self, *args, **kwargs):
151 - return self.get_tokenizer().encode(*args, **kwargs)
152 -
155 def get_mm_processor(self) -> "BaseMultiModalProcessor":
156 if self.mm_processor is None:
157 raise ValueError("Multi-modal processor not available for text-only models")
@@ -414,31 +416,64 @@ class BaseRenderer(ABC, Generic[_T]):
416 return self.render_messages(messages, params)
417
418 # Step 2: Tokenize prompts if necessary
417 - def _tokenize_prompt(
419 + def _can_produce_offsets(self) -> bool:
420 + """Whether this renderer's tokenizer can emit char-level offsets.
421 +
422 + Defaults to False; only renderers backed by an HF fast tokenizer
423 + (see ``HfRenderer``) can produce ``offset_mapping``.
424 + """
425 + return False
426 +
427 + def _wants_offsets(
428 self,
419 - prompt: TextPrompt,
420 - params: TokenizeParams,
421 - ) -> TokensPrompt:
422 - tokenizer = self.get_tokenizer()
423 - prompt_token_ids = tokenizer.encode(
424 - prompt["prompt"],
425 - **params.get_encode_kwargs(),
429 + prompt: "TextPrompt",
430 + params: "TokenizeParams",
431 + ) -> bool:
432 + return (
433 + params.return_token_offsets
434 + and self._can_produce_offsets()
435 + and not prompt.get("multi_modal_data")
436 + and not prompt.get("multi_modal_uuids")
437 )
438
428 - return TokensPrompt(prompt_token_ids=prompt_token_ids, **prompt)
439 + @staticmethod
440 + def _build_tokens_prompt(
441 + token_ids: Sequence[int],
442 + prompt: "TextPrompt",
443 + *,
444 + offset_mapping: Sequence[tuple[int, int]] | None = None,
445 + ) -> "TokensPrompt":
446 + """Build a TokensPrompt from already-extracted token ids.
447
430 - async def _tokenize_prompt_async(
448 + ``offset_mapping`` is the per-token ``(start, end)`` sequence from
449 + a BatchEncoding; pass it only when offsets were requested, and it
450 + is attached as ``prompt_token_offsets``.
451 + """
452 + if offset_mapping is not None:
453 + return TokensPrompt(
454 + prompt_token_ids=list(token_ids),
455 + prompt_token_offsets=[(int(s), int(e)) for s, e in offset_mapping],
456 + **prompt,
457 + )
458 + return TokensPrompt(prompt_token_ids=list(token_ids), **prompt)
459 +
460 + def _tokenize_prompt(
461 self,
462 prompt: TextPrompt,
463 params: TokenizeParams,
464 ) -> TokensPrompt:
435 - prompt_token_ids = await self._async_tokenizer_encode(
436 - prompt["prompt"],
437 - **params.get_encode_kwargs(),
465 + tokenizer = self.get_tokenizer()
466 + want_offsets = self._wants_offsets(prompt, params)
467 + kwargs = params.get_encode_kwargs()
468 + if want_offsets:
469 + kwargs = {**kwargs, "return_offsets_mapping": True}
470 + encoding = tokenizer(prompt["prompt"], **kwargs)
471 + return self._build_tokens_prompt(
472 + encoding["input_ids"],
473 + prompt,
474 + offset_mapping=encoding["offset_mapping"] if want_offsets else None,
475 )
476
440 - return TokensPrompt(prompt_token_ids=prompt_token_ids, **prompt)
441 -
477 def _detokenize_prompt(self, prompt: TokensPrompt) -> TokensPrompt:
478 tokenizer = self.get_tokenizer()
479 prompt["prompt"] = tokenizer.decode(prompt["prompt_token_ids"])
@@ -747,6 +782,11 @@ class BaseRenderer(ABC, Generic[_T]):
782 engine_input["prompt"] = prompt_text
783 if cache_salt := prompt.get("cache_salt"):
784 engine_input["cache_salt"] = cache_salt
785 + # Narrow the union — `prompt_token_offsets` is only on TokensInput.
786 + if engine_input["type"] == "token" and (
787 + (offsets := prompt.get("prompt_token_offsets")) is not None
788 + ):
789 + engine_input["prompt_token_offsets"] = offsets
790
791 return engine_input
792
@@ -805,6 +845,11 @@ class BaseRenderer(ABC, Generic[_T]):
845 engine_input["prompt"] = prompt_text
846 if cache_salt := prompt.get("cache_salt"):
847 engine_input["cache_salt"] = cache_salt
848 + # Narrow the union — `prompt_token_offsets` is only on TokensInput.
849 + if engine_input["type"] == "token" and (
850 + (offsets := prompt.get("prompt_token_offsets")) is not None
851 + ):
852 + engine_input["prompt_token_offsets"] = offsets
853
854 return engine_input
855
vllm/renderers/hf.py
+5
@@ -882,6 +882,11 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
882 self.tokenizer, config.model_config.renderer_num_workers + 1
883 )
884
885 + def _can_produce_offsets(self) -> bool:
886 + # HF tokenizers may be slow (use_fast=False); only fast tokenizers
887 + # expose offset_mapping.
888 + return self.tokenizer is not None and self.tokenizer.is_fast
889 +
890 def render_messages(
891 self,
892 messages: list[ChatCompletionMessageParam],
vllm/renderers/params.py
+5
@@ -167,6 +167,11 @@ class TokenizeParams:
167 add_special_tokens: bool = True
168 """Whether to add special tokens."""
169
170 + return_token_offsets: bool = False
171 + """If true, request char-level (start, end) offsets per token. Honored
172 + only for Fast (Rust-backed) tokenizers with text input and no multimodal
173 + data; otherwise silently ignored."""
174 +
175 needs_detokenization: bool = False
176 """
177 Whether the tokenized prompt needs to contain the original text.