mirror of
https://github.com/vllm-project/vllm.git
synced 2026-10-28 09:45:26 +08:00
vhagorandIsotr0py
3dc7a68ce4
Signed-off-by: Juqi Li <2223621784@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
15 lines
508 B
Python
15 lines
508 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.warmup.mamba_triton_warmup import _warm_batch_memcpy_kernel
|
|
from vllm.platforms import current_platform
|
|
|
|
|
|
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA is required")
|
|
def test_mamba_batch_memcpy_kernel_compiles_on_gpu() -> None:
|
|
_warm_batch_memcpy_kernel(torch.device("cuda"))
|
|
torch.accelerator.synchronize("cuda")
|