3abf858443 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/basic_correctness/test_cumem.py

+10 -8

@@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes

13

14 from ..utils import create_new_process_for_each_test, requires_fp8

15

16 +DEVICE_TYPE = current_platform.device_type

17 +

18

19 @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")

20 def test_python_error():

@@ -26,13 +28,13 @@ def test_python_error():

28 tensors = []

29 with allocator.use_memory_pool():

30 # allocate 70% of the total memory

29 - x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")

31 + x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)

32 tensors.append(x)

33 # release the memory

34 allocator.sleep()

35

36 # allocate more memory than the total memory

35 - y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")

37 + y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)

38 tensors.append(y)

39 with pytest.raises(RuntimeError):

40 # when the allocator is woken up, it should raise an error

@@ -44,17 +46,17 @@ def test_python_error():

46 def test_basic_cumem():

47 # some tensors from default memory pool

48 shape = (1024, 1024)

47 - x = torch.empty(shape, device="cuda")

49 + x = torch.empty(shape, device=DEVICE_TYPE)

50 x.zero_()

51

52 # some tensors from custom memory pool

53 allocator = CuMemAllocator.get_instance()

54 with allocator.use_memory_pool():

55 # custom memory pool

54 - y = torch.empty(shape, device="cuda")

56 + y = torch.empty(shape, device=DEVICE_TYPE)

57 y.zero_()

58 y += 1

57 - z = torch.empty(shape, device="cuda")

59 + z = torch.empty(shape, device=DEVICE_TYPE)

60 z.zero_()

61 z += 2

62

@@ -77,16 +79,16 @@ def test_basic_cumem():

79 def test_cumem_with_cudagraph():

80 allocator = CuMemAllocator.get_instance()

81 with allocator.use_memory_pool():

80 - weight = torch.eye(1024, device="cuda")

82 + weight = torch.eye(1024, device=DEVICE_TYPE)

83 with allocator.use_memory_pool(tag="discard"):

82 - cache = torch.empty(1024, 1024, device="cuda")

84 + cache = torch.empty(1024, 1024, device=DEVICE_TYPE)

85

86 def model(x):

87 out = x @ weight

88 cache[: out.size(0)].copy_(out)

89 return out + 1

90

89 - x = torch.empty(128, 1024, device="cuda")

91 + x = torch.empty(128, 1024, device=DEVICE_TYPE)

92

93 # warmup

94 model(x)

tests/compile/passes/distributed/test_async_tp.py

+3 -2

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform

31 from vllm.utils.system_utils import update_environment_variables

32 from vllm.utils.torch_utils import set_random_seed

33

34 +DEVICE_TYPE = current_platform.device_type

35 FP8_DTYPE = current_platform.fp8_dtype()

36

37 prompts = [

@@ -299,7 +300,7 @@ def async_tp_pass_on_test_model(

300 ):

301 set_random_seed(0)

302

302 - device = torch.device(f"cuda:{local_rank}")

303 + device = torch.device(f"{DEVICE_TYPE}:{local_rank}")

304 torch.accelerator.set_device_index(device)

305 torch.set_default_device(device)

306 torch.set_default_dtype(dtype)

@@ -324,7 +325,7 @@ def async_tp_pass_on_test_model(

325 fuse_gemm_comms=True,

326 ),

327 )

327 - vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

328 + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))

329

330 # this is a fake model name to construct the model config

331 # in the vllm_config, it's not really used.

tests/compile/passes/distributed/test_fusion_all_reduce.py

+4 -2

@@ -37,6 +37,8 @@ from vllm.platforms import current_platform

37 from vllm.utils.system_utils import update_environment_variables

38 from vllm.utils.torch_utils import set_random_seed

39

40 +DEVICE_TYPE = current_platform.device_type

41 +

42

43 class TestAllReduceRMSNormModel(torch.nn.Module):

44 def __init__(

@@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model(

270 ):

271 set_random_seed(0)

272

271 - device = torch.device(f"cuda:{local_rank}")

273 + device = torch.device(f"{DEVICE_TYPE}:{local_rank}")

274 torch.accelerator.set_device_index(device)

275 torch.set_default_device(device)

276 torch.set_default_dtype(dtype)

@@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model(

302 vllm_config.compilation_config.pass_config = PassConfig(

303 fuse_allreduce_rms=True, eliminate_noops=True

304 )

303 - vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

305 + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))

306 vllm_config.parallel_config.rank = local_rank # Setup rank for debug path

307

308 # this is a fake model name to construct the model config

tests/compile/passes/distributed/test_sequence_parallelism.py

+4 -2

@@ -35,6 +35,8 @@ from vllm.platforms import current_platform

35 from vllm.utils.system_utils import update_environment_variables

36 from vllm.utils.torch_utils import set_random_seed

37

38 +DEVICE_TYPE = current_platform.device_type

39 +

40 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")

41

42 FP8_DTYPE = current_platform.fp8_dtype()

@@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model(

230 ):

231 set_random_seed(0)

232

231 - device = torch.device(f"cuda:{local_rank}")

233 + device = torch.device(f"{DEVICE_TYPE}:{local_rank}")

234 torch.accelerator.set_device_index(device)

235 torch.set_default_device(device)

236 torch.set_default_dtype(dtype)

@@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model(

260 eliminate_noops=True,

261 ),

262 ) # NoOp needed for fusion

261 - device_config = DeviceConfig(device=torch.device("cuda"))

263 + device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))

264

265 # this is a fake model name to construct the model config

266 # in the vllm_config, it's not really used.

tests/compile/passes/test_fusion_attn.py

+2 -1

@@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata

41 from vllm.v1.attention.backends.registry import AttentionBackendEnum

42 from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode

43

44 +DEVICE_TYPE = current_platform.device_type

45 FP8_DTYPE = current_platform.fp8_dtype()

46 FP4_DTYPE = torch.uint8

47

@@ -300,7 +301,7 @@ def test_attention_quant_pattern(

301

302 custom_ops_list = custom_ops.split(",") if custom_ops else []

303

303 - device = torch.device("cuda:0")

304 + device = torch.device(f"{DEVICE_TYPE}:0")

305 torch.set_default_dtype(dtype)

306 torch.manual_seed(42)

307

tests/compile/passes/test_mla_attn_quant_fusion.py

+2 -1

@@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec

45

46 FP8_DTYPE = current_platform.fp8_dtype()

47 FP4_DTYPE = torch.uint8

48 +DEVICE_TYPE = current_platform.device_type

49

50

51 class MLAAttentionQuantPatternModel(torch.nn.Module):

@@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern(

357

358 custom_ops_list = custom_ops.split(",") if custom_ops else []

359

359 - device = torch.device("cuda:0")

360 + device = torch.device(f"{DEVICE_TYPE}:0")

361 torch.set_default_dtype(dtype)

362 torch.manual_seed(42)

363

tests/compile/passes/test_noop_elimination.py

+5 -2

@@ -8,6 +8,9 @@ import vllm

8 from tests.compile.backend import TestBackend

9 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass

10 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig

11 +from vllm.platforms import current_platform

12 +

13 +DEVICE_TYPE = current_platform.device_type

14

15

16 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])

@@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf

20 )

21 @pytest.mark.parametrize("hidden_size", [64, 4096])

22 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):

20 - torch.set_default_device("cuda")

23 + torch.set_default_device(DEVICE_TYPE)

24 torch.set_default_dtype(dtype)

25 torch.manual_seed(1)

26

@@ -88,7 +91,7 @@ def test_non_noop_slice_preserved():

91 Regression test for a bug where end=-1 was treated like an inferred

92 dimension (reshape semantics) leading to incorrect elimination.

93 """

91 - torch.set_default_device("cuda")

94 + torch.set_default_device(DEVICE_TYPE)

95 x = torch.randn(16, 16)

96

97 class SliceModel(torch.nn.Module):

tests/compile/passes/test_scatter_split_replace.py

+4 -1

@@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import (

13 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass

14 from vllm.config import CompilationConfig, CompilationMode, VllmConfig

15 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

16 +from vllm.platforms import current_platform

17 +

18 +DEVICE_TYPE = current_platform.device_type

19

20

21 class ScatterSplitReplacementModel(nn.Module):

@@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module):

64

65 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

66 def test_scatter_split_replace(dtype):

64 - torch.set_default_device("cuda")

67 + torch.set_default_device(DEVICE_TYPE)

68 torch.set_default_dtype(dtype)

69 torch.manual_seed(0)

70

tests/compile/passes/test_split_coalescing.py

+4 -1

@@ -8,6 +8,9 @@ import vllm

8 from tests.compile.backend import TestBackend

9 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass

10 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig

11 +from vllm.platforms import current_platform

12 +

13 +DEVICE_TYPE = current_platform.device_type

14

15

16 class SplitCoalescingModel(torch.nn.Module):

@@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module):

31

32 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

33 def test_split_coalescing(dtype):

31 - torch.set_default_device("cuda")

34 + torch.set_default_device(DEVICE_TYPE)

35 torch.set_default_dtype(dtype)

36 torch.manual_seed(0)

37

tests/compile/test_config.py

+4 -2

@@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

31 # This import automatically registers `torch.ops.silly.attention`

32 from . import silly_attention # noqa: F401

33

34 +DEVICE_TYPE = current_platform.device_type

35 +

36

37 def test_version():

38 # Test the version comparison logic using the private function

@@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config):

458 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

459

460 dtype = torch.bfloat16

459 - device = torch.device("cuda:0")

461 + device = torch.device(f"{DEVICE_TYPE}:0")

462 batch_size, num_qo_heads, head_size = 8, 16, 128

463

464 # access and cache default compilation config

@@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config):

480 query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)

481 query_quant = torch.compile(query_quant)

482

481 - _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")

483 + _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE)

484 query = torch.randn(

485 batch_size, num_qo_heads * head_size, dtype=dtype, device=device

486 )

tests/compile/test_graph_partition.py

+5 -2

@@ -15,10 +15,13 @@ from vllm.compilation.backends import (

15 split_graph,

16 )

17 from vllm.compilation.passes.fx_utils import find_op_nodes

18 +from vllm.platforms import current_platform

19

20 # This import automatically registers `torch.ops.silly.attention`

21 from . import silly_attention # noqa: F401

22

23 +DEVICE_TYPE = current_platform.device_type

24 +

25

26 def test_getitem_moved_to_producer_subgraph():

27 """

@@ -151,7 +154,7 @@ def test_consecutive_ops_in_split():

154 final_result = torch.sigmoid(attn_inout)

155 return final_result

156

154 - torch.set_default_device("cuda")

157 + torch.set_default_device(DEVICE_TYPE)

158

159 # Create the traced FX graph for the model

160 x = torch.randn(8, 4)

@@ -329,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged():

332 "Expected two builtin empty_like nodes in merged non-splitting subgraph"

333 )

334

332 - x = torch.randn(2, 3, device="cuda")

335 + x = torch.randn(2, 3, device=DEVICE_TYPE)

336 output_original = gm(x)

337 output_split = split_gm(x)

338 assert torch.allclose(output_original, output_split), "Output mismatch after split"

tests/compile/test_rotary_embedding_compile.py

+3 -1

@@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode

16 from vllm.model_executor.layers.rotary_embedding import get_rope

17 from vllm.platforms import current_platform

18

19 +DEVICE_TYPE = current_platform.device_type

20 +

21

22 @support_torch_compile

23 class RotaryEmbeddingCompileModule(torch.nn.Module):

@@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch):

47 monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1")

48 monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")

49

48 - device = "cuda"

50 + device = DEVICE_TYPE

51 positions = torch.arange(16, device=device)

52 query = torch.randn(16, 32, device=device, dtype=torch.bfloat16)

53 key = torch.randn(16, 32, device=device, dtype=torch.bfloat16)

tests/compile/test_structured_logging.py

+3 -1

@@ -17,8 +17,10 @@ from vllm.config.compilation import (

17 )

18 from vllm.config.scheduler import SchedulerConfig

19 from vllm.forward_context import set_forward_context

20 +from vllm.platforms import current_platform

21

22 MLP_SIZE = 64

23 +DEVICE_TYPE = current_platform.device_type

24

25

26 @support_torch_compile

@@ -71,7 +73,7 @@ class TraceStructuredCapture:

73 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")

74 def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):

75 """Test that all expected vLLM artifacts are logged during compilation."""

74 - torch.set_default_device("cuda")

76 + torch.set_default_device(DEVICE_TYPE)

77

78 capture = TraceStructuredCapture()

79

tests/model_executor/test_eagle_quantization.py

+3 -2

@@ -10,9 +10,10 @@ from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig

10 from vllm.model_executor.models.utils import get_draft_quant_config

11 from vllm.platforms import current_platform

12

13 +DEVICE_TYPE = current_platform.device_type

14 DEVICES = (

14 - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)]

15 - if current_platform.is_cuda_alike()

15 + [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]

16 + if not current_platform.is_cpu()

17 else ["cpu"]

18 )

19

tests/models/multimodal/pooling/test_intern_vit.py

+8 -4

@@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download

7 from transformers import AutoConfig, AutoModel, CLIPImageProcessor

8

9 from vllm.distributed import cleanup_dist_env_and_memory

10 +from vllm.platforms import current_platform

11 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

12

13 from ....conftest import ImageTestAssets

@@ -15,6 +16,8 @@ from ....conftest import ImageTestAssets

16 # dynamic_module and trust_remote_code for hf_runner

17 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]

18

19 +DEVICE_TYPE = current_platform.device_type

20 +

21

22 @torch.inference_mode()

23 def run_intern_vit_test(

@@ -39,9 +42,9 @@ def run_intern_vit_test(

42

43 hf_model = AutoModel.from_pretrained(

44 model, dtype=torch_dtype, trust_remote_code=True

42 - ).to("cuda")

45 + ).to(DEVICE_TYPE)

46 hf_outputs_per_image = [

44 - hf_model(pixel_value.to("cuda")).last_hidden_state

47 + hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state

48 for pixel_value in pixel_values

49 ]

50

@@ -53,9 +56,10 @@ def run_intern_vit_test(

56 del hf_model

57 cleanup_dist_env_and_memory()

58

56 - vllm_model = vllm_model.to("cuda", torch_dtype)

59 + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)

60 vllm_outputs_per_image = [

58 - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values

61 + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))

62 + for pixel_value in pixel_values

63 ]

64 del vllm_model

65 cleanup_dist_env_and_memory()

tests/models/multimodal/pooling/test_radio.py

+8 -4

@@ -8,6 +8,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor

8

9 from vllm.distributed import cleanup_dist_env_and_memory

10 from vllm.model_executor.models.radio import RadioModel

11 +from vllm.platforms import current_platform

12 from vllm.transformers_utils.configs.radio import RadioConfig

13 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

14

@@ -17,6 +18,8 @@ from ....conftest import ImageTestAssets

18 # dynamic_module and trust_remote_code for hf_runner

19 DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]

20

21 +DEVICE_TYPE = current_platform.device_type

22 +

23

24 @torch.inference_mode()

25 def run_radio_test(

@@ -51,7 +54,7 @@ def run_radio_test(

54 config=hf_config,

55 dtype=torch_dtype,

56 trust_remote_code=True,

54 - ).to("cuda")

57 + ).to(DEVICE_TYPE)

58 hf_model.eval()

59

60 # A HF model has image normalization as a part of model's forward

@@ -62,7 +65,7 @@ def run_radio_test(

65 hf_model.make_preprocessor_external()

66

67 hf_outputs_per_image = [

65 - hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values

68 + hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values

69 ]

70

71 vllm_config = RadioConfig(

@@ -71,10 +74,11 @@ def run_radio_test(

74 )

75 vllm_model = RadioModel(vllm_config)

76 vllm_model.load_weights(hf_model.state_dict())

74 - vllm_model = vllm_model.to("cuda", torch_dtype)

77 + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype)

78

79 vllm_outputs_per_image = [

77 - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values

80 + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE))

81 + for pixel_value in pixel_values

82 ]

83 del vllm_model, hf_model

84 cleanup_dist_env_and_memory()

tests/models/test_utils.py

+8 -2

@@ -10,6 +10,8 @@ from vllm.model_executor.models.utils import (

10 )

11 from vllm.platforms import current_platform

12

13 +DEVICE_TYPE = current_platform.device_type

14 +

15

16 class ModuleWithBatchNorm(torch.nn.Module):

17 def __init__(self):

@@ -174,8 +176,12 @@ class raise_if_cuda_sync:

176

177 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")

178 def test_merge_multimodal_embeddings_no_sync():

177 - inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0")

178 - multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")]

179 + inputs_embeds = torch.zeros(

180 + [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0"

181 + )

182 + multimodal_embeddings = [

183 + torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0")

184 + ]

185 is_multimodal = torch.tensor([True, False, True, True, False], device="cpu")

186 with raise_if_cuda_sync():

187 _merge_multimodal_embeddings(

tests/quantization/test_fp8.py

+7 -3

@@ -24,6 +24,8 @@ from vllm.model_executor.layers.quantization.fp8 import (

24 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

25 from vllm.platforms import current_platform

26

27 +DEVICE_TYPE = current_platform.device_type

28 +

29 MODELS = [

30 "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV",

31 # The checkpoint below was removed from the HF.

@@ -314,7 +316,7 @@ def test_scaled_fp8_quant(dtype) -> None:

316

317 # Note that we use a shape % 4 != 0 to cover edge cases,

318 # because scaled_fp8_quant is vectorized by 4.

317 - x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype)

319 + x = (torch.randn(size=(11, 11), device=DEVICE_TYPE) * 13).to(dtype)

320

321 # Dynamic quantization

322 ref_y, inv_scale = ops.scaled_fp8_quant(x, None)

@@ -338,7 +340,9 @@ def test_scaled_fp8_quant(dtype) -> None:

340

341 # non-contiguous input with padding

342 m, n, padded_stride = 975, 512, 576

341 - padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype)

343 + padded_tensor = (torch.randn(size=(m, padded_stride), device=DEVICE_TYPE) * 13).to(

344 + dtype

345 + )

346 x_nc = padded_tensor[:, :n] # shape (m, n) with stride (padded_stride, 1)

347

348 assert not x_nc.is_contiguous()

@@ -409,7 +413,7 @@ def test_fp8_reloading(

413

414 # Set model config as model_config.dtype is required in Fp8LinearMethod.

415 default_vllm_config.model_config = ModelConfig()

412 - with torch.device("cuda:0"):

416 + with torch.device(f"{DEVICE_TYPE}:0"):

417 config = Fp8Config(

418 is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,

419 weight_block_size=weight_block_size,

tests/quantization/test_per_token_kv_cache.py

+8 -6

@@ -25,11 +25,13 @@ from vllm.platforms import current_platform

25 from vllm.utils.torch_utils import set_random_seed

26 from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache

27

28 +DEVICE_TYPE = current_platform.device_type

29 +

30 # Skip entire module if no CUDA/ROCm GPU available

31 pytestmark = [

32 pytest.mark.skipif(

31 - not current_platform.is_cuda_alike(),

32 - reason="Per-token-head KV cache tests require CUDA or ROCm GPU.",

33 + current_platform.is_cpu(),

34 + reason="Per-token-head KV cache tests require GPU.",

35 ),

36 ]

37

@@ -166,7 +168,7 @@ def test_reshape_and_cache_per_token_head(

168 )

169

170 set_random_seed(seed)

169 - torch.set_default_device("cuda")

171 + torch.set_default_device(DEVICE_TYPE)

172

173 num_blocks = (num_tokens + block_size - 1) // block_size + 4

174

@@ -260,7 +262,7 @@ def test_per_token_head_round_trip_accuracy(

262 triton_reshape_and_cache_flash_per_token_head_quant,

263 )

264

263 - torch.set_default_device("cuda")

265 + torch.set_default_device(DEVICE_TYPE)

266 set_random_seed(42)

267

268 num_blocks = (num_tokens + block_size - 1) // block_size + 2

@@ -323,7 +325,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig):

325 triton_reshape_and_cache_flash_per_token_head_quant,

326 )

327

326 - torch.set_default_device("cuda")

328 + torch.set_default_device(DEVICE_TYPE)

329 num_tokens = 4

330 num_heads = 2

331 head_size = 64

@@ -430,7 +432,7 @@ def test_triton_unified_attention_per_token_head_scale(

432 from vllm.utils.math_utils import next_power_of_2

433 from vllm.v1.attention.ops.triton_unified_attention import unified_attention

434

433 - torch.set_default_device("cuda")

435 + torch.set_default_device(DEVICE_TYPE)

436 set_random_seed(0)

437

438 num_seqs = len(seq_lens)

tests/quantization/test_quark.py

+7 -5

@@ -36,6 +36,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(

36 importlib.metadata.version("amd-quark")

37 ) >= version.parse(QUARK_MXFP4_MIN_VERSION)

38

39 +DEVICE_TYPE = current_platform.device_type

40 +

41 if QUARK_MXFP4_AVAILABLE:

42 from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer

43 from quark.torch.kernel import mx as mx_kernel

@@ -309,7 +311,7 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in

311 torch.manual_seed(0)

312

313 hidden_size = 64 * 32

312 - inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 0.5) * 2

314 + inp = (torch.rand(1, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2

315 for i in range(hidden_size // 32):

316 inp[:, i * 32 : (i + 1) * 32] = (

317 inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)]

@@ -353,15 +355,15 @@ def test_mxfp4_dequant_kernel_match_quark(

355 reorder=False,

356 real_quantized=True,

357 float_dtype=float_dtype,

356 - device="cuda",

358 + device=DEVICE_TYPE,

359 )

360

359 - observer = qspec.observer_cls(qspec, device="cuda")

361 + observer = qspec.observer_cls(qspec, device=DEVICE_TYPE)

362

363 hidden_size = 512

364 shape = (11008, hidden_size)

365

364 - w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2

366 + w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2

367

368 # Make it so that different groups have different scales.

369 for i in range(hidden_size // 32):

@@ -373,7 +375,7 @@ def test_mxfp4_dequant_kernel_match_quark(

375 scale, _ = observer._calculate_qparams()

376 weight_quantizer.scale = scale

377

376 - w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda")

378 + w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to(DEVICE_TYPE)

379 weight_quantizer.maybe_convert_and_transpose_scale()

380

381 scale = weight_quantizer.scale

tests/quantization/test_torchao.py

+9 -8

@@ -8,6 +8,7 @@ import torch

8 from vllm.model_executor.model_loader import get_model_loader

9 from vllm.platforms import current_platform

10

11 +DEVICE_TYPE = current_platform.device_type

12 DTYPE = ["bfloat16"]

13

14 TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None

@@ -33,7 +34,7 @@ def test_pre_quantized_model(vllm_runner):

34 @pytest.mark.parametrize(

35 "pt_load_map_location",

36 [

36 - "cuda:0",

37 + f"{DEVICE_TYPE}:0",

38 # {"": "cuda"},

39 ],

40 )

@@ -60,7 +61,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):

61 model_name=model_name,

62 quantization="torchao",

63 dtype="bfloat16",

63 - pt_load_map_location="cuda:0",

64 + pt_load_map_location=f"{DEVICE_TYPE}:0",

65 enforce_eager=True,

66 ) as llm:

67 output = llm.generate_greedy(["The capital of France is"], max_tokens=4)

@@ -81,7 +82,7 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):

82 model_name=model_name,

83 quantization="torchao",

84 dtype="bfloat16",

84 - pt_load_map_location="cuda:0",

85 + pt_load_map_location=f"{DEVICE_TYPE}:0",

86 ) as llm:

87 output = llm.generate_greedy(["The capital of France is"], max_tokens=4)

88

@@ -112,7 +113,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle):

113 with vllm_runner(

114 model_name=model_name,

115 dtype="bfloat16",

115 - pt_load_map_location="cuda:0",

116 + pt_load_map_location=f"{DEVICE_TYPE}:0",

117 quantization="torchao",

118 hf_overrides=hf_overrides,

119 enforce_eager=True,

@@ -158,7 +159,7 @@ def test_online_quant_config_file(vllm_runner):

159 with vllm_runner(

160 model_name=model_name,

161 dtype="bfloat16",

161 - pt_load_map_location="cuda:0",

162 + pt_load_map_location=f"{DEVICE_TYPE}:0",

163 quantization="torchao",

164 hf_overrides=hf_overrides,

165 enforce_eager=True,

@@ -248,7 +249,7 @@ def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner):

249 torch._dynamo.reset()

250 model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev"

251 with vllm_runner(

251 - model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0"

252 + model_name=model_name, dtype="bfloat16", pt_load_map_location=f"{DEVICE_TYPE}:0"

253 ) as llm:

254 output = llm.generate_greedy(["The capital of France is"], max_tokens=4)

255

@@ -278,7 +279,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel(vllm_runner, monkeypat

279 model_name=model_name,

280 quantization="torchao",

281 dtype="bfloat16",

281 - pt_load_map_location="cuda:0",

282 + pt_load_map_location=f"{DEVICE_TYPE}:0",

283 enforce_eager=True,

284 ) as llm:

285

@@ -357,7 +358,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant(

358 model_name=model_name,

359 quantization="torchao",

360 dtype="bfloat16",

360 - pt_load_map_location="cuda:0",

361 + pt_load_map_location=f"{DEVICE_TYPE}:0",

362 hf_overrides=hf_overrides,

363 enforce_eager=True,

364 ) as llm:

tests/test_config.py

+4 -2

@@ -34,6 +34,8 @@ from vllm.config.vllm import (

34 )

35 from vllm.platforms import current_platform

36

37 +DEVICE_TYPE = current_platform.device_type

38 +

39

40 def test_compile_config_repr_succeeds():

41 # setup: VllmBackend mutates the config object

@@ -504,8 +506,8 @@ def test_generation_config_loading():

506 @pytest.mark.parametrize(

507 "pt_load_map_location",

508 [

507 - "cuda",

508 - {"": "cuda"},

509 + DEVICE_TYPE,

510 + {"": DEVICE_TYPE},

511 ],

512 )

513 def test_load_config_pt_load_map_location(pt_load_map_location):

tests/v1/sample/test_topk_topp_sampler.py

+1 -1

@@ -127,7 +127,7 @@ def test_flashinfer_sampler():

127 # =============================================================================

128

129

130 -@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU 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

tests/v1/test_tensor_ipc_queue.py

+6 -3

@@ -14,6 +14,7 @@ import pytest

14 import torch

15 import torch.multiprocessing as torch_mp

16

17 +from vllm.platforms import current_platform

18 from vllm.v1.engine.tensor_ipc import (

19 TensorIpcData,

20 TensorIpcReceiver,

@@ -21,6 +22,8 @@ from vllm.v1.engine.tensor_ipc import (

22 )

23 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

24

25 +DEVICE_TYPE = current_platform.device_type

26 +

27

28 @pytest.fixture(scope="module", autouse=True)

29 def setup_multiprocessing():

@@ -53,7 +56,7 @@ def encoder_process(

56 encoder = MsgpackEncoder(oob_tensor_consumer=sender)

57

58 if torch.cuda.is_available():

56 - device = "cuda:0"

59 + device = f"{DEVICE_TYPE}:0"

60 tensor = torch.randn(

61 *tensor_data["shape"], dtype=tensor_data["dtype"], device=device

62 )

@@ -384,7 +387,7 @@ def mixed_tensor_encoder_process(

387

388 # Create only CUDA tensor for IPC (CPU will be serialized)

389 # But actually, let's just send CUDA tensor directly

387 - cuda_tensor = torch.randn(4, 5, device="cuda:0")

390 + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")

391

392 # Manually send via IPC to test the mechanism

393 cuda_tensor_shared = cuda_tensor.share_memory_()

@@ -651,7 +654,7 @@ def test_ipc_disabled_mode():

654

655 # If CUDA is available, test with CUDA tensor too

656 if torch.cuda.is_available():

654 - cuda_tensor = torch.randn(4, 5, device="cuda:0")

657 + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0")

658 encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor})

659 assert len(encoded_cuda) > 0

660 assert tensor_queues[0].empty(), (

[Test] Refactor hard coded device string in test files under compile/quantization/models/model_executor folders (#38901)