| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi K3 model — hardware-isolated entry point. |
| 4 | |
| 5 | The implementation lives under ``nvidia/`` and ``amd/``; this module picks the |
| 6 | right one for the current platform and re-exports the public classes used by |
| 7 | the model registry. (Mirrors ``vllm.models.minimax_m3``.) |
| 8 | """ |
| 9 | |
| 10 | from typing import TYPE_CHECKING |
| 11 | |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | # The NVIDIA branch is the static default that type-checkers see; the ROCm |
| 15 | # branch overrides it at runtime (kept type-compatible via type: ignore). |
| 16 | # TPU plugins import the shared ``common`` modules through this package, but |
| 17 | # register their own model classes. Do not eagerly import a GPU implementation. |
| 18 | if TYPE_CHECKING: |
| 19 | from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM |
| 20 | from .nvidia.mtp import KimiK3MTP |
| 21 | elif current_platform.device_type == "tpu": |
| 22 | pass |
| 23 | elif current_platform.is_rocm(): |
| 24 | from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment] |
| 25 | from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment] |
| 26 | from .amd.mtp import KimiK3MTP # type: ignore[assignment] |
| 27 | else: |
| 28 | from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM |
| 29 | from .nvidia.mtp import KimiK3MTP |
| 30 | |
| 31 | __all__ = [ |
| 32 | "KimiK3ForConditionalGeneration", |
| 33 | "KimiK3MTP", |
| 34 | "KimiLinearForCausalLM", |
| 35 | ] |