vllm.models.kimi_k3.nvidia.model_state

Kimi K3's model runner state: the hybrid (KDA + MLA) state, plus the lifecycle of the specialized decode when decode_backend=latency_specialized selects it.

Classes:

• KimiK3ModelState –

  Builds the specialized decode once the weights are loaded and the aux layers

KimiK3ModelState

Bases: MambaHybridModelState

Builds the specialized decode once the weights are loaded and the aux layers set, and prepares it for every bound KV cache (the profiling cache, then the real one).

Source code in vllm/models/kimi_k3/nvidia/model_state.py

class KimiK3ModelState(MambaHybridModelState):
    """Builds the specialized decode once the weights are loaded and the aux layers
    set, and prepares it for every bound KV cache (the profiling cache, then the
    real one)."""

    def __init__(self, vllm_config, model, encoder_cache, device):
        super().__init__(vllm_config, model, encoder_cache, device)
        if hasattr(model, "get_language_model"):
            model = model.get_language_model()
        self.language_model = model
        if model.specialized_decode_config is not None:
            model.build_specialized_decode()

    def initialize_kv_cache(self, kv_cache_config, block_tables):
        super().initialize_kv_cache(kv_cache_config, block_tables)
        decode = self.language_model.specialized_decode
        if decode is not None:
            decode.prepare_kv()