[UT] fix device of test_outputs.py (#52237)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Aug 13, 2026 at 20:29 UTC
d18bb7b3908800f5e7f4c271de4eaf82d0f6f3e6
1 file changed
+11
-6
tests/v1/test_outputs.py
+11
-6
@@ -5,10 +5,13 @@ from unittest import TestCase
5
import numpy as np
6
import torch
7
8
+from vllm.platforms import current_platform
9
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
12
13
+DEVICE_TYPE = current_platform.device_type
14
+
15
16
def test_logprobs_tensors_cat():
17
first = LogprobsTensors(
@@ -71,9 +74,9 @@ def test_sampling_mask_tensors_from_logits():
74
[3.0, 4.0, float("-inf")],
75
[float("-inf"), 5.0, 6.0],
76
],
74
- device="cuda",
77
+ device=DEVICE_TYPE,
78
),
76
- num_sampled_tokens=torch.tensor([1, 0, 1], device="cuda"),
79
+ num_sampled_tokens=torch.tensor([1, 0, 1], device=DEVICE_TYPE),
80
)
81
82
result = tensors.tolists(np.array([1, 0, 1]))
@@ -85,9 +88,11 @@ def test_sampling_mask_tensors_from_logits():
88
89
def test_sampling_mask_matches_processed_top_k_top_p_support():
90
processed_logits = apply_top_k_top_p(
88
- logits=torch.tensor([[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device="cuda"),
89
- k=torch.tensor([3], device="cuda"),
90
- p=torch.tensor([0.9], device="cuda"),
91
+ logits=torch.tensor(
92
+ [[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE
93
+ ),
94
+ k=torch.tensor([3], device=DEVICE_TYPE),
95
+ p=torch.tensor([0.9], device=DEVICE_TYPE),
96
)
97
expected_token_ids = (
98
torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
@@ -96,7 +101,7 @@ def test_sampling_mask_matches_processed_top_k_top_p_support():
101
102
tensors = SamplingMaskTensors.from_logits(
103
processed_logits,
99
- num_sampled_tokens=torch.tensor([1], device="cuda"),
104
+ num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
105
)
106
result = tensors.tolists(np.array([1]))
107