A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[TPU][V1] Disable per-request seed/Generator (#16172)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Apr 10, 2025 at 23:05 UTC
3cc9af88ff76c16498cfa85579620e8db63fede9
4 files changed
+24
-18
tests/v1/tpu/test_sampler.py
+5
@@ -34,3 +34,8 @@ def test_sampler_different(model_name: str):
34
sampling_params = SamplingParams(temperature=0.1, min_p=0.8, max_tokens=64)
35
output2 = llm.generate(prompts, sampling_params)
36
assert output[0].outputs[0].text != output2[0].outputs[0].text
37
+
38
+ with pytest.raises(ValueError):
39
+ # Unsupported `seed` param.
40
+ sampling_params = SamplingParams(temperature=0.3, seed=42)
41
+ output2 = llm.generate(prompts, sampling_params)
vllm/platforms/tpu.py
+8
-5
@@ -7,7 +7,7 @@ import torch
7
import vllm.envs as envs
8
from vllm.inputs import PromptType
9
from vllm.logger import init_logger
10
-from vllm.sampling_params import SamplingParams
10
+from vllm.sampling_params import SamplingParams, SamplingType
11
12
from .interface import Platform, PlatformEnum, _Backend
13
@@ -145,7 +145,10 @@ class TpuPlatform(Platform):
145
params: Union[SamplingParams, PoolingParams],
146
) -> None:
147
"""Raises if this request is unsupported on this platform"""
148
- if isinstance(params,
149
- SamplingParams) and params.guided_decoding is not None:
150
- raise ValueError("Structured output is not supported on "
151
- f"{cls.device_name}.")
148
+ if isinstance(params, SamplingParams):
149
+ if params.guided_decoding is not None:
150
+ raise ValueError("Structured output is not supported on "
151
+ f"{cls.device_name}.")
152
+ if params.sampling_type == SamplingType.RANDOM_SEED:
153
+ raise ValueError(
154
+ "Torch XLA does not support per-request seed.")
vllm/v1/sample/tpu/metadata.py
+10
-6
@@ -33,10 +33,6 @@ class TPUSupportedSamplingMetadata:
33
# Greedy sampling flag for compiling single xla graph.
34
all_greedy: bool = True
35
36
- # Generator not supported by xla
37
- generators: dict[int,
38
- torch.Generator] = field(default_factory=lambda: dict())
39
-
36
# unsupported, you need to return an extra tensor of static size BxV
37
max_num_logprobs = None
38
@@ -57,6 +53,15 @@ class TPUSupportedSamplingMetadata:
53
allowed_token_ids_mask = None
54
bad_words_token_ids = None
55
56
+ # Generator not supported by xla
57
+ _generators: dict[int,
58
+ torch.Generator] = field(default_factory=lambda: dict())
59
+
60
+ @property
61
+ def generators(self) -> dict[int, torch.Generator]:
62
+ # Generator not supported by torch/xla. This field must be immutable.
63
+ return self._generators
64
+
65
@classmethod
66
def from_input_batch(
67
cls,
@@ -109,5 +114,4 @@ class TPUSupportedSamplingMetadata:
114
top_p=None, # input_batch.top_p[:padded_num_reqs],
115
top_k=None, # input_batch.top_k[:padded_num_reqs],
116
min_p=input_batch.min_p_cpu_tensor[:padded_num_reqs].to(
112
- xla_device),
113
- generators=input_batch.generators)
117
+ xla_device))
vllm/v1/worker/tpu_model_runner.py
+1
-7
@@ -23,7 +23,6 @@ from vllm.model_executor.model_loader import get_model
23
from vllm.multimodal import MULTIMODAL_REGISTRY
24
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
25
from vllm.multimodal.utils import group_mm_inputs_by_modality
26
-from vllm.sampling_params import SamplingType
26
from vllm.sequence import IntermediateTensors
27
from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
28
from vllm.v1.attention.backends.pallas import (PallasAttentionBackend,
@@ -267,11 +266,6 @@ class TPUModelRunner:
266
for new_req_data in scheduler_output.scheduled_new_reqs:
267
req_id = new_req_data.req_id
268
sampling_params = new_req_data.sampling_params
270
- if sampling_params.sampling_type == SamplingType.RANDOM_SEED:
271
- generator = torch.Generator(device=self.device)
272
- generator.manual_seed(sampling_params.seed)
273
- else:
274
- generator = None
269
270
self.requests[req_id] = CachedRequestState(
271
req_id=req_id,
@@ -280,7 +274,7 @@ class TPUModelRunner:
274
mm_inputs=new_req_data.mm_inputs,
275
mm_positions=new_req_data.mm_positions,
276
sampling_params=sampling_params,
283
- generator=generator,
277
+ generator=None,
278
block_ids=new_req_data.block_ids,
279
num_computed_tokens=new_req_data.num_computed_tokens,
280
output_token_ids=[],