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, 13:30 UTC
[TPU][V1] Enable Top-P (#16843)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Nicolò Lucchesi committed
Apr 22, 2025 at 02:46 UTC
fa3bba2a538de76c630f75de160bfb43e4e1cd4b
2 files changed
+9
-10
tests/v1/tpu/test_sampler.py
+4
-3
@@ -42,7 +42,7 @@ def test_sampler_different(model_name: str):
42
sampling_params = SamplingParams(temperature=0.3, seed=42)
43
output2 = llm.generate(prompts, sampling_params)
44
45
- # Batch-case with TopK
45
+ # Batch-case with TopK/P
46
for B in [4, 16]:
47
p = prompts * B
48
sampling_params = [
@@ -51,9 +51,10 @@ def test_sampler_different(model_name: str):
51
min_p=0.8,
52
max_tokens=64,
53
# Vary number of ks
54
- top_k=random.randint(4, 12)) for _ in range(B)
54
+ top_k=random.randint(4, 12),
55
+ top_p=random.random()) for _ in range(B)
56
]
56
- # Make sure first two reqs have the same K
57
+ # Make sure first two reqs have the same K/P
58
sampling_params[0] = sampling_params[1]
59
output = llm.generate(p, sampling_params)
60
assert output[0].outputs[0].text == output[1].outputs[0].text
vllm/v1/sample/tpu/metadata.py
+5
-7
@@ -11,7 +11,7 @@ DEFAULT_SAMPLING_PARAMS = dict(
11
min_p=0.0,
12
# strictly disabled for now
13
top_k=0,
14
- # top_p=0.0,
14
+ top_p=1.0,
15
# frequency_penalties=0.0,
16
# presence_penalties=0.0,
17
# repetition_penalties=0.0,
@@ -26,11 +26,9 @@ class TPUSupportedSamplingMetadata:
26
temperature: torch.Tensor = None
27
28
min_p: torch.Tensor = None
29
- # Still too slow on forward_native!
29
top_k: torch.Tensor = None
30
top_p: torch.Tensor = None
31
33
- # Greedy sampling flag for compiling single xla graph.
32
all_greedy: bool = True
33
34
# unsupported, you need to return an extra tensor of static size BxV
@@ -103,9 +101,8 @@ class TPUSupportedSamplingMetadata:
101
DEFAULT_SAMPLING_PARAMS["min_p"])
102
fill_slice(input_batch.top_k_cpu_tensor,
103
DEFAULT_SAMPLING_PARAMS["top_k"])
106
- # TODO Temporarily disabled until sampling options are enabled
107
- # fill_slice(input_batch.top_p_cpu_tensor,
108
- # DEFAULT_SAMPLING_PARAMS["top_p"])
104
+ fill_slice(input_batch.top_p_cpu_tensor,
105
+ DEFAULT_SAMPLING_PARAMS["top_p"])
106
107
# Slice persistent device tensors to a fixed pre-compiled padded shape.
108
return cls(
@@ -113,7 +110,8 @@ class TPUSupportedSamplingMetadata:
110
to(xla_device),
111
all_greedy=input_batch.all_greedy,
112
# TODO enable more and avoid returning None values
116
- top_p=None, # input_batch.top_p[:padded_num_reqs],
113
+ top_p=input_batch.top_p_cpu_tensor[:padded_num_reqs].to(
114
+ xla_device),
115
top_k=input_batch.top_k_cpu_tensor[:padded_num_reqs].to(
116
xla_device),
117
min_p=input_batch.min_p_cpu_tensor[:padded_num_reqs].to(