__init__.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging

from vllm.kvprune.models.llama3 import LlamaForCausalLM
from vllm.kvprune.models.qwen3 import Qwen3ForCausalLM

logger = logging.getLogger(__name__)

MODEL_REGISTRY = {
    "llama": LlamaForCausalLM,
    "qwen3": Qwen3ForCausalLM,
}

try:
    from vllm.kvprune.models.qwen3_moe import Qwen3MoeForCausalLM
except Exception as exc:
    logger.debug("Skipping qwen3_moe registration due to import error: %s", exc)
else:
    MODEL_REGISTRY["qwen3_moe"] = Qwen3MoeForCausalLM