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

• dep –

  The specialized decode at DEP (data-parallel attention at TP1 + expert-parallel

• hybrid –

  The specialized decode of the hybrid decomposition: attention (KDA, DSA) and

• prefetch –

  L2 weight prefetch for the specialized decode.

• providers –

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

Functions:

• dsa_fi_kernels –

  The FlashInfer sparse MLA kernels of the DSA layers on a per-tensor FP8 or

• dsa_unsupported –

  Why the K3.1 DSA layers' specialized kernels cannot run: the query and index

• kda_dtypes –

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

• prims_ts_rows –

  The prims_ts plans' row capacities: the deployment's CUDA graph sizes and the

• specialized_decode_config –

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

dsa_fi_kernels(vllm_config)

The FlashInfer sparse MLA kernels of the DSA layers on a per-tensor FP8 or nvfp4_ds_mla cache (none for other models and caches).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py

def dsa_fi_kernels(vllm_config) -> tuple[str, ...]:
    """The FlashInfer sparse MLA kernels of the DSA layers on a per-tensor FP8 or
    nvfp4_ds_mla cache (none for other models and caches)."""
    text = vllm_config.model_config.hf_text_config
    dtype = vllm_config.cache_config.cache_dtype
    if getattr(text, "dsa_config", None) is None or dtype not in FI_CACHE_DTYPES:
        return ()
    from ..cutedsl.dsa.attn_fi import fi_kernels

    return fi_kernels(dtype)

dsa_unsupported(vllm_config)

Why the K3.1 DSA layers' specialized kernels cannot run: the query and index fronts implement the K3.1 indexer (no indexer RoPE, LayerNorm keys; an FP8 wq_index is refused when the layer is built), and on a FlashInfer sparse cache the per-batch kernel policy needs FlashInfer's prims_ts kernel.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py

def dsa_unsupported(vllm_config) -> str | None:
    """Why the K3.1 DSA layers' specialized kernels cannot run: the query and index
    fronts implement the K3.1 indexer (no indexer RoPE, LayerNorm keys; an FP8
    wq_index is refused when the layer is built), and on a FlashInfer sparse cache
    the per-batch kernel policy needs FlashInfer's prims_ts kernel."""
    dsa = vllm_config.model_config.hf_text_config.dsa_config
    if getattr(dsa, "index_rope_dim", 0) != 0:
        return "DSA indexer RoPE (index_rope_dim > 0)"
    if getattr(dsa, "index_k_norm_type", None) != "layernorm":
        return "DSA indexer keys normalized other than by LayerNorm"
    dtype = vllm_config.cache_config.cache_dtype
    if dtype in FI_CACHE_DTYPES:
        return f"DSA layers on a {dtype} cache without FlashInfer prims_ts"
    return None

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

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,
    )

prims_ts_rows(vllm_config)

The prims_ts plans' row capacities: the deployment's CUDA graph sizes and the powers of two up to its widest call (FLASHINFER_MLA_SPARSE's buckets).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/runtime/__init__.py

def prims_ts_rows(vllm_config) -> tuple[int, ...]:
    """The prims_ts plans' row capacities: the deployment's CUDA graph sizes and the
    powers of two up to its widest call (FLASHINFER_MLA_SPARSE's buckets)."""
    # FLASHINFER_MLA_SPARSE does not run prims_ts yet.
    return ()

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

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
    if spec is not None:
        # The draft model loads after this target: build it as the specialized
        # drafter (its own module, ``drafter.model``).
        from ..drafter.model import register

        register()
    return DecodeConfig(
        max_context=vllm_config.model_config.max_model_len,
        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 ()
        ),
        fi_kernels=dsa_fi_kernels(vllm_config),
        prims_ts_rows=prims_ts_rows(vllm_config),
        dsa=getattr(vllm_config.model_config.hf_text_config, "dsa_config", None)
        is not None,
    )