Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
23 lines
618 B
Python
23 lines
618 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Sparse attention indexer (kpool) — hardware-isolated entry point.
|
|
|
|
The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers
|
|
in ``common/``); this module picks the right one for the current platform.
|
|
"""
|
|
|
|
from vllm.platforms import current_platform
|
|
|
|
if current_platform.is_rocm():
|
|
from .amd.sparse_indexer import (
|
|
SparseAttnIndexerKpool,
|
|
)
|
|
else:
|
|
from .nvidia.sparse_indexer import (
|
|
SparseAttnIndexerKpool,
|
|
)
|
|
|
|
__all__ = [
|
|
"SparseAttnIndexerKpool",
|
|
]
|