[K3] Allow tpu to import kimi_k3.common (#51529)
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Jeff (Junze) Ma committed
Aug 9, 2026 at 02:06 UTC
04d13b5d653725651f7a750f0022f06d095881dc
1 file changed
+9
-2
vllm/models/kimi_k3/__init__.py
+9
-2
@@ -13,13 +13,20 @@ 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
-if TYPE_CHECKING or not current_platform.is_rocm():
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
19
-else:
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",