Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
5de8d9f1
Unverified
Commit
5de8d9f1
authored
Jul 12, 2025
by
Maximilien de Bayser
Committed by
GitHub
Jul 12, 2025
Browse files
Remove extra tensor on CPU (#20693)
Signed-off-by:
Max de Bayser
<
mbayser@br.ibm.com
>
parent
c1c8ca57
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
13 additions
and
5 deletions
+13
-5
vllm/v1/sample/logits_processor.py
vllm/v1/sample/logits_processor.py
+13
-5
No files found.
vllm/v1/sample/logits_processor.py
View file @
5de8d9f1
...
...
@@ -234,10 +234,16 @@ class MinPLogitsProcessor(LogitsProcessor):
device
=
"cpu"
,
pin_memory
=
pin_memory
)
self
.
min_p_cpu
=
self
.
min_p_cpu_tensor
.
numpy
()
# Pre-allocated device tensor
self
.
min_p_device
:
torch
.
Tensor
=
torch
.
empty
((
max_num_reqs
,
),
dtype
=
torch
.
float32
,
device
=
device
)
self
.
use_double_tensor
=
torch
.
device
(
"cpu"
)
!=
torch
.
device
(
device
)
if
self
.
use_double_tensor
:
# Pre-allocated device tensor
self
.
min_p_device
:
torch
.
Tensor
=
torch
.
empty
((
max_num_reqs
,
),
dtype
=
torch
.
float32
,
device
=
device
)
else
:
self
.
min_p_device
=
self
.
min_p_cpu_tensor
# Current slice of the device tensor
self
.
min_p
:
torch
.
Tensor
=
self
.
min_p_device
[:
0
]
...
...
@@ -284,7 +290,9 @@ class MinPLogitsProcessor(LogitsProcessor):
size
=
batch_update
.
batch_size
if
self
.
min_p_count
and
(
needs_update
or
self
.
min_p
.
shape
[
0
]
!=
size
):
self
.
min_p
=
self
.
min_p_device
[:
size
]
self
.
min_p
.
copy_
(
self
.
min_p_cpu_tensor
[:
size
],
non_blocking
=
True
)
if
self
.
use_double_tensor
:
self
.
min_p
.
copy_
(
self
.
min_p_cpu_tensor
[:
size
],
non_blocking
=
True
)
self
.
min_p
.
unsqueeze_
(
1
)
def
apply
(
self
,
logits
:
torch
.
Tensor
)
->
torch
.
Tensor
:
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment