vllm.models.kimi_k3.amd.specialized_decode.model_state

Kimi-K3's model runner state on ROCm: 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 prepares it

KimiK3ModelState

Bases: MambaHybridModelState

Builds the specialized decode once the weights are loaded, and prepares it for every bound KV cache (the CUDA-graph memory-profiling cache, then the real one).

Source code in vllm/models/kimi_k3/amd/specialized_decode/model_state.py

class KimiK3ModelState(MambaHybridModelState):
    """Builds the specialized decode once the weights are loaded, and prepares it
    for every bound KV cache (the CUDA-graph memory-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()