main
py 22 lines 618 Bytes
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Sparse attention indexer (kpool) — hardware-isolated entry point.
4
5 The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers
6 in ``common/``); this module picks the right one for the current platform.
7 """
8
9 from vllm.platforms import current_platform
10
11 if current_platform.is_rocm():
12 from .amd.sparse_indexer import (
13 SparseAttnIndexerKpool,
14 )
15 else:
16 from .nvidia.sparse_indexer import (
17 SparseAttnIndexerKpool,
18 )
19
20 __all__ = [
21 "SparseAttnIndexerKpool",
22 ]