__init__.py 863 Bytes
Newer Older
chenzk's avatar
chenzk committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared helpers: Triton compat, layout bridge, context, sequences."""

from vllm.kvprune.utils.layout_bridge import (
    block_table_to_global_page_table,
    build_batch_mapping,
    build_page_table_head_major,
    flatten_kv_cache_head_major,
    flatten_kv_cache_plane,
    write_head_major_flat_to_interleaved,
)
from vllm.kvprune.utils.triton_compat import (
    autotune as triton_autotune,
    cuda_capability_geq,
    maybe_set_allocator,
)

__all__ = [
    "block_table_to_global_page_table",
    "build_batch_mapping",
    "build_page_table_head_major",
    "cuda_capability_geq",
    "flatten_kv_cache_head_major",
    "flatten_kv_cache_plane",
    "write_head_major_flat_to_interleaved",
    "maybe_set_allocator",
    "triton_autotune",
]