| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.warmup.mamba_triton_warmup import _warm_batch_memcpy_kernel |
| 8 | from vllm.platforms import current_platform |
| 9 | |
| 10 | |
| 11 | @pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required") |
| 12 | def test_mamba_batch_memcpy_kernel_compiles_on_gpu() -> None: |
| 13 | _warm_batch_memcpy_kernel(torch.device("cuda")) |
| 14 | torch.accelerator.synchronize("cuda") |