@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 32e0c0bfa2
Fork 0

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 5, 14:45 UTC

refactor hard coded device string in test files under tests/v1 and tests/lora (#37566)

Signed-off-by: Liao, Wei <wei.liao@intel.com>

wliao2 committed Apr 2, 2026 at 20:21 UTC 32e0c0bfa210cc6e26bebb967c24aa402919d0cb
28 files changed +239 -146
tests/lora/test_fused_moe_lora_kernel.py
+1 -1
@@ -637,7 +637,7 @@ def use_fused_moe_lora_kernel_tensor_parallel(
637
638 set_random_seed(seed)
639
640 - device = torch.device(f"cuda:{local_rank}")
640 + device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
641 torch.accelerator.set_device_index(device)
642 torch.set_default_device(device)
643 torch.set_default_dtype(dtype)
tests/lora/test_layers.py
+6 -2
@@ -60,8 +60,12 @@ pytestmark = pytest.mark.skipif(
60 reason="Backend not supported",
61 )
62
63 +DEVICE_TYPE = current_platform.device_type
64 DEVICES = (
64 - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
65 + [
66 + f"{DEVICE_TYPE}:{i}"
67 + for i in range(1 if torch.accelerator.device_count() == 1 else 2)
68 + ]
69 if current_platform.is_cuda_alike()
70 else ["cpu"]
71 )
@@ -196,7 +200,7 @@ def create_random_inputs(
200 input_size: tuple[int, ...],
201 input_range: tuple[float, float],
202 input_type: torch.dtype = torch.int,
199 - device: torch.device = "cuda",
203 + device: torch.device = DEVICE_TYPE,
204 ) -> tuple[list[torch.Tensor], list[int], list[int]]:
205 """Creates random inputs.
206
tests/lora/test_lora_manager.py
+2 -2
@@ -35,9 +35,9 @@ EMBEDDING_MODULES = {
35 "lm_head": "output_embeddings",
36 }
37
38 -
38 +DEVICE_TYPE = current_platform.device_type
39 DEVICES = (
40 - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]
40 + [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
41 if current_platform.is_cuda_alike()
42 else ["cpu"]
43 )
tests/lora/test_moe_lora_align_sum.py
+16 -6
@@ -6,6 +6,9 @@ import pytest
6 import torch
7
8 from vllm import _custom_ops as ops
9 +from vllm.platforms import current_platform
10 +
11 +DEVICE_TYPE = current_platform.device_type
12
13
14 def round_up(x, base):
@@ -27,7 +30,7 @@ def sample_data(num_experts, max_loras, num_tokens, topk_num):
30 topk_ids[i, j] = pool[j]
31 token_lora_mapping[i] = random.randint(0, max_loras - 1)
32
30 - return topk_ids.to("cuda"), token_lora_mapping.to("cuda")
33 + return topk_ids.to(DEVICE_TYPE), token_lora_mapping.to(DEVICE_TYPE)
34
35
36 @pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) # 81920
@@ -56,14 +59,21 @@ def test_moe_lora_align_block_size(
59 (max_loras * max_num_tokens_padded,),
60 topk_ids.numel(),
61 dtype=torch.int32,
59 - device="cuda",
62 + device=DEVICE_TYPE,
63 )
64 expert_ids = torch.full(
62 - (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
65 + (max_loras * max_num_m_blocks,),
66 + num_experts,
67 + dtype=torch.int32,
68 + device=DEVICE_TYPE,
69 + )
70 + num_tokens_post_pad = torch.zeros(
71 + (max_loras,), dtype=torch.int32, device=DEVICE_TYPE
72 + )
73 + adapter_enabled = torch.ones(
74 + (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
75 )
64 - num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
65 - adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
66 - lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda")
76 + lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device=DEVICE_TYPE)
77
78 # call kernel
79 ops.moe_lora_align_block_size(
tests/lora/test_punica_ops.py
+10 -3
@@ -9,10 +9,13 @@ import vllm.lora.ops.torch_ops as torch_ops
9 import vllm.lora.ops.triton_ops as triton_ops
10 from vllm.lora.ops.triton_ops import LoRAKernelMeta
11 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
12 +from vllm.platforms import current_platform
13 from vllm.utils.torch_utils import set_random_seed
14
15 from .utils import PunicaTensors, assert_close, generate_data_for_nslices
16
17 +DEVICE_TYPE = current_platform.device_type
18 +
19
20 @pytest.fixture(autouse=True)
21 def reset_device(reset_default_device):
@@ -146,7 +149,9 @@ def check_lora_shrink_kernel(
149
150 # Setup metadata information for the LoRA kernel.
151 lora_meta = LoRAKernelMeta.make(
149 - max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
152 + max_loras=num_loras,
153 + max_num_tokens=token_nums,
154 + device=DEVICE_TYPE,
155 )
156 lora_meta.prepare_tensors(data.token_lora_mapping)
157
@@ -219,7 +224,9 @@ def check_lora_expand_kernel(
224
225 # Setup metadata information for the LoRA kernel.
226 lora_meta = LoRAKernelMeta.make(
222 - max_loras=num_loras, max_num_tokens=token_nums, device="cuda"
227 + max_loras=num_loras,
228 + max_num_tokens=token_nums,
229 + device=DEVICE_TYPE,
230 )
231 lora_meta.prepare_tensors(data.token_lora_mapping)
232
@@ -367,7 +374,7 @@ test_params = {
374 }
375
376 DTYPES = [torch.float16, torch.bfloat16]
370 -DEVICES = [f"cuda:{0}"]
377 +DEVICES = [f"{DEVICE_TYPE}:{0}"]
378 SEED = [0]
379
380
tests/lora/test_punica_ops_fp8.py
+3 -1
@@ -28,9 +28,11 @@ from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import (
28 _SHRINK_LORA_SCALE_PTR_DICT,
29 )
30 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
31 +from vllm.platforms import current_platform
32 from vllm.utils.torch_utils import set_random_seed
33
33 -DEVICES = [f"cuda:{0}"]
34 +DEVICE_TYPE = current_platform.device_type
35 +DEVICES = [f"{DEVICE_TYPE}:{0}"]
36 SEED = [0]
37
38 _dict_lock = Lock()
tests/lora/test_worker.py
+4 -1
@@ -19,11 +19,14 @@ from vllm.config.load import LoadConfig
19 from vllm.config.lora import LoRAConfig
20 from vllm.lora.model_manager import LoRAMapping
21 from vllm.lora.request import LoRARequest
22 +from vllm.platforms import current_platform
23 from vllm.v1.worker.gpu_worker import Worker
24
25 MODEL_PATH = "Qwen/Qwen3-0.6B"
26 NUM_LORAS = 16
27
28 +DEVICE_TYPE = current_platform.device_type
29 +
30
31 @patch.dict(os.environ, {"RANK": "0"})
32 def test_worker_apply_lora(qwen3_lora_files):
@@ -61,7 +64,7 @@ def test_worker_apply_lora(qwen3_lora_files):
64 max_num_seqs=32,
65 max_num_partial_prefills=32,
66 ),
64 - device_config=DeviceConfig("cuda"),
67 + device_config=DeviceConfig(DEVICE_TYPE),
68 cache_config=CacheConfig(
69 block_size=16,
70 cache_dtype="auto",
tests/lora/utils.py
+6 -3
@@ -9,10 +9,13 @@ import torch
9 from safetensors.torch import save_file
10
11 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
12 +from vllm.platforms import current_platform
13 +
14 +DEVICE_TYPE = current_platform.device_type
15
16
17 class DummyLoRAManager:
15 - def __init__(self, device: torch.device = "cuda:0"):
18 + def __init__(self, device: torch.device = f"{DEVICE_TYPE}:0"):
19 super().__init__()
20 self._loras: dict[str, LoRALayerWeights] = {}
21 self._device = device
@@ -57,8 +60,8 @@ class DummyLoRAManager:
60 module_name,
61 rank=rank,
62 lora_alpha=1,
60 - lora_a=torch.rand([rank, input_dim], device="cuda"),
61 - lora_b=torch.rand([output_dim, input_dim], device="cuda"),
63 + lora_a=torch.rand([rank, input_dim], device=DEVICE_TYPE),
64 + lora_b=torch.rand([output_dim, input_dim], device=DEVICE_TYPE),
65 embeddings_tensor=embeddings_tensor,
66 )
67 self.set_module_lora(module_name, lora)
tests/v1/attention/test_attention_backends.py
+3 -1
@@ -40,6 +40,8 @@ BACKENDS_TO_TEST = [
40 "FLEX_ATTENTION_SLOW",
41 ]
42
43 +DEVICE_TYPE = current_platform.device_type
44 +
45 # Remove flashinfer from the list if it's not available
46 try:
47 import flashinfer # noqa: F401
@@ -366,7 +368,7 @@ def _test_backend_correctness(
368 num_gpu_blocks=8192,
369 hf_config_override=hf_config_override,
370 )
369 - device = torch.device("cuda:0")
371 + device = torch.device(f"{DEVICE_TYPE}:0")
372
373 kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
374
tests/v1/attention/test_chunked_local_attention.py
+4 -1
@@ -7,6 +7,7 @@ import pytest
7 import torch
8
9 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10 +from vllm.platforms import current_platform
11 from vllm.v1.attention.backends.utils import make_local_attention_virtual_batches
12
13
@@ -22,6 +23,8 @@ class LocalAttentionTestData:
23 expected_local_block_table: list[list[int]]
24
25
26 +DEVICE_TYPE = current_platform.device_type
27 +
28 test_data_list = [
29 # Same as example in docstring of make_local_attention_virtual_batches
30 # except block table has 9 columns instead of 10
@@ -151,7 +154,7 @@ test_data_list = [
154
155 @pytest.mark.parametrize("test_data", test_data_list)
156 def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):
154 - device = torch.device("cuda:0")
157 + device = torch.device(f"{DEVICE_TYPE}:0")
158 batch_spec = test_data.batch_spec
159 attn_chunk_size = test_data.attn_chunk_size
160 block_size = test_data.block_size
tests/v1/attention/test_mla_backends.py
+3 -1
@@ -42,6 +42,8 @@ BACKENDS_TO_TEST = [
42 AttentionBackendEnum.TRITON_MLA,
43 ]
44
45 +DEVICE_TYPE = current_platform.device_type
46 +
47 # Remove sm100 backends from the list if not using sm100
48 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
49 BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
@@ -763,7 +765,7 @@ def test_backend_correctness(
765 method="ngram", num_speculative_tokens=query_len - 1
766 )
767
766 - device = torch.device("cuda:0")
768 + device = torch.device(f"{DEVICE_TYPE}:0")
769
770 # 1. Setup
771 batch_size = batch_spec.batch_size
tests/v1/attention/test_sparse_mla_backends.py
+6 -4
@@ -64,6 +64,8 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
64 seq_lens=[256] * 2, query_lens=[256] * 2
65 )
66
67 +DEVICE_TYPE = current_platform.device_type
68 +
69
70 def _float_to_e8m0_truncate(f: float) -> float:
71 """Simulate SM100's float -> e8m0 -> bf16 scale conversion.
@@ -222,7 +224,7 @@ def test_sparse_backend_decode_correctness(
224 batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
225 use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
226
225 - device = torch.device("cuda")
227 + device = torch.device(DEVICE_TYPE)
228 dtype = torch.bfloat16
229
230 # Model hyper-parameters (kept intentionally small for the unit test)
@@ -586,7 +588,7 @@ def _triton_convert_reference_impl(
588 def test_triton_convert_req_index_to_global_index_decode_only(
589 block_size, num_topk_tokens
590 ):
589 - device = torch.device("cuda")
591 + device = torch.device(DEVICE_TYPE)
592 num_tokens = 8
593 num_requests = 4
594 max_blocks_per_req = 10
@@ -639,7 +641,7 @@ def test_triton_convert_req_index_to_global_index_decode_only(
641 reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
642 )
643 def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size):
642 - device = torch.device("cuda")
644 + device = torch.device(DEVICE_TYPE)
645 num_requests = 4
646 max_blocks_per_req = 8
647 num_topk_tokens = 128
@@ -794,7 +796,7 @@ def test_split_indexer_prefill_chunks_single_request_overflow():
796
797 def test_triton_convert_returns_valid_counts():
798 """Test that return_valid_counts correctly counts non-negative indices."""
797 - device = torch.device("cuda")
799 + device = torch.device(DEVICE_TYPE)
800 num_tokens = 8
801 num_requests = 2
802 max_blocks_per_req = 10
tests/v1/attention/test_trtllm_attention_integration.py
+2 -1
@@ -55,6 +55,7 @@ class MockAttentionLayer:
55 MODEL = "Qwen/Qwen2.5-0.5B"
56 BLOCK_SIZE = 16
57 NUM_GPU_BLOCKS = 8192
58 +DEVICE_TYPE = current_platform.device_type
59
60 BATCH_SPECS = {
61 "decode_only": BatchSpec(
@@ -172,7 +173,7 @@ def _run_trtllm_integration(batch_spec):
173 """Run TRTLLM attention through the full FlashInfer pipeline
174 and compare against an SDPA reference."""
175 set_random_seed(42)
175 - device = torch.device("cuda:0")
176 + device = torch.device(f"{DEVICE_TYPE}:0")
177
178 vllm_config = create_vllm_config(
179 model_name=MODEL,
tests/v1/cudagraph/test_cudagraph_dispatch.py
+11 -9
@@ -23,6 +23,8 @@ from vllm.forward_context import BatchDescriptor, set_forward_context
23 from vllm.platforms import current_platform
24 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
25
26 +DEVICE_TYPE = current_platform.device_type
27 +
28
29 # Helper MLP for testing
30 class SimpleMLP(nn.Module):
@@ -269,9 +271,9 @@ class TestCudagraphDispatcher:
271 class TestCUDAGraphWrapper:
272 def setup_method(self):
273 self.vllm_config = _create_vllm_config(CompilationConfig())
272 - self.model = SimpleMLP().to("cuda")
273 - self.persistent_input_buffer = torch.zeros(1, 10, device="cuda")
274 - self.input_tensor = torch.randn(1, 10, device="cuda")
274 + self.model = SimpleMLP().to(DEVICE_TYPE)
275 + self.persistent_input_buffer = torch.zeros(1, 10, device=DEVICE_TYPE)
276 + self.input_tensor = torch.randn(1, 10, device=DEVICE_TYPE)
277
278 def test_capture_and_replay(self):
279 wrapper = CUDAGraphWrapper(
@@ -428,10 +430,10 @@ class TestCudagraphIntegration:
430
431 @create_new_process_for_each_test("spawn")
432 def test_capture_replay_bypass_logic(self):
431 - model = SimpleMLP().to("cuda")
433 + model = SimpleMLP().to(DEVICE_TYPE)
434 full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
435 max_bs = 16
434 - persistent_input_buffer = torch.zeros(max_bs, 10, device="cuda")
436 + persistent_input_buffer = torch.zeros(max_bs, 10, device=DEVICE_TYPE)
437 input_1 = persistent_input_buffer[:1]
438 input_2 = persistent_input_buffer[:2]
439 input_3 = persistent_input_buffer[:3]
@@ -486,17 +488,17 @@ class TestCudagraphIntegration:
488 @create_new_process_for_each_test("spawn")
489 def test_nested_wrappers(self):
490 """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
489 - model = SimpleMLP().to("cuda")
491 + model = SimpleMLP().to(DEVICE_TYPE)
492 full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
491 - input_1 = torch.randn(1, 10, device="cuda")
493 + input_1 = torch.randn(1, 10, device=DEVICE_TYPE)
494
495 # Setup: Inner model is wrapped with PIECEWISE, outer with FULL
494 - inner_model = SimpleMLP().to("cuda")
496 + inner_model = SimpleMLP().to(DEVICE_TYPE)
497 piecewise_wrapper = CUDAGraphWrapper(
498 inner_model, self.vllm_config, CUDAGraphMode.PIECEWISE
499 )
500 inner_model.forward = MagicMock(wraps=inner_model.forward)
499 - outer_model = SimpleMLP().to("cuda")
501 + outer_model = SimpleMLP().to(DEVICE_TYPE)
502 # When outer model is called, it calls the piecewise_wrapper
503 outer_model.forward = MagicMock(
504 wraps=outer_model.forward, side_effect=piecewise_wrapper
tests/v1/determinism/test_rms_norm_batch_invariant.py
+10 -7
@@ -13,6 +13,9 @@ from utils import skip_unsupported
13
14 from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
15 from vllm.model_executor.layers.layernorm import RMSNorm
16 +from vllm.platforms import current_platform
17 +
18 +DEVICE_TYPE = current_platform.device_type
19
20
21 @skip_unsupported
@@ -34,7 +37,7 @@ def test_rms_norm_batch_invariant_vs_standard(
37 equivalent results to the standard CUDA implementation across various
38 configurations.
39 """
37 - device = torch.device("cuda")
40 + device = torch.device(DEVICE_TYPE)
41
42 # Create test input and weight
43 torch.manual_seed(42)
@@ -81,7 +84,7 @@ def test_rms_norm_3d_input(
84 Ensures that the batch-invariant RMS norm correctly handles multi-dimensional
85 inputs that are common in transformer models.
86 """
84 - device = torch.device("cuda")
87 + device = torch.device(DEVICE_TYPE)
88 dtype = torch.bfloat16
89 eps = 1e-6
90
@@ -120,7 +123,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):
123 Ensures that both implementations handle edge cases like very small or large
124 values without producing NaN or Inf.
125 """
123 - device = torch.device("cuda")
126 + device = torch.device(DEVICE_TYPE)
127 dtype = torch.float16
128 eps = 1e-6
129 hidden_size = 2048
@@ -179,7 +182,7 @@ def test_rms_norm_formula(default_vllm_config):
182
183 Verifies: output = input / sqrt(mean(input^2) + eps) * weight
184 """
182 - device = torch.device("cuda")
185 + device = torch.device(DEVICE_TYPE)
186 dtype = torch.float32 # Use float32 for higher precision in formula check
187 eps = 1e-6
188 hidden_size = 1024
@@ -214,7 +217,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
217 The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it
218 correctly handles hidden sizes both smaller and larger than the block size.
219 """
217 - device = torch.device("cuda")
220 + device = torch.device(DEVICE_TYPE)
221 dtype = torch.bfloat16
222 eps = 1e-6
223 batch_size = 16
@@ -251,7 +254,7 @@ def test_rms_norm_determinism(default_vllm_config):
254 Runs the same input through the kernel multiple times and verifies
255 identical outputs.
256 """
254 - device = torch.device("cuda")
257 + device = torch.device(DEVICE_TYPE)
258 dtype = torch.bfloat16
259 eps = 1e-6
260 hidden_size = 4096
@@ -283,7 +286,7 @@ if __name__ == "__main__":
286 # Run a quick smoke test
287 print("Running quick smoke test of RMS norm implementations...")
288
286 - device = torch.device("cuda")
289 + device = torch.device(DEVICE_TYPE)
290 batch_size = 8
291 hidden_size = 4096
292 dtype = torch.bfloat16
tests/v1/e2e/general/test_mamba_prefix_cache.py
+15 -5
@@ -16,6 +16,7 @@ from vllm import LLM, SamplingParams, TokensPrompt
16 from vllm.config import CacheConfig
17 from vllm.distributed import cleanup_dist_env_and_memory
18 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
19 +from vllm.platforms import current_platform
20 from vllm.sequence import IntermediateTensors
21 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
22 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
@@ -48,6 +49,7 @@ num_accepted_tokens = 1
49 prompt_token_ids: list[int] = []
50 MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
51 BLOCK_SIZE = 560
52 +DEVICE_TYPE = current_platform.device_type
53 NUM_HIDDEN_LAYERS = 1
54 cur_step_action_idx = 0
55 cur_step_action: StepAction | None = None
@@ -71,7 +73,7 @@ def get_fake_sample_fn() -> SamplerOutput:
73 return SamplerOutput(
74 sampled_token_ids=torch.tensor(
75 [[prompt_token_ids[first_token_id_index]]],
74 - device="cuda",
76 + device=DEVICE_TYPE,
77 dtype=torch.int32,
78 ),
79 logprobs_tensors=None,
@@ -83,7 +85,9 @@ def get_fake_sample_fn() -> SamplerOutput:
85 sampled_token_ids = accepted_tokens
86 return SamplerOutput(
87 sampled_token_ids=torch.tensor(
86 - [sampled_token_ids], device="cuda", dtype=torch.int32
88 + [sampled_token_ids],
89 + device=DEVICE_TYPE,
90 + dtype=torch.int32,
91 ),
92 logprobs_tensors=None,
93 )
@@ -128,17 +132,23 @@ def get_fake_propose_draft_token_ids_fn():
132 - 1
133 + num_accepted_tokens
134 ],
131 - device="cuda",
135 + device=DEVICE_TYPE,
136 dtype=torch.int32,
137 )
138
139 valid_sampled_tokens_count = torch.tensor(
136 - [num_accepted_tokens], device="cuda", dtype=torch.int32
140 + [num_accepted_tokens],
141 + device=DEVICE_TYPE,
142 + dtype=torch.int32,
143 )
144
145 self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
146
141 - return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32)
147 + return torch.tensor(
148 + proposed_draft_token_ids,
149 + device=DEVICE_TYPE,
150 + dtype=torch.int32,
151 + )
152
153 return fake_propose_draft_token_ids_fn
154
tests/v1/kv_offload/test_cpu_gpu.py
+4 -2
@@ -6,6 +6,7 @@ import time
6 import pytest
7 import torch
8
9 +from vllm.platforms import current_platform
10 from vllm.utils.torch_utils import set_random_seed
11 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
12 from vllm.v1.kv_offload.spec import (
@@ -21,7 +22,8 @@ GPU_PAGE_SIZES = [512, 1024]
22 BLOCK_SIZE_FACTORS = [1, 3]
23 NUM_TENSORS = [4]
24 SEEDS = [0]
24 -CUDA_DEVICES = ["cuda:0"]
25 +DEVICE_TYPE = current_platform.device_type
26 +DEVICES = [f"{DEVICE_TYPE}:0"]
27 NUM_MAPPINGS = [3]
28
29
@@ -33,7 +35,7 @@ NUM_MAPPINGS = [3]
35 @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
36 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
37 @pytest.mark.parametrize("seed", SEEDS)
36 -@pytest.mark.parametrize("device", CUDA_DEVICES)
38 +@pytest.mark.parametrize("device", DEVICES)
39 @torch.inference_mode()
40 def test_transfer(
41 default_vllm_config,
tests/v1/logits_processors/test_correctness.py
+4 -3
@@ -39,8 +39,9 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available()
39 MAX_NUM_REQS = 256
40 VOCAB_SIZE = 1024
41 NUM_OUTPUT_TOKENS = 20
42 -CUDA_DEVICES = [
43 - f"{current_platform.device_type}:{i}"
42 +DEVICE_TYPE = current_platform.device_type
43 +DEVICES = [
44 + f"{DEVICE_TYPE}:{i}"
45 for i in range(1 if current_platform.device_count() == 1 else 2)
46 ]
47 MAX_NUM_PROMPT_TOKENS = 64
@@ -801,7 +802,7 @@ def _assert_valid(
802
803
804 @create_new_process_for_each_test()
804 -@pytest.mark.parametrize("device", CUDA_DEVICES)
805 +@pytest.mark.parametrize("device", DEVICES)
806 @pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])
807 @pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())
808 def test_logitsprocs(
tests/v1/sample/test_rejection_sampler.py
+60 -30
@@ -19,7 +19,7 @@ from vllm.v1.sample.rejection_sampler import (
19 from vllm.v1.sample.sampler import Sampler, SamplerOutput
20 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
21
22 -DEVICE = current_platform.device_type
22 +DEVICE_TYPE = current_platform.device_type
23
24
25 @pytest.fixture
@@ -57,7 +57,7 @@ def create_logits_tensor(
57 will produce desired token ids on argmax"""
58 token_ids = [tokens[:-1] for tokens in output_token_ids]
59 num_total_tokens = sum(len(tokens) for tokens in token_ids)
60 - logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE)
60 + logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE_TYPE)
61 start_loc = 0
62 for tokens in token_ids:
63 for j, token_id in enumerate(tokens):
@@ -99,9 +99,9 @@ def create_sampling_metadata(
99 assert output_token_ids
100 assert len(output_token_ids) > 0
101
102 - frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE)
103 - presence_penalties = torch.tensor(presence_penalties, device=DEVICE)
104 - repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE)
102 + frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE_TYPE)
103 + presence_penalties = torch.tensor(presence_penalties, device=DEVICE_TYPE)
104 + repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE_TYPE)
105 else:
106 no_penalties = True
107 frequency_penalties = torch.tensor([])
@@ -320,14 +320,27 @@ def test_deterministic_when_seeded(
320 n_rep: int,
321 ):
322 num_tokens = batch_size * k
323 - draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
323 + draft_probs = torch.rand(
324 + num_tokens,
325 + vocab_size,
326 + dtype=torch.float32,
327 + device=DEVICE_TYPE,
328 + )
329 draft_probs = F.softmax(draft_probs, dim=-1)
330 target_logits = torch.rand_like(draft_probs)
331 bonus_token_ids = torch.randint(
327 - low=0, high=vocab_size, size=(batch_size, 1), dtype=torch.int64, device=DEVICE
332 + low=0,
333 + high=vocab_size,
334 + size=(batch_size, 1),
335 + dtype=torch.int64,
336 + device=DEVICE_TYPE,
337 )
338 draft_token_ids = torch.randint(
330 - low=0, high=vocab_size, size=(batch_size, k), dtype=torch.int64, device=DEVICE
339 + low=0,
340 + high=vocab_size,
341 + size=(batch_size, k),
342 + dtype=torch.int64,
343 + device=DEVICE_TYPE,
344 )
345
346 seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded
@@ -335,12 +348,12 @@ def test_deterministic_when_seeded(
348 results = []
349 for _ in range(n_rep):
350 seeded_seqs = {
338 - i: torch.Generator(device=DEVICE).manual_seed(i)
351 + i: torch.Generator(device=DEVICE_TYPE).manual_seed(i)
352 for i in range(batch_size)
353 if seeded_mask[i]
354 }
355
343 - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
356 + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
357 sampling_metadata = create_sampling_metadata(
358 all_greedy=False, temperature=temperature, generators=seeded_seqs
359 )
@@ -387,7 +400,7 @@ def test_rejection_sampling_approximates_target_distribution():
400 much more than the distance improvement between the observed
401 distribution and the random distribution.
402 """
390 - torch.set_default_device(DEVICE)
403 + torch.set_default_device(DEVICE_TYPE)
404 vocab_size = 10
405 k = 2
406 num_reference_probs = 100
@@ -410,7 +423,7 @@ def test_rejection_sampling_approximates_target_distribution():
423 rej_sample_probs = estimate_rejection_sampling_pdf(
424 draft_probs, target_logits, k, vocab_size, num_samples
425 )
413 - rej_sample_probs = rej_sample_probs.to(DEVICE)
426 + rej_sample_probs = rej_sample_probs.to(DEVICE_TYPE)
427
428 # Average distance from reference probs.
429 reference_vs_rejsample_dist = (
@@ -491,11 +504,11 @@ def estimate_rejection_sampling_pdf(
504 draft_probs = draft_probs.view(num_tokens, vocab_size)
505
506 # Bonus tokens not used but required.
494 - bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE).repeat(
507 + bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE_TYPE).repeat(
508 num_samples, 1
509 )
510
498 - temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE)
511 + temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE_TYPE)
512 sampling_metadata = create_sampling_metadata(
513 all_greedy=False, temperature=temperature
514 )
@@ -600,7 +613,7 @@ def _test_masked_logits(
613
614 # Create random draft probabilities.
615 draft_probs = torch.rand(
603 - (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE
616 + (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE_TYPE
617 )
618 draft_probs = F.softmax(draft_probs, dim=-1)
619
@@ -610,7 +623,11 @@ def _test_masked_logits(
623 draft_token_ids = draft_token_ids.tolist()
624
625 # Bonus tokens not used but required
613 - bonus_token_ids = torch.zeros((batch_size, 1), dtype=torch.int64, device=DEVICE)
626 + bonus_token_ids = torch.zeros(
627 + (batch_size, 1),
628 + dtype=torch.int64,
629 + device=DEVICE_TYPE,
630 + )
631
632 # Create spec decode metadata
633 spec_decode_metadata = create_spec_decode_metadata(draft_token_ids, target_logits)
@@ -645,12 +662,13 @@ def test_top_k(rejection_sampler, top_k):
662
663 # Randomly create top-k indices.
664 top_k_indices = [
648 - torch.randperm(vocab_size, device=DEVICE)[:top_k] for _ in range(num_tokens)
665 + torch.randperm(vocab_size, device=DEVICE_TYPE)[:top_k]
666 + for _ in range(num_tokens)
667 ]
668 top_k_indices = torch.stack(top_k_indices)
669
670 # Create logits with the uniform distribution.
653 - target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE)
671 + target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE_TYPE)
672
673 # Increment the logits for top-k indices, a little bit more than the other
674 # ones. If the masking is effective, the non-topk indices will never be
@@ -659,11 +677,11 @@ def test_top_k(rejection_sampler, top_k):
677 target_logits[i, top_k_indices[i]] += 0.1
678
679 # Create sampling metadata
662 - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
680 + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
681 sampling_metadata = create_sampling_metadata(
682 all_greedy=False,
683 temperature=temperature,
666 - top_k=torch.tensor([top_k] * batch_size, device=DEVICE, dtype=torch.int64),
684 + top_k=torch.tensor([top_k] * batch_size, device=DEVICE_TYPE, dtype=torch.int64),
685 )
686
687 _test_masked_logits(
@@ -686,8 +704,8 @@ def test_top_p(rejection_sampler, top_p):
704 num_tokens = batch_size * num_draft_tokens
705
706 # Create logits with the uniform distribution.
689 - target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE)
690 - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
707 + target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE_TYPE)
708 + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
709 rescaled_logits = target_logits / temperature
710
711 logits_sort, logits_idx = rescaled_logits.sort(dim=-1, descending=False)
@@ -706,7 +724,11 @@ def test_top_p(rejection_sampler, top_p):
724 sampling_metadata = create_sampling_metadata(
725 all_greedy=False,
726 temperature=temperature,
709 - top_p=torch.tensor([top_p] * batch_size, device=DEVICE, dtype=torch.float32),
727 + top_p=torch.tensor(
728 + [top_p] * batch_size,
729 + device=DEVICE_TYPE,
730 + dtype=torch.float32,
731 + ),
732 )
733
734 _test_masked_logits(
@@ -732,7 +754,10 @@ def test_frequency_penalties(rejection_sampler):
754 all_greedy=True,
755 output_token_ids=[[2], [3], [4]],
756 spec_token_ids=spec_tokens,
735 - prompt_token_ids=torch.tensor([[5, 6, 7], [6, 7, 8], [7, 8, 9]], device=DEVICE),
757 + prompt_token_ids=torch.tensor(
758 + [[5, 6, 7], [6, 7, 8], [7, 8, 9]],
759 + device=DEVICE_TYPE,
760 + ),
761 frequency_penalties=[1.5, 1.5, 0.7],
762 presence_penalties=[0.0] * num_requests,
763 repetition_penalties=[1.0] * num_requests,
@@ -858,21 +883,26 @@ def test_sample_recovered_tokens(
883 num_tokens = batch_size * max_spec_len
884
885 # Create random draft probabilities.
861 - draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
886 + draft_probs = torch.rand(
887 + num_tokens,
888 + vocab_size,
889 + dtype=torch.float32,
890 + device=DEVICE_TYPE,
891 + )
892 draft_probs = F.softmax(draft_probs, dim=-1)
893
894 # Create random target probabilities.
895 target_logits = torch.rand(
866 - num_tokens, vocab_size, dtype=torch.float32, device=DEVICE
896 + num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
897 )
898 target_probs = F.softmax(target_logits, dim=-1)
899
900 # Randomly sample draft token ids from draft probs
901 draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32)
902
873 - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
903 + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
904 generators = {
875 - i: torch.Generator(device=DEVICE).manual_seed(i) for i in range(batch_size)
905 + i: torch.Generator(device=DEVICE_TYPE).manual_seed(i) for i in range(batch_size)
906 }
907 sampling_metadata = create_sampling_metadata(
908 all_greedy=False, temperature=temperature, generators=generators
@@ -890,7 +920,7 @@ def test_sample_recovered_tokens(
920 None if no_draft_probs else draft_probs,
921 target_probs,
922 sampling_metadata,
893 - device=DEVICE,
923 + device=DEVICE_TYPE,
924 )
925 recovered_token_ids = sample_recovered_tokens(
926 max_spec_len,
@@ -900,6 +930,6 @@ def test_sample_recovered_tokens(
930 None if no_draft_probs else draft_probs,
931 target_probs,
932 sampling_metadata,
903 - device=DEVICE,
933 + device=DEVICE_TYPE,
934 )
935 assert torch.equal(recovered_token_ids, ref_recovered_token_ids)
tests/v1/sample/test_sampler.py
+8 -7
@@ -17,8 +17,9 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available()
17 MAX_NUM_REQS = 256
18 VOCAB_SIZE = 1024
19 NUM_OUTPUT_TOKENS = 20
20 -CUDA_DEVICES = [
21 - f"{current_platform.device_type}:{i}"
20 +DEVICE_TYPE = current_platform.device_type
21 +DEVICES = [
22 + f"{DEVICE_TYPE}:{i}"
23 for i in range(1 if current_platform.device_count() == 1 else 2)
24 ]
25 MAX_NUM_PROMPT_TOKENS = 64
@@ -199,7 +200,7 @@ def _create_weighted_output_token_list(
200 return output_token_ids, sorted_token_ids_in_output
201
202
202 -@pytest.mark.parametrize("device", CUDA_DEVICES)
203 +@pytest.mark.parametrize("device", DEVICES)
204 @pytest.mark.parametrize("batch_size", [1, 2, 32])
205 @pytest.mark.parametrize("presence_penalty", [-2.0, 2.0])
206 def test_sampler_presence_penalty(
@@ -249,7 +250,7 @@ def test_sampler_presence_penalty(
250 assert penalized_token_id not in output_token_ids[batch_idx]
251
252
252 -@pytest.mark.parametrize("device", CUDA_DEVICES)
253 +@pytest.mark.parametrize("device", DEVICES)
254 @pytest.mark.parametrize("batch_size", [1, 2, 32])
255 @pytest.mark.parametrize("frequency_penalty", [-2.0, 2.0])
256 def test_sampler_frequency_penalty(
@@ -305,7 +306,7 @@ def test_sampler_frequency_penalty(
306 assert penalized_token_id not in distinct_sorted_token_ids_in_output
307
308
308 -@pytest.mark.parametrize("device", CUDA_DEVICES)
309 +@pytest.mark.parametrize("device", DEVICES)
310 @pytest.mark.parametrize("batch_size", [1, 2, 32])
311 @pytest.mark.parametrize("repetition_penalty", [0.1, 1.9])
312 def test_sampler_repetition_penalty(
@@ -363,7 +364,7 @@ def test_sampler_repetition_penalty(
364 )
365
366
366 -@pytest.mark.parametrize("device", CUDA_DEVICES)
367 +@pytest.mark.parametrize("device", DEVICES)
368 @pytest.mark.parametrize("batch_size", [1, 2, 32])
369 @pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])
370 def test_sampler_allowed_token_ids(
@@ -409,7 +410,7 @@ def test_sampler_allowed_token_ids(
410 assert logits_for_req[token_id] != -float("inf")
411
412
412 -@pytest.mark.parametrize("device", CUDA_DEVICES)
413 +@pytest.mark.parametrize("device", DEVICES)
414 @pytest.mark.parametrize("batch_size", [1, 2, 32])
415 @pytest.mark.parametrize("bad_words_lengths", [(1,), (1, 3), (2, 2)])
416 def test_sampler_bad_words(
tests/v1/sample/test_topk_topp_sampler.py
+8 -9
@@ -7,8 +7,7 @@ from torch import Generator
7 from vllm.platforms import current_platform
8 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
9
10 -CUDA_DEVICE = "cuda" if current_platform.is_cuda() else None
11 -DEVICE = current_platform.device_type
10 +DEVICE_TYPE = current_platform.device_type
11
12 BATCH_SIZE = 1024
13 VOCAB_SIZE = 128 * 1024
@@ -26,8 +25,8 @@ def reset_default_device():
25
26
27 def test_topk_impl_equivalence():
29 - torch.set_default_device(DEVICE)
30 - generator = Generator(device=DEVICE).manual_seed(33)
28 + torch.set_default_device(DEVICE_TYPE)
29 + generator = Generator(device=DEVICE_TYPE).manual_seed(33)
30
31 logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
32
@@ -76,8 +75,8 @@ def test_flashinfer_sampler():
75 if not FLASHINFER_ENABLED:
76 pytest.skip("FlashInfer not installed or not available on this platform.")
77
79 - torch.set_default_device(DEVICE)
80 - generator = Generator(device=DEVICE).manual_seed(42)
78 + torch.set_default_device(DEVICE_TYPE)
79 + generator = Generator(device=DEVICE_TYPE).manual_seed(42)
80
81 # Generate random logits
82 logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
@@ -128,15 +127,15 @@ def test_flashinfer_sampler():
127 # =============================================================================
128
129
131 -@pytest.mark.skipif(CUDA_DEVICE is None, reason="CUDA not available")
130 +@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available")
131 class TestTritonTopkTopp:
132 """Tests for the Triton top-k/top-p kernel."""
133
134 @pytest.fixture(autouse=True)
135 def setup(self):
136 """Set up test fixtures."""
138 - torch.set_default_device(CUDA_DEVICE)
139 - self.generator = Generator(device=CUDA_DEVICE).manual_seed(42)
137 + torch.set_default_device(DEVICE_TYPE)
138 + self.generator = Generator(device=DEVICE_TYPE).manual_seed(42)
139
140 def _compare_results(
141 self,
tests/v1/spec_decode/test_eagle.py
+10 -9
@@ -42,6 +42,7 @@ dflash_target_dir = "Qwen/Qwen3-8B"
42 dflash_dir = "z-lab/Qwen3-8B-DFlash-b16"
43
44 BLOCK_SIZE = 16
45 +DEVICE_TYPE = current_platform.device_type
46
47
48 def _create_proposer(
@@ -92,7 +93,7 @@ def _create_proposer(
93 # Overwrite pard_token to avoid crash during init
94 speculative_config.draft_model_config.hf_config.pard_token = 0
95
95 - device = current_platform.device_type
96 + device = DEVICE_TYPE
97 vllm_config = VllmConfig(
98 model_config=model_config,
99 cache_config=CacheConfig(block_size=16),
@@ -124,7 +125,7 @@ def test_prepare_next_token_ids():
125 either the GPU tensor of sampled_token_ids with -1 for rejected tokens,
126 or the CPU python list[list[int]] with the rejected tokens removed.
127 """
127 - device = torch.device(current_platform.device_type)
128 + device = torch.device(DEVICE_TYPE)
129
130 num_requests = 4
131 num_speculative_tokens = 4
@@ -207,7 +208,7 @@ def test_prepare_inputs():
208 a, a + 1, ..., a + b - n2 - 1,
209 a + b, a + b + 1, ..., a + b + c - n3 - 1]
210 """
210 - device = torch.device(current_platform.device_type)
211 + device = torch.device(DEVICE_TYPE)
212
213 # q1 = 4, q2 = 7, q3 = 5
214 # n1 = 1, n2 = 3, n3 = 2
@@ -300,7 +301,7 @@ def test_prepare_inputs_padded():
301 from the original indices to sample from.
302 """
303
303 - device = torch.device(current_platform.device_type)
304 + device = torch.device(DEVICE_TYPE)
305
306 expected_token_indices_to_sample = torch.tensor(
307 [1, 5, 6], dtype=torch.int32, device=device
@@ -370,7 +371,7 @@ def test_set_inputs_first_pass_default_eagle():
371 - After inserting next_tokens [100, 200, 300]:
372 [a2, a3, 100, b2, 200, c2, c3, c4, 300]
373 """
373 - device = torch.device(current_platform.device_type)
374 + device = torch.device(DEVICE_TYPE)
375
376 num_speculative_tokens = 3
377 proposer = _create_proposer("eagle", num_speculative_tokens)
@@ -471,7 +472,7 @@ def test_set_inputs_first_pass_draft_model():
472 - idx 5: token 21, pos 1
473 - idx 6: token 200, pos 2 (bonus token)
474 """
474 - device = torch.device(current_platform.device_type)
475 + device = torch.device(DEVICE_TYPE)
476
477 num_speculative_tokens = 2
478 block_size = BLOCK_SIZE
@@ -609,7 +610,7 @@ def test_set_inputs_first_pass_parallel_drafting():
610 - idx 9: bonus token 200
611 - idx 10-11: parallel_drafting_tokens, is_masked=True
612 """
612 - device = torch.device(current_platform.device_type)
613 + device = torch.device(DEVICE_TYPE)
614
615 num_speculative_tokens = 3
616 block_size = BLOCK_SIZE
@@ -859,7 +860,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
860 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
861
862 # Use GPU device
862 - device = torch.device(current_platform.device_type)
863 + device = torch.device(DEVICE_TYPE)
864
865 # Setup test parameters
866 batch_size = 2
@@ -1030,7 +1031,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
1031 )
1032 def test_propose_tree(spec_token_tree):
1033 # Get GPU device.
1033 - device = torch.device(current_platform.device_type)
1034 + device = torch.device(DEVICE_TYPE)
1035
1036 # Setup test parameters.
1037 batch_size = 2
tests/v1/spec_decode/test_eagle_step_kernel.py
+6 -3
@@ -5,11 +5,14 @@
5 import pytest
6 import torch
7
8 +from vllm.platforms import current_platform
9 from vllm.v1.spec_decode.utils import (
10 PADDING_SLOT_ID,
11 eagle_step_update_slot_mapping_and_metadata,
12 )
13
14 +DEVICE_TYPE = current_platform.device_type
15 +
16 # Skip if no CUDA - Triton kernel requires GPU
17 pytest.importorskip("triton")
18 if not torch.cuda.is_available():
@@ -47,7 +50,7 @@ def _reference_eagle_step_slot_mapping(
50
51 def test_eagle_step_slot_mapping_kernel():
52 """Test fused kernel matches Python reference for slot mapping and metadata."""
50 - device = torch.device("cuda")
53 + device = torch.device(DEVICE_TYPE)
54 batch_size = 32
55 block_size = 16
56 max_model_len = 4096
@@ -93,7 +96,7 @@ def test_eagle_step_slot_mapping_kernel():
96
97 def test_eagle_step_slot_mapping_kernel_exceeds_max():
98 """Test fused kernel when position exceeds max_model_len."""
96 - device = torch.device("cuda")
99 + device = torch.device(DEVICE_TYPE)
100 batch_size = 4
101 block_size = 16
102 max_model_len = 100
@@ -130,7 +133,7 @@ def test_eagle_step_slot_mapping_kernel_exceeds_max():
133 def test_eagle_step_slot_mapping_kernel_cudagraph_padding():
134 """Test that padding threads write PADDING_SLOT_ID when
135 input_batch_size > batch_size (cudagraph padding)."""
133 - device = torch.device("cuda")
136 + device = torch.device(DEVICE_TYPE)
137 batch_size = 4
138 input_batch_size = 8
139 block_size = 16
tests/v1/spec_decode/test_extract_hidden_states.py
+6 -5
@@ -27,6 +27,7 @@ from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesPropose
27 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
28
29 model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
30 +DEVICE_TYPE = current_platform.device_type
31
32
33 def _create_proposer(
@@ -51,7 +52,7 @@ def _create_proposer(
52 },
53 )
54
54 - device = current_platform.device_type
55 + device = DEVICE_TYPE
56 vllm_config = VllmConfig(
57 model_config=model_config,
58 cache_config=CacheConfig(),
@@ -101,7 +102,7 @@ def test_proposer_initialization_missing_layer_ids():
102 },
103 )
104
104 - device = current_platform.device_type
105 + device = DEVICE_TYPE
106 vllm_config = VllmConfig(
107 model_config=model_config,
108 cache_config=CacheConfig(),
@@ -130,7 +131,7 @@ def test_prepare_next_token_ids_padded():
131 For each request we either use the sampled token (if valid and not discarded)
132 or a backup token from the request state.
133 """
133 - device = torch.device(current_platform.device_type)
134 + device = torch.device(DEVICE_TYPE)
135
136 num_requests = 4
137 req_ids = [f"req_{i + 1}" for i in range(num_requests)]
@@ -197,7 +198,7 @@ def test_propose():
198 2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1])
199 3. Cache the hidden states in the model's KV cache
200 """
200 - device = torch.device(current_platform.device_type)
201 + device = torch.device(DEVICE_TYPE)
202
203 # Setup test parameters
204 batch_size = 2
@@ -273,7 +274,7 @@ def test_propose():
274 @pytest.mark.parametrize("num_hidden_layers", [1, 4, 8])
275 def test_propose_different_layer_counts(num_hidden_layers):
276 """Test that propose works correctly with different numbers of hidden layers."""
276 - device = torch.device(current_platform.device_type)
277 + device = torch.device(DEVICE_TYPE)
278
279 batch_size = 2
280 num_tokens = 5
tests/v1/spec_decode/test_mtp.py
+4 -3
@@ -28,6 +28,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
28 from vllm.v1.spec_decode.eagle import EagleProposer
29
30 mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base"
31 +DEVICE_TYPE = current_platform.device_type
32
33
34 def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
@@ -48,7 +49,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
49 model_config=model_config,
50 cache_config=CacheConfig(),
51 speculative_config=speculative_config,
51 - device_config=DeviceConfig(device=current_platform.device_type),
52 + device_config=DeviceConfig(device=DEVICE_TYPE),
53 parallel_config=ParallelConfig(),
54 load_config=LoadConfig(),
55 scheduler_config=SchedulerConfig(
@@ -57,7 +58,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:
58 ),
59 )
60
60 - return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type)
61 + return EagleProposer(vllm_config=vllm_config, device=DEVICE_TYPE)
62
63
64 @mock.patch("vllm.v1.spec_decode.eagle.get_pp_group")
@@ -118,7 +119,7 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro
119 def test_mtp_propose(num_speculative_tokens, monkeypatch):
120 """Test that MTP's forward method returns hidden states directly"""
121
121 - device = torch.device(current_platform.device_type)
122 + device = torch.device(DEVICE_TYPE)
123 batch_size = 2
124 seq_lens = [5, 3]
125 total_tokens = sum(seq_lens)
tests/v1/spec_decode/test_tree_attention.py
+5 -3
@@ -18,6 +18,8 @@ from vllm.v1.attention.backend import CommonAttentionMetadata
18 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
19 from vllm.v1.attention.backends.registry import AttentionBackendEnum
20
21 +DEVICE_TYPE = current_platform.device_type
22 +
23 if not is_flash_attn_varlen_func_available():
24 pytest.skip(
25 "This test requires flash_attn_varlen_func, but it's not available.",
@@ -170,9 +172,9 @@ def _get_available_reference_backends() -> list[AttentionBackendEnum]:
172
173
174 class MockAttentionLayer(torch.nn.Module):
173 - _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
174 - _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
175 - _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
175 + _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
176 + _k_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
177 + _v_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)
178 layer_name = "mock_layer"
179
180 def __init__(self):
tests/v1/worker/test_gpu_input_batch.py
+5 -7
@@ -22,10 +22,8 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
22 VOCAB_SIZE = 1024
23 NUM_OUTPUT_TOKENS = 20
24 MAX_PROMPT_SIZE = 100
25 -CUDA_DEVICES = [
26 - f"{current_platform.device_type}:{i}"
27 - for i in range(min(current_platform.device_count(), 2))
28 -]
25 +DEVICE_TYPE = current_platform.device_type
26 +DEVICES = [f"{DEVICE_TYPE}:{i}" for i in range(min(current_platform.device_count(), 2))]
27 MAX_NUM_PROMPT_TOKENS = 64
28
29
@@ -219,7 +217,7 @@ def _construct_cached_request_state(req_id_suffix: int):
217 )
218
219
222 -@pytest.mark.parametrize("device", CUDA_DEVICES)
220 +@pytest.mark.parametrize("device", DEVICES)
221 @pytest.mark.parametrize("batch_size", [1, 2, 32, 64])
222 def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
223 """
@@ -313,7 +311,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
311 )
312
313
316 -@pytest.mark.parametrize("device", CUDA_DEVICES)
314 +@pytest.mark.parametrize("device", DEVICES)
315 @pytest.mark.parametrize("batch_size", [32])
316 @pytest.mark.parametrize("swap_list", [((0, 1),)])
317 def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list):
@@ -400,7 +398,7 @@ def _construct_pooling_request(req_id_suffix: int, pooling_params=None):
398 )
399
400
403 -@pytest.mark.parametrize("device", CUDA_DEVICES)
401 +@pytest.mark.parametrize("device", DEVICES)
402 def test_pooling_prompt_lens_not_aliased(device: str):
403 """Verify that prompt_lens in PoolingMetadata does not share memory
404 with the internal num_prompt_tokens pinned buffer. Guards against possible
tests/v1/worker/test_gpu_model_runner.py
+17 -17
@@ -45,7 +45,7 @@ from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
45
46 BLOCK_SIZE = 16
47 NUM_BLOCKS = 10
48 -DEVICE = current_platform.device_type
48 +DEVICE_TYPE = current_platform.device_type
49
50
51 def initialize_kv_cache(runner: GPUModelRunner):
@@ -121,7 +121,7 @@ def model_runner():
121 vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
122 num_heads, head_size, 0.1
123 )
124 - runner = GPUModelRunner(vllm_config, DEVICE)
124 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
125 initialize_kv_cache(runner)
126 yield runner
127
@@ -340,7 +340,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
340 [1.0, 2.0, 3.0],
341 [3.0, 2.0, 1.0],
342 ],
343 - device=DEVICE,
343 + device=DEVICE_TYPE,
344 )
345 result = model_runner._get_nans_in_logits(logits)
346 assert result == {"req_0": 0, "req_1": 0}
@@ -350,7 +350,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
350 [1.0, float("nan"), 3.0],
351 [4.0, float("nan"), float("nan")],
352 ],
353 - device=DEVICE,
353 + device=DEVICE_TYPE,
354 )
355 result = model_runner._get_nans_in_logits(logits)
356 assert result == {"req_0": 1, "req_1": 2}
@@ -360,7 +360,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
360 [1.0, 2.0, 3.0],
361 [4.0, float("nan"), float("nan")],
362 ],
363 - device=DEVICE,
363 + device=DEVICE_TYPE,
364 )
365 result = model_runner._get_nans_in_logits(logits)
366 assert result == {"req_0": 0, "req_1": 2}
@@ -372,7 +372,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
372 [
373 [1.0, float("nan"), 3.0],
374 ],
375 - device=DEVICE,
375 + device=DEVICE_TYPE,
376 )
377 result = model_runner._get_nans_in_logits(logits)
378 assert result == {"req_0": 1, "req_1": 0}
@@ -383,7 +383,7 @@ def test_get_nans_in_logits(model_runner, dist_init):
383 [1.0, 2.0, 3.0],
384 [float("nan"), 2.0, 3.0],
385 ],
386 - device=DEVICE,
386 + device=DEVICE_TYPE,
387 )
388 result = model_runner._get_nans_in_logits(logits)
389 assert result == {"req_0": 2, "req_1": 0}
@@ -643,7 +643,7 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
643 # Set high context length to test max context length estimation
644 vllm_config.model_config.max_model_len = 3_000_000
645 vllm_ctx = vllm_config.compilation_config.static_forward_context
646 - runner = GPUModelRunner(vllm_config, DEVICE)
646 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
647 kv_cache_spec = runner.get_kv_cache_spec()
648 assert len(kv_cache_spec) == 2
649 assert len(runner.shared_kv_cache_layers) == 0
@@ -711,7 +711,7 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
711 # Set high context length to test max context length estimation
712 vllm_config.model_config.max_model_len = 3_000_000
713 vllm_ctx = vllm_config.compilation_config.static_forward_context
714 - runner = GPUModelRunner(vllm_config, DEVICE)
714 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
715 kv_cache_spec = runner.get_kv_cache_spec()
716 assert len(kv_cache_spec) == 1
717 assert layer_0 in kv_cache_spec
@@ -850,7 +850,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
850 assert fwd_context is not None
851 vllm_ctx = vllm_config.compilation_config.static_forward_context
852
853 - runner = GPUModelRunner(vllm_config, DEVICE)
853 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
854 current_platform.update_block_size_for_backend(vllm_config)
855 kv_cache_spec = runner.get_kv_cache_spec()
856
@@ -896,13 +896,13 @@ def test_hybrid_attention_mamba_tensor_shapes():
896 ssm_constant_shape = ssm_shape[1:]
897
898 attn_blocks_constant = torch.full(
899 - (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33
899 + (test_block_size, *attn_constant_shape), device=DEVICE_TYPE, fill_value=3.33
900 )
901 conv_blocks_constant = torch.full(
902 - (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66
902 + (test_block_size, *conv_constant_shape), device=DEVICE_TYPE, fill_value=6.66
903 )
904 ssm_blocks_constant = torch.full(
905 - (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99
905 + (test_block_size, *ssm_constant_shape), device=DEVICE_TYPE, fill_value=9.99
906 )
907
908 # Fill attention blocks with constants using kv block indices
@@ -997,7 +997,7 @@ def test_hybrid_block_table_initialization():
997 max_num_blocks_per_req=max_num_blocks_per_req,
998 max_num_batched_tokens=max_num_batched_tokens,
999 pin_memory=False,
1000 - device=torch.device(DEVICE),
1000 + device=torch.device(DEVICE_TYPE),
1001 kernel_block_size=kernel_block_sizes[0],
1002 cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
1003 )
@@ -1036,7 +1036,7 @@ def test_input_batch_with_kernel_block_sizes():
1036 max_num_reqs = 10
1037 max_model_len = 512
1038 max_num_batched_tokens = 512
1039 - device = torch.device(DEVICE)
1039 + device = torch.device(DEVICE_TYPE)
1040 pin_memory = False
1041 vocab_size = 50272
1042
@@ -1083,7 +1083,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
1083 num_heads, head_size, 0.1
1084 )
1085
1086 - runner = GPUModelRunner(vllm_config, DEVICE)
1086 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
1087
1088 # Initialize KV cache with configuration
1089 attn_spec = FullAttentionSpec(
@@ -1306,7 +1306,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
1306 )
1307 assert fwd_context is not None
1308
1309 - runner = GPUModelRunner(vllm_config, DEVICE)
1309 + runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
1310 current_platform.update_block_size_for_backend(vllm_config)
1311 kv_cache_spec = runner.get_kv_cache_spec()
1312