vllm.model_executor.models.config ¶

class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
    """Apply the attention contract declared by a ColQwen3.5 checkpoint."""

    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
        configs = {
            id(config): config
            for config in (
                model_config.hf_config,
                model_config.hf_text_config,
            )
        }
        declarations = [
            contract
            for config in configs.values()
            if (contract := getattr(config, "retrieval_attention_contract", None))
            is not None
        ]
        supported = {"causal", "bidirectional"}
        if not declarations:
            raise ValueError(
                "ColQwen3.5 checkpoints must declare "
                "retrieval_attention_contract as 'causal' or 'bidirectional'"
            )
        if (
            any(not isinstance(contract, str) for contract in declarations)
            or any(contract not in supported for contract in declarations)
            or len(set(declarations)) != 1
        ):
            raise ValueError(
                "unsupported or conflicting ColQwen3.5 "
                f"retrieval_attention_contract declarations: {declarations!r}"
            )
        is_causal = declarations[0] == "causal"
        for config in configs.values():
            config.is_causal = is_causal

class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig):
    @classmethod
    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
        """Set up the diffusion config and defaults for DiffusionGemma.

        Auto-creates DiffusionConfig from the HF config when the user
        didn't pass ``--diffusion-config``. Diffusion sampling params are
        read straight from generation_config.json at sampler-build time
        (see DiffusionGemma's custom_sampler), not injected here.
        """
        # Inherit Gemma4's attention backend selection (FA4 on Hopper,
        # TRITON_ATTN fallback for heterogeneous head dims).
        Gemma4Config.verify_and_update_config(vllm_config)

        from vllm.v1.attention.backends.registry import AttentionBackendEnum

        attention_config = vllm_config.attention_config
        if attention_config.backend == AttentionBackendEnum.FLASHINFER:
            raise ValueError(
                "FlashInfer does not support DiffusionGemma's mixed "
                "causal/bidirectional attention. Use --attention-backend "
                "FLASH_ATTN or TRITON_ATTN instead."
            )
        if attention_config.backend is None and not attention_config.use_non_causal:
            attention_config.use_non_causal = True
            logger.info(
                "DiffusionGemma uses mixed causal/bidirectional attention "
                "within a batch; setting use_non_causal=True to exclude "
                "FlashInfer from auto-selection."
            )

        # Auto-create DiffusionConfig from HF config if not provided.
        if vllm_config.diffusion_config is None:
            from vllm.config.diffusion import DiffusionConfig

            hf_config = vllm_config.model_config.hf_config
            canvas_length = getattr(hf_config, "canvas_length", 256)
            vllm_config.diffusion_config = DiffusionConfig(
                canvas_length=canvas_length,
            )

        # The diffusion sampler materializes [num_seqs, canvas_length, vocab]
        # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200).
        # Default to 8 when the user didn't pass --max-num-seqs.
        # We can't see the original None here (the engine already filled a generic
        # default), so use >= DEFAULT_MAX_NUM_SEQS as a proxy, (the default is much
        # larger than any deliberate value for this model)
        from vllm.config.scheduler import SchedulerConfig

        sc = vllm_config.scheduler_config
        if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS:
            sc.max_num_seqs = 8

        # Remove the model's generation_config.json cap on max_new_tokens
        # (256) so DiffusionGemma behaves like every other model: no
        # server-wide limit, each request controls its own output length
        # via max_tokens.  Setting to None causes get_diff_sampling_param
        # to skip this key entirely.
        model_config = vllm_config.model_config
        if "max_new_tokens" not in model_config.override_generation_config:
            model_config.override_generation_config["max_new_tokens"] = None
            logger.info(
                "DiffusionGemma: removing server-wide max_new_tokens cap "
                "from generation_config.json (use "
                "--override-generation-config to set a custom limit).",
            )

class Gemma4Config(VerifyAndUpdateConfig):
    @staticmethod
    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
        """Configure attention for heterogeneous head dimensions.

        Gemma4 uses different head dimensions for sliding window vs full attention
        layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes
        mixed backend selection and numerical divergence.

        When FA4 is available we force it for ALL layers, giving a uniform kernel path
        and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back
        to Triton.
        """
        model_config = vllm_config.model_config
        arch_config = model_config.model_arch_config
        layer_types = getattr(model_config.hf_text_config, "layer_types", None) or []
        head_dims = {
            layer_types[i]: arch_config[i].head_size
            for i in range(min(arch_config.total_num_hidden_layers, len(layer_types)))
        }

        if len(set(head_dims.values())) <= 1:
            return

        from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
        from vllm.v1.attention.backends.registry import AttentionBackendEnum

        max_head_dim = max(head_dims.values())

        if is_fa_version_supported(4) and max_head_dim <= 512:
            if (
                vllm_config.attention_config.flash_attn_version is None
                and vllm_config.attention_config.backend
                in (None, AttentionBackendEnum.FLASH_ATTN)
            ):
                vllm_config.attention_config.flash_attn_version = 4
                logger.info(
                    "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "
                    "all layers to avoid mixed FA3/FA4 penalty.",
                    head_dims,
                )
        elif vllm_config.attention_config.backend is None:
            vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
            logger.info(
                "Gemma4 model has heterogeneous head dimensions "
                "%s. FA4 not available, forcing TRITON_ATTN backend.",
                head_dims,
            )

class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
    @classmethod
    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
        """
        Perform early validation and setup for hybrid attention/mamba models.

        Block size alignment with mamba page sizes is handled later by
        Platform.update_block_size_for_backend(), which runs after model
        layers are constructed and the attention backend is known.

        Args:
            vllm_config: vLLM Config
        """
        # Enable FULL_AND_PIECEWISE by default
        MambaModelConfig.verify_and_update_config(vllm_config)

class JinaEmbeddingsV5ModelConfig(VerifyAndUpdateConfig):
    """Config handler for Jina Embeddings V5 embedding models."""

    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
        """Enable the bidirectional encoder backbone for -nano checkpoints.

        The V5 family ships more than one backbone under a single
        `architectures` entry: the `-small` variants are Qwen3 decoders, while
        `-nano` is a bidirectional EuroBERT encoder. Upstream ships a separate
        `configuration_*.py` per repository, so the config carries no backbone
        field and the encoder variants are only identifiable by
        `is_decoder=False`. For encoder checkpoints, set `is_causal=False` so the
        Llama backbone uses EncoderOnlyAttention; `JinaEmbeddingsV5Model` then
        dispatches to its encoder implementation.
        """
        if getattr(model_config.hf_config, "is_decoder", True):
            return

        model_config.hf_config.is_causal = False

class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig):
    """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface.

    Kimi-K3 ships its routed experts as compressed-tensors
    ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which
    lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel
    selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config``
    (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4
    checkpoint is left untouched. Covers both the main model and the MTP draft.

    ``model_arch_config.quantization_config`` is a separate dict, snapshotted in
    ``ModelConfig.__init__`` before this hook runs, and it is what
    ``_verify_quantization`` reads when resolving the quant method. Patch it
    alongside the hf configs so the rewrite lands before resolution; otherwise
    the main model still resolves to compressed-tensors.
    """

    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
        for cfg in (
            model_config.hf_config,
            model_config.hf_text_config,
            model_config.model_arch_config,
        ):
            quant_config = getattr(cfg, "quantization_config", None)
            if (
                isinstance(quant_config, dict)
                and quant_config.get("quant_method") == "compressed-tensors"
                and quant_config.get("format") == "mxfp4-pack-quantized"
            ):
                quant_config["quant_method"] = "mxfp4"

class LlamaNemotronVLConfig(VerifyAndUpdateConfig):
    """Config handler for LlamaNemotronVL embedding models."""

    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
        from vllm.config.pooler import SequencePoolingType

        hf_config = model_config.hf_config

        # Set bidirectional attention on the language model config
        hf_config.is_causal = False
        if hasattr(hf_config, "llm_config"):
            hf_config.llm_config.is_causal = False

        if hasattr(hf_config, "vision_config"):
            hf_config.patch_size = hf_config.vision_config.patch_size

        # Set up pooling type
        pooling_type_map: dict[str, SequencePoolingType] = {
            "avg": "MEAN",
            "cls": "CLS",
            "last": "LAST",
        }

        # Get pooling type from config (check both top-level and llm_config)
        pooling = getattr(hf_config, "pooling", None)
        if pooling is None and hasattr(hf_config, "llm_config"):
            pooling = getattr(hf_config.llm_config, "pooling", "avg")

        assert isinstance(pooling, str)

        pooling_type = pooling_type_map.get(pooling)
        if pooling_type is None:
            raise ValueError(f"pool_type {pooling!r} not supported")

        assert model_config.pooler_config is not None
        model_config.pooler_config.seq_pooling_type = pooling_type

class MambaModelConfig(VerifyAndUpdateConfig):
    @classmethod
    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
        """
        Enable FULL_AND_PIECEWISE cuda graph mode by default (required
        to get good performance for mamba layers in V1).

        Args:
            vllm_config: vLLM Config
        """
        model_config = vllm_config.model_config
        cache_config = vllm_config.cache_config

        if cache_config.enable_prefix_caching:
            if cache_config.mamba_cache_mode == "none":
                cache_config.mamba_cache_mode = "align"
                logger.info(
                    "Mamba cache mode is set to '%s' for %s by default "
                    "when prefix caching is enabled",
                    cache_config.mamba_cache_mode,
                    model_config.architecture,
                )
            if (
                cache_config.mamba_cache_mode == "all"
                and not model_config.supports_mamba_prefix_caching
            ):
                cache_config.mamba_cache_mode = "align"
                logger.warning(
                    "Hybrid or mamba-based model detected without support "
                    "for prefix caching with Mamba cache 'all' mode: "
                    "falling back to 'align' mode."
                )
            if cache_config.mamba_cache_mode == "align":
                assert vllm_config.scheduler_config.enable_chunked_prefill, (
                    "Chunked prefill is required for mamba cache mode 'align'."
                )
            # By default, mamba block size will be set to max_model_len (see
            # below). When enabling prefix caching, we align mamba block size
            # to the block size as the basic granularity for prefix caching.
            if cache_config.mamba_block_size is None:
                cache_config.mamba_block_size = cache_config.block_size
        else:
            if cache_config.mamba_cache_mode != "none":
                cache_config.mamba_cache_mode = "none"
                logger.warning(
                    "Mamba cache mode is set to 'none' when prefix caching is disabled"
                )
            if cache_config.mamba_block_size is None:
                cache_config.mamba_block_size = model_config.max_model_len

class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
    DEFAULT_MAMBA_SSM_CACHE_DTYPE: "MambaDType" = "float32"
    """Only `float32` is known to have no accuracy issues by default."""

    @classmethod
    def update_mamba_ssm_cache_dtype(
        cls, *, cache_config: "CacheConfig", hf_config: "PretrainedConfig"
    ) -> None:
        """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto'
        (or not explicitly set), to the value specified in the HF config, or to
        `float32` if not specified.
        """
        if cache_config.mamba_ssm_cache_dtype == "auto":
            mamba_ssm_cache_dtype: MambaDType = getattr(
                hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE
            )
            logger.info(
                "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model",
                mamba_ssm_cache_dtype,
            )
            cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype

    @classmethod
    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
        cls.update_mamba_ssm_cache_dtype(
            cache_config=vllm_config.cache_config,
            hf_config=vllm_config.model_config.hf_config,
        )

class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
    @staticmethod
    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
        """Update mamba_ssm_cache_dtype for Qwen3.5 models when set to 'auto'
        (or not explicitly set), to the value specified in the HF config's
        mamba_ssm_dtype field. Warn if the user explicitly overrides it to a
        different value.
        """
        cache_config = vllm_config.cache_config
        hf_text_config = vllm_config.model_config.hf_text_config
        mamba_ssm_dtype = getattr(hf_text_config, "mamba_ssm_dtype", None)
        if cache_config.mamba_ssm_cache_dtype == "auto":
            if mamba_ssm_dtype is not None:
                cache_config.mamba_ssm_cache_dtype = mamba_ssm_dtype
        elif (
            mamba_ssm_dtype is not None
            and cache_config.mamba_ssm_cache_dtype != mamba_ssm_dtype
        ):
            logger.warning(
                "Qwen3.5 model specifies mamba_ssm_dtype='%s' in its config, "
                "but --mamba-ssm-cache-dtype='%s' was passed. "
                "Using the user-specified value.",
                mamba_ssm_dtype,
                cache_config.mamba_ssm_cache_dtype,
            )

class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig):
    @staticmethod
    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
        """Configure Unlimited-OCR attention backends for R-SWA and vision.

        Backend selection — controlled by the standard ``--attention-config``
        CLI argument (priority order):

          1. ``--attention-config '{"backend": "FLASH_ATTN"}'``
             → FA4 + rswa_mask_mod.  Exact token-level R-SWA.
               ``flash_attn_version`` is forced to 4 if not already set (R-SWA
               mask_mod requires FA4; FA3 cannot express it).  Raises if FA4 is
               not available on this device.

          2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'``
             → FlexAttention R-SWA via Triton block mask.

          3. ``--attention-config '{"backend": "TRITON_ATTN"}'``
             → Triton unified attention with an R-SWA decode mask.

          4. ``--attention-config '{"backend": "auto"}'`` (or omitted)
             → Auto-detect: FA4 if available (H20/H100 SM90), else TritonAttention.

        Regardless of backend, prefix caching is disabled for this model: R-SWA
        decode-phase KV is not a pure causal function of the prefix (so decode
        blocks are not reusable), and single-turn image-led OCR prompts rarely
        hit the prefix cache.

        Example — force FlexAttention even on a machine with FA4::

            vllm serve baidu/Unlimited-OCR \\
                --attention-config '{"backend": "FLEX_ATTENTION"}'
        """
        from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
        from vllm.v1.attention.backends.registry import AttentionBackendEnum

        attn_config = vllm_config.attention_config
        fa4_available = is_fa_version_supported(4)

        # ── step 1: resolve backend ─────────────────────────────────────────
        # None means the user did not explicitly specify a backend; auto-select.
        if attn_config.backend is None:
            attn_config.backend = (
                AttentionBackendEnum.FLASH_ATTN
                if fa4_available
                else AttentionBackendEnum.TRITON_ATTN
            )
            logger.info(
                "Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).",
                attn_config.backend.value,
                fa4_available,
            )

        # ── step 2: configure the chosen backend ────────────────────────────
        if attn_config.backend == AttentionBackendEnum.FLASH_ATTN:
            if not fa4_available:
                raise RuntimeError(
                    "Unlimited-OCR: --attention-config backend=FLASH_ATTN "
                    "requires FA4 (rswa_mask_mod), but FA4 is not available on "
                    "this device/installation. Use backend=TRITON_ATTN or "
                    "FLEX_ATTENTION, or upgrade vllm-flash-attn."
                )
            # On SM90 (H20), the default FA version is FA3 regardless of FA4
            # availability (FA4 is only auto-upgraded when head_size > 256).
            # The R-SWA mask_mod requires FA4, so force the version globally.
            if attn_config.flash_attn_version is None:
                attn_config.flash_attn_version = 4
            elif attn_config.flash_attn_version < 4:
                logger.warning(
                    "Unlimited-OCR: flash_attn_version=%d cannot express the "
                    "R-SWA mask_mod; upgrading to 4.",
                    attn_config.flash_attn_version,
                )
                attn_config.flash_attn_version = 4
            logger.info(
                "Unlimited-OCR: FlashAttention FA%d + rswa_mask_mod — exact R-SWA.",
                attn_config.flash_attn_version,
            )

        elif attn_config.backend == AttentionBackendEnum.TRITON_ATTN:
            logger.info(
                "Unlimited-OCR: TritonAttention — R-SWA via unified attention mask."
            )

        elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION:
            logger.info(
                "Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.",
                ""
                if not fa4_available
                else (
                    " (FA4 available but not used; pass backend=FLASH_ATTN to upgrade)"
                ),
            )

        else:
            raise ValueError(
                f"Unlimited-OCR: unsupported attention backend "
                f"{attn_config.backend!r} for R-SWA. "
                "Use FLASH_ATTN (FA4), TRITON_ATTN or FLEX_ATTENTION."
            )

        # R-SWA windows the *generated* tokens, so a decode-token's KV is not a
        # pure causal function of the prefix and cannot be safely reused across
        # requests via prefix caching. Only the prompt/image prefix is cacheable,
        # but OCR is single-turn with image-led prompts that rarely share a
        # prefix, so prefix caching brings little benefit while complicating the
        # KV cache manager. Disable it for this model.
        cache_config = vllm_config.cache_config
        if cache_config.enable_prefix_caching:
            cache_config.enable_prefix_caching = False
            logger.info(
                "Unlimited-OCR: disabling prefix caching (R-SWA decode KV is not "
                "cacheable, and single-turn image-led prompts rarely hit the "
                "prefix cache)."
            )

        mm_config = getattr(vllm_config.model_config, "multimodal_config", None)
        if mm_config is not None:
            if mm_config.mm_encoder_attn_backend is None:
                mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN
            elif mm_config.mm_encoder_attn_backend == AttentionBackendEnum.FLASHINFER:
                logger.warning(
                    "Unlimited-OCR: FlashInfer is not supported for the vision "
                    "encoder (the CLIP stage runs full attention without "
                    "cu_seqlens); falling back to FlashAttention."
                )
                mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN

    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
        text_config = model_config.hf_config.text_config
        text_config.architectures = ["DeepseekV2ForCausalLM"]
        if getattr(model_config.hf_config, "rswa_window", None) is None:
            model_config.hf_config.rswa_window = 128
        # Propagate rswa_window to text_config so that DeepseekAttention (which
        # receives text_config as its vllm_config.model_config.hf_config via
        # init_vllm_registered_model) can read it and create RSWAAttention.
        rswa_window = model_config.hf_config.rswa_window
        text_config.rswa_window = rswa_window