vllm/tests/kernels/conftest.py
Lily Liu 43c413ec57
[Kernel] Use flashinfer for decoding (#4353)
Co-authored-by: LiuXiaoxuanPKU <llilyliupku@gmail.com>
2024-05-03 15:51:27 -07:00

15 lines
314 B
Python

import pytest
from vllm.utils import (create_kv_caches_with_random,
create_kv_caches_with_random_flash)
@pytest.fixture()
def kv_cache_factory():
return create_kv_caches_with_random
@pytest.fixture()
def kv_cache_factory_flashinfer():
return create_kv_caches_with_random_flash