dotcache/backends/__init__.py

from .cpu_ref import mix_page_cpu_ref, score_page_cpu_ref
from .torch_cuda import (
    cuda_available,
    decode_grouped_multiquery_step_cuda_tensor,
    decode_grouped_multiquery_step_prepared_cuda_tensor,
    decode_grouped_multiquery_step_prepared_cuda_tensor_output_only,
    decode_multi_query_step_cuda,
    decode_multi_query_step_cuda_tensor,
    decode_step_cuda,
    mix_page_cuda,
    page_supported_cuda,
    prepare_page_cuda,
    prepare_pages_cuda,
    PreparedPageCUDA,
    score_page_cuda,
    score_pages_cuda,
)
from .torch_mps import (
    clear_prepared_chunk_cache,
    configure_prepared_chunk_cache,
    decode_grouped_multiquery_step_prepared_mps_tensor,
    decode_grouped_multiquery_step_mps_tensor,
    decode_grouped_multiquery_step_prepared_torch_tensor,
    decode_grouped_multiquery_step_prepared_torch_tensor_output_only,
    decode_grouped_multiquery_step_torch_tensor,
    decode_multi_query_step_torch,
    decode_multi_query_step_torch_tensor,
    decode_step_torch,
    PreparedPageMPS,
    PreparedPageTorch,
    decode_multi_query_step_mps,
    decode_multi_query_step_mps_tensor,
    decode_step_mps,
    mix_page_mps,
    mps_available,
    page_supported_mps,
    page_supported_torch,
    prepare_m0_affine_pages_from_tensor_torch,
    prepare_page_mps,
    prepare_page_torch,
    prepare_pages_mps,
    prepare_pages_torch,
    prepared_chunk_cache_resident_bytes,
    score_pages_mps,
    score_pages_torch,
    score_page_mps,
    score_page_torch,
    set_prepared_chunk_cache_budget_override,
    torch_device_available,
    mix_page_torch,
)

__all__ = [
    "clear_prepared_chunk_cache",
    "configure_prepared_chunk_cache",
    "cuda_available",
    "decode_grouped_multiquery_step_cuda_tensor",
    "decode_grouped_multiquery_step_prepared_cuda_tensor",
    "decode_grouped_multiquery_step_prepared_cuda_tensor_output_only",
    "PreparedPageMPS",
    "PreparedPageCUDA",
    "PreparedPageTorch",
    "decode_grouped_multiquery_step_prepared_torch_tensor",
    "decode_grouped_multiquery_step_prepared_torch_tensor_output_only",
    "decode_grouped_multiquery_step_prepared_mps_tensor",
    "decode_grouped_multiquery_step_torch_tensor",
    "decode_grouped_multiquery_step_mps_tensor",
    "decode_multi_query_step_cuda",
    "decode_multi_query_step_cuda_tensor",
    "decode_multi_query_step_torch",
    "decode_multi_query_step_torch_tensor",
    "decode_multi_query_step_mps",
    "decode_multi_query_step_mps_tensor",
    "decode_step_cuda",
    "decode_step_torch",
    "decode_step_mps",
    "mix_page_cpu_ref",
    "mix_page_cuda",
    "mix_page_torch",
    "mix_page_mps",
    "mps_available",
    "page_supported_cuda",
    "page_supported_mps",
    "page_supported_torch",
    "prepare_m0_affine_pages_from_tensor_torch",
    "prepare_page_cuda",
    "prepare_page_torch",
    "prepare_page_mps",
    "prepare_pages_cuda",
    "prepare_pages_torch",
    "prepare_pages_mps",
    "prepared_chunk_cache_resident_bytes",
    "score_pages_cuda",
    "score_pages_torch",
    "score_pages_mps",
    "score_page_cpu_ref",
    "score_page_cuda",
    "score_page_torch",
    "score_page_mps",
    "set_prepared_chunk_cache_budget_override",
    "torch_device_available",
]