Skip to content

vllm.models.kimi_k3.nvidia.specialized_decode.runtime

Kimi K3 specialized decode: a CuTe decode executor over the stock backbone, enabled by additional_config={"decode_backend": "latency_specialized"}.

Modules:

  • decode –

    The specialized decode executor over the stock Kimi K3 backbone: one routing

  • prefetch –

    L2 weight prefetch for the specialized decode.

  • providers –

    Per-layer attention providers (input projection, KDA/MLA core, o-projection) and

Functions:

  • kda_dtypes –

    (conv, recurrent) KDA state dtypes the deployment's cache config selects.

  • specialized_decode_config –

    The deployment's specialized decode settings, None for the stock decode.

kda_dtypes(vllm_config)

(conv, recurrent) KDA state dtypes the deployment's cache config selects.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py 
43
44
45
46
47
48
49
50
51
52
def kda_dtypes(vllm_config):
    """(conv, recurrent) KDA state dtypes the deployment's cache config selects."""
    from vllm.model_executor.layers.mamba.mamba_utils import MambaStateDtypeCalculator

    cache = vllm_config.cache_config
    return MambaStateDtypeCalculator.kda_state_dtype(
        vllm_config.model_config.dtype,
        cache.mamba_cache_dtype,
        cache.mamba_ssm_cache_dtype,
    )

specialized_decode_config(vllm_config)

The deployment's specialized decode settings, None for the stock decode. Raises when the specialized decode is requested for an unsupported deployment.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def specialized_decode_config(vllm_config) -> DecodeConfig | None:
    """The deployment's specialized decode settings, None for the stock decode.
    Raises when the specialized decode is requested for an unsupported deployment."""
    extra = vllm_config.additional_config
    backend = extra.get("decode_backend", "auto") if isinstance(extra, dict) else "auto"
    if backend == "auto":
        return None
    if backend != "latency_specialized":
        raise ValueError("Supported decode_backend values are auto/latency_specialized")
    reason = unsupported(vllm_config)
    if reason is not None:
        raise ValueError(f"Kimi K3 specialized decode does not support {reason}")
    spec = vllm_config.speculative_config
    return DecodeConfig(
        max_context=vllm_config.model_config.max_model_len,
        mla_core=extra.get("mla_core", "auto"),
        dcp=vllm_config.parallel_config.decode_context_parallel_size,
        q=1 + spec.num_speculative_tokens if spec is not None else 1,
        capture_sizes=tuple(
            vllm_config.compilation_config.cudagraph_capture_sizes or ()
        ),
    )