| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TYPE_CHECKING |
| 4 | |
| 5 | from vllm.logger import init_logger |
| 6 | from vllm.utils.math_utils import round_up |
| 7 | |
| 8 | if TYPE_CHECKING: |
| 9 | from transformers import PreTrainedConfig |
| 10 | |
| 11 | from vllm.config import CacheConfig, ModelConfig, VllmConfig |
| 12 | from vllm.config.cache import MambaDType |
| 13 | |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | class VerifyAndUpdateConfig: |
| 19 | @staticmethod |
| 20 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 21 | return |
| 22 | |
| 23 | @staticmethod |
| 24 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 25 | return |
| 26 | |
| 27 | |
| 28 | class DeepseekV32ForCausalLM(VerifyAndUpdateConfig): |
| 29 | @classmethod |
| 30 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 31 | hf_config = vllm_config.model_config.hf_config |
| 32 | |
| 33 | # Mirror the check in vllm/model_executor/models/deepseek_v2.py |
| 34 | is_v32 = hasattr(hf_config, "index_topk") |
| 35 | assert is_v32 |
| 36 | |
| 37 | cache_config = vllm_config.cache_config |
| 38 | if cache_config.cache_dtype == "bfloat16": |
| 39 | cache_config.cache_dtype = "auto" |
| 40 | logger.info("Using bfloat16 kv-cache for DeepSeekV3.2") |
| 41 | |
| 42 | |
| 43 | class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig): |
| 44 | @staticmethod |
| 45 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 46 | # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs |
| 47 | # in most cases. |
| 48 | vllm_config.parallel_config.set_dcp_defaults( |
| 49 | comm_backend="a2a", q_replicate=True |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig): |
| 54 | @staticmethod |
| 55 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 56 | # Ernie4.5-VL conditionally executes text/vision MoE branches, so |
| 57 | # fast_moe_cold_start can silently produce incorrect execution order. |
| 58 | vllm_config.compilation_config.fast_moe_cold_start = False |
| 59 | |
| 60 | |
| 61 | class Gemma3TextModelConfig(VerifyAndUpdateConfig): |
| 62 | @staticmethod |
| 63 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 64 | hf_config = model_config.hf_config |
| 65 | hf_config.is_causal = not hf_config.use_bidirectional_attention |
| 66 | |
| 67 | |
| 68 | class UnlimitedOCRForCausalLMConfig(VerifyAndUpdateConfig): |
| 69 | @staticmethod |
| 70 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 71 | """Configure Unlimited-OCR attention backends for R-SWA and vision. |
| 72 | |
| 73 | Backend selection — controlled by the standard ``--attention-config`` |
| 74 | CLI argument (priority order): |
| 75 | |
| 76 | 1. ``--attention-config '{"backend": "FLASH_ATTN"}'`` |
| 77 | → FA4 + rswa_mask_mod. Exact token-level R-SWA. |
| 78 | ``flash_attn_version`` is forced to 4 if not already set (R-SWA |
| 79 | mask_mod requires FA4; FA3 cannot express it). Raises if FA4 is |
| 80 | not available on this device. |
| 81 | |
| 82 | 2. ``--attention-config '{"backend": "FLEX_ATTENTION"}'`` |
| 83 | → FlexAttention R-SWA via Triton block mask. |
| 84 | |
| 85 | 3. ``--attention-config '{"backend": "TRITON_ATTN"}'`` |
| 86 | → Triton unified attention with an R-SWA decode mask. |
| 87 | |
| 88 | 4. ``--attention-config '{"backend": "auto"}'`` (or omitted) |
| 89 | → Auto-detect: FA4 if available (H20/H100 SM90), else TritonAttention. |
| 90 | |
| 91 | Regardless of backend, prefix caching is disabled for this model: R-SWA |
| 92 | decode-phase KV is not a pure causal function of the prefix (so decode |
| 93 | blocks are not reusable), and single-turn image-led OCR prompts rarely |
| 94 | hit the prefix cache. |
| 95 | |
| 96 | Example — force FlexAttention even on a machine with FA4:: |
| 97 | |
| 98 | vllm serve baidu/Unlimited-OCR \\ |
| 99 | --attention-config '{"backend": "FLEX_ATTENTION"}' |
| 100 | """ |
| 101 | from vllm.v1.attention.backends.fa_utils import is_fa_version_supported |
| 102 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 103 | |
| 104 | attn_config = vllm_config.attention_config |
| 105 | fa4_available = is_fa_version_supported(4) |
| 106 | |
| 107 | # ── step 1: resolve backend ───────────────────────────────────────── |
| 108 | # None means the user did not explicitly specify a backend; auto-select. |
| 109 | if attn_config.backend is None: |
| 110 | attn_config.backend = ( |
| 111 | AttentionBackendEnum.FLASH_ATTN |
| 112 | if fa4_available |
| 113 | else AttentionBackendEnum.TRITON_ATTN |
| 114 | ) |
| 115 | logger.info( |
| 116 | "Unlimited-OCR: auto-selected attention backend=%s (fa4_available=%s).", |
| 117 | attn_config.backend.value, |
| 118 | fa4_available, |
| 119 | ) |
| 120 | |
| 121 | # ── step 2: configure the chosen backend ──────────────────────────── |
| 122 | if attn_config.backend == AttentionBackendEnum.FLASH_ATTN: |
| 123 | if not fa4_available: |
| 124 | raise RuntimeError( |
| 125 | "Unlimited-OCR: --attention-config backend=FLASH_ATTN " |
| 126 | "requires FA4 (rswa_mask_mod), but FA4 is not available on " |
| 127 | "this device/installation. Use backend=TRITON_ATTN or " |
| 128 | "FLEX_ATTENTION, or upgrade vllm-flash-attn." |
| 129 | ) |
| 130 | # On SM90 (H20), the default FA version is FA3 regardless of FA4 |
| 131 | # availability (FA4 is only auto-upgraded when head_size > 256). |
| 132 | # The R-SWA mask_mod requires FA4, so force the version globally. |
| 133 | if attn_config.flash_attn_version is None: |
| 134 | attn_config.flash_attn_version = 4 |
| 135 | elif attn_config.flash_attn_version < 4: |
| 136 | logger.warning( |
| 137 | "Unlimited-OCR: flash_attn_version=%d cannot express the " |
| 138 | "R-SWA mask_mod; upgrading to 4.", |
| 139 | attn_config.flash_attn_version, |
| 140 | ) |
| 141 | attn_config.flash_attn_version = 4 |
| 142 | logger.info( |
| 143 | "Unlimited-OCR: FlashAttention FA%d + rswa_mask_mod — exact R-SWA.", |
| 144 | attn_config.flash_attn_version, |
| 145 | ) |
| 146 | |
| 147 | elif attn_config.backend == AttentionBackendEnum.TRITON_ATTN: |
| 148 | logger.info( |
| 149 | "Unlimited-OCR: TritonAttention — R-SWA via unified attention mask." |
| 150 | ) |
| 151 | |
| 152 | elif attn_config.backend == AttentionBackendEnum.FLEX_ATTENTION: |
| 153 | logger.info( |
| 154 | "Unlimited-OCR: FlexAttention — R-SWA via Triton block mask%s.", |
| 155 | "" |
| 156 | if not fa4_available |
| 157 | else ( |
| 158 | " (FA4 available but not used; pass backend=FLASH_ATTN to upgrade)" |
| 159 | ), |
| 160 | ) |
| 161 | |
| 162 | else: |
| 163 | raise ValueError( |
| 164 | f"Unlimited-OCR: unsupported attention backend " |
| 165 | f"{attn_config.backend!r} for R-SWA. " |
| 166 | "Use FLASH_ATTN (FA4), TRITON_ATTN or FLEX_ATTENTION." |
| 167 | ) |
| 168 | |
| 169 | # R-SWA windows the *generated* tokens, so a decode-token's KV is not a |
| 170 | # pure causal function of the prefix and cannot be safely reused across |
| 171 | # requests via prefix caching. Only the prompt/image prefix is cacheable, |
| 172 | # but OCR is single-turn with image-led prompts that rarely share a |
| 173 | # prefix, so prefix caching brings little benefit while complicating the |
| 174 | # KV cache manager. Disable it for this model. |
| 175 | cache_config = vllm_config.cache_config |
| 176 | if cache_config.enable_prefix_caching: |
| 177 | cache_config.enable_prefix_caching = False |
| 178 | logger.info( |
| 179 | "Unlimited-OCR: disabling prefix caching (R-SWA decode KV is not " |
| 180 | "cacheable, and single-turn image-led prompts rarely hit the " |
| 181 | "prefix cache)." |
| 182 | ) |
| 183 | |
| 184 | mm_config = getattr(vllm_config.model_config, "multimodal_config", None) |
| 185 | if mm_config is not None: |
| 186 | if mm_config.mm_encoder_attn_backend is None: |
| 187 | mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN |
| 188 | elif mm_config.mm_encoder_attn_backend == AttentionBackendEnum.FLASHINFER: |
| 189 | logger.warning( |
| 190 | "Unlimited-OCR: FlashInfer is not supported for the vision " |
| 191 | "encoder (the CLIP stage runs full attention without " |
| 192 | "cu_seqlens); falling back to FlashAttention." |
| 193 | ) |
| 194 | mm_config.mm_encoder_attn_backend = AttentionBackendEnum.FLASH_ATTN |
| 195 | |
| 196 | @staticmethod |
| 197 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 198 | text_config = model_config.hf_config.text_config |
| 199 | text_config.architectures = ["DeepseekV2ForCausalLM"] |
| 200 | if getattr(model_config.hf_config, "rswa_window", None) is None: |
| 201 | model_config.hf_config.rswa_window = 128 |
| 202 | # Propagate rswa_window to text_config so that DeepseekAttention (which |
| 203 | # receives text_config as its vllm_config.model_config.hf_config via |
| 204 | # init_vllm_registered_model) can read it and create RSWAAttention. |
| 205 | rswa_window = model_config.hf_config.rswa_window |
| 206 | text_config.rswa_window = rswa_window |
| 207 | |
| 208 | |
| 209 | class Gemma4Config(VerifyAndUpdateConfig): |
| 210 | @staticmethod |
| 211 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 212 | """Configure attention for heterogeneous head dimensions. |
| 213 | |
| 214 | Gemma4 uses different head dimensions for sliding window vs full attention |
| 215 | layers. The default FA3 on Hopper cannot handle head_dim > 256. |
| 216 | |
| 217 | On SM90 with FP8 KV cache, use FA3 for supported layers and let the generic |
| 218 | FlashAttention selector upgrade larger head dimensions to FA4. |
| 219 | The multimodal-prefix composite routes image masks to Triton and causal |
| 220 | requests to this per-layer FA3/FA4 selection. For other configurations, |
| 221 | force FA4 for all layers to avoid the mixed |
| 222 | FA3+FA4 penalty. |
| 223 | When FA4 is not available, fall back to Triton. |
| 224 | """ |
| 225 | model_config = vllm_config.model_config |
| 226 | arch_config = model_config.model_arch_config |
| 227 | layer_types = getattr(model_config.hf_text_config, "layer_types", None) or [] |
| 228 | head_dims = { |
| 229 | layer_types[i]: arch_config[i].head_size |
| 230 | for i in range(min(arch_config.total_num_hidden_layers, len(layer_types))) |
| 231 | } |
| 232 | |
| 233 | if len(set(head_dims.values())) <= 1: |
| 234 | return |
| 235 | |
| 236 | from vllm.platforms import current_platform |
| 237 | from vllm.v1.attention.backends.fa_utils import is_fa_version_supported |
| 238 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 239 | |
| 240 | max_head_dim = max(head_dims.values()) |
| 241 | |
| 242 | if is_fa_version_supported(4) and max_head_dim <= 512: |
| 243 | if ( |
| 244 | vllm_config.attention_config.flash_attn_version is None |
| 245 | and vllm_config.attention_config.backend |
| 246 | in (None, AttentionBackendEnum.FLASH_ATTN) |
| 247 | ): |
| 248 | use_per_layer_fa = current_platform.is_device_capability_family( |
| 249 | 90 |
| 250 | ) and vllm_config.cache_config.cache_dtype.startswith("fp8") |
| 251 | if use_per_layer_fa: |
| 252 | logger.info( |
| 253 | "Gemma4 model has heterogeneous head dimensions %s. Using " |
| 254 | "per-layer FA3/FA4 selection for FP8 KV cache on SM90.", |
| 255 | head_dims, |
| 256 | ) |
| 257 | else: |
| 258 | vllm_config.attention_config.flash_attn_version = 4 |
| 259 | logger.info( |
| 260 | "Gemma4 model has heterogeneous head dimensions %s. Using FA4 " |
| 261 | "for all layers to avoid mixed FA3/FA4 penalty.", |
| 262 | head_dims, |
| 263 | ) |
| 264 | elif vllm_config.attention_config.backend is None: |
| 265 | vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN |
| 266 | logger.info( |
| 267 | "Gemma4 model has heterogeneous head dimensions " |
| 268 | "%s. FA4 not available, forcing TRITON_ATTN backend.", |
| 269 | head_dims, |
| 270 | ) |
| 271 | |
| 272 | |
| 273 | class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): |
| 274 | @classmethod |
| 275 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 276 | """Set up the diffusion config and defaults for DiffusionGemma. |
| 277 | |
| 278 | Auto-creates DiffusionConfig from the HF config when the user |
| 279 | didn't pass ``--diffusion-config``. Diffusion sampling params are |
| 280 | read straight from generation_config.json at sampler-build time |
| 281 | (see DiffusionGemma's custom_sampler), not injected here. |
| 282 | """ |
| 283 | # Inherit Gemma4's attention backend selection (FA4 on Hopper, |
| 284 | # TRITON_ATTN fallback for heterogeneous head dims). |
| 285 | Gemma4Config.verify_and_update_config(vllm_config) |
| 286 | |
| 287 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 288 | |
| 289 | attention_config = vllm_config.attention_config |
| 290 | if attention_config.backend == AttentionBackendEnum.FLASHINFER: |
| 291 | raise ValueError( |
| 292 | "FlashInfer does not support DiffusionGemma's mixed " |
| 293 | "causal/bidirectional attention. Use --attention-backend " |
| 294 | "FLASH_ATTN or TRITON_ATTN instead." |
| 295 | ) |
| 296 | if attention_config.backend is None and not attention_config.use_non_causal: |
| 297 | # Only the sm_100 default backend order reads this flag (it puts |
| 298 | # FlashInfer first for causal models). Elsewhere Gemma4Config above |
| 299 | # keeps FlashInfer out: FA4 on every layer, else TRITON_ATTN. |
| 300 | attention_config.use_non_causal = True |
| 301 | logger.info( |
| 302 | "DiffusionGemma mixes causal and bidirectional attention in " |
| 303 | "one batch; setting use_non_causal=True so the default backend " |
| 304 | "order prefers FlashAttention over FlashInfer." |
| 305 | ) |
| 306 | |
| 307 | # Auto-create DiffusionConfig from HF config if not provided. |
| 308 | if vllm_config.diffusion_config is None: |
| 309 | from vllm.config.diffusion import DiffusionConfig |
| 310 | |
| 311 | hf_config = vllm_config.model_config.hf_config |
| 312 | canvas_length = getattr(hf_config, "canvas_length", 256) |
| 313 | vllm_config.diffusion_config = DiffusionConfig( |
| 314 | canvas_length=canvas_length, |
| 315 | ) |
| 316 | |
| 317 | # The diffusion sampler materializes [num_seqs, canvas_length, vocab] |
| 318 | # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200). |
| 319 | # Default to 8 when the user didn't pass --max-num-seqs. |
| 320 | # We can't see the original None here (the engine already filled a generic |
| 321 | # default), so use >= DEFAULT_MAX_NUM_SEQS as a proxy, (the default is much |
| 322 | # larger than any deliberate value for this model) |
| 323 | from vllm.config.scheduler import SchedulerConfig |
| 324 | |
| 325 | sc = vllm_config.scheduler_config |
| 326 | if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS: |
| 327 | sc.max_num_seqs = 8 |
| 328 | |
| 329 | # Remove the model's generation_config.json cap on max_new_tokens |
| 330 | # (256) so DiffusionGemma behaves like every other model: no |
| 331 | # server-wide limit, each request controls its own output length |
| 332 | # via max_tokens. Setting to None causes get_diff_sampling_param |
| 333 | # to skip this key entirely. |
| 334 | model_config = vllm_config.model_config |
| 335 | if "max_new_tokens" not in model_config.override_generation_config: |
| 336 | model_config.override_generation_config["max_new_tokens"] = None |
| 337 | logger.info( |
| 338 | "DiffusionGemma: removing server-wide max_new_tokens cap " |
| 339 | "from generation_config.json (use " |
| 340 | "--override-generation-config to set a custom limit).", |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): |
| 345 | @staticmethod |
| 346 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 347 | quant_config = getattr(model_config.hf_config, "quantization_config", None) |
| 348 | if quant_config is not None and quant_config.get("quant_method") == "fp8": |
| 349 | model_type = getattr(model_config.hf_config, "model_type", None) |
| 350 | if model_type in ("deepseek_v4", "deepseek_v41"): |
| 351 | model_config.hf_config.quantization_config["quant_method"] = ( |
| 352 | "deepseek_v4_fp8" |
| 353 | ) |
| 354 | |
| 355 | hf_text_quant_config = getattr( |
| 356 | model_config.hf_text_config, "quantization_config", None |
| 357 | ) |
| 358 | if ( |
| 359 | hf_text_quant_config is not None |
| 360 | and hf_text_quant_config.get("quant_method") == "fp8" |
| 361 | ): |
| 362 | model_type = getattr(model_config.hf_text_config, "model_type", None) |
| 363 | if model_type in ("deepseek_v4", "deepseek_v41"): |
| 364 | model_config.hf_text_config.quantization_config["quant_method"] = ( |
| 365 | "deepseek_v4_fp8" |
| 366 | ) |
| 367 | |
| 368 | |
| 369 | class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig): |
| 370 | """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface. |
| 371 | |
| 372 | Kimi-K3 ships its routed experts as compressed-tensors |
| 373 | ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which |
| 374 | lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel |
| 375 | selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config`` |
| 376 | (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4 |
| 377 | checkpoint is left untouched. Covers both the main model and the MTP draft. |
| 378 | |
| 379 | ``model_arch_config.quantization_config`` is a separate dict, snapshotted in |
| 380 | ``ModelConfig.__init__`` before this hook runs, and it is what |
| 381 | ``_verify_quantization`` reads when resolving the quant method. Patch it |
| 382 | alongside the hf configs so the rewrite lands before resolution; otherwise |
| 383 | the main model still resolves to compressed-tensors. |
| 384 | """ |
| 385 | |
| 386 | @staticmethod |
| 387 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 388 | for cfg in ( |
| 389 | model_config.hf_config, |
| 390 | model_config.hf_text_config, |
| 391 | model_config.model_arch_config, |
| 392 | ): |
| 393 | quant_config = getattr(cfg, "quantization_config", None) |
| 394 | if ( |
| 395 | isinstance(quant_config, dict) |
| 396 | and quant_config.get("quant_method") == "compressed-tensors" |
| 397 | and quant_config.get("format") == "mxfp4-pack-quantized" |
| 398 | ): |
| 399 | quant_config["quant_method"] = "mxfp4" |
| 400 | |
| 401 | |
| 402 | class GptOssForCausalLMConfig(VerifyAndUpdateConfig): |
| 403 | @staticmethod |
| 404 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 405 | quant_config = getattr(model_config.hf_config, "quantization_config", None) |
| 406 | if quant_config is not None and quant_config.get("quant_method") == "mxfp4": |
| 407 | model_config.hf_config.quantization_config["quant_method"] = "gpt_oss_mxfp4" |
| 408 | |
| 409 | hf_text_quant_config = getattr( |
| 410 | model_config.hf_text_config, "quantization_config", None |
| 411 | ) |
| 412 | if ( |
| 413 | hf_text_quant_config is not None |
| 414 | and hf_text_quant_config.get("quant_method") == "mxfp4" |
| 415 | ): |
| 416 | model_config.hf_text_config.quantization_config["quant_method"] = ( |
| 417 | "gpt_oss_mxfp4" |
| 418 | ) |
| 419 | |
| 420 | @staticmethod |
| 421 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 422 | structured_outputs_config = vllm_config.structured_outputs_config |
| 423 | if structured_outputs_config.reasoning_parser == "": |
| 424 | structured_outputs_config.reasoning_parser = "openai_gptoss" |
| 425 | |
| 426 | # Increase the max capture size from 512 to 1024 for performance. |
| 427 | # NOTE(woosuk): This will increase the number of CUDA graphs |
| 428 | # from 67 to 83. |
| 429 | compilation_config = vllm_config.compilation_config |
| 430 | # Only override when the user has not set either of |
| 431 | # cudagraph_capture_sizes or max_cudagraph_capture_size. |
| 432 | if ( |
| 433 | compilation_config.cudagraph_capture_sizes is None |
| 434 | and compilation_config.max_cudagraph_capture_size is None |
| 435 | ): |
| 436 | compilation_config.max_cudagraph_capture_size = 1024 |
| 437 | logger.info( |
| 438 | "Overriding max cuda graph capture size to %d for performance.", 1024 |
| 439 | ) |
| 440 | |
| 441 | |
| 442 | class GteNewModelConfig(VerifyAndUpdateConfig): |
| 443 | @staticmethod |
| 444 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 445 | config = model_config.hf_config |
| 446 | |
| 447 | assert config.__class__.__name__ == "NewConfig" |
| 448 | assert config.hidden_act == "gelu" |
| 449 | |
| 450 | config.hidden_act = "geglu" |
| 451 | |
| 452 | head_dim = config.hidden_size // config.num_attention_heads |
| 453 | rotary_dim = getattr(config, "rotary_emb_dim", head_dim) |
| 454 | config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim |
| 455 | config.rotary_kwargs = { |
| 456 | "head_size": head_dim, |
| 457 | "max_position": config.max_position_embeddings, |
| 458 | "rope_parameters": config.rope_parameters, |
| 459 | } |
| 460 | |
| 461 | |
| 462 | class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig): |
| 463 | @classmethod |
| 464 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 465 | """Perform early validation and setup for hybrid attention/mamba models. |
| 466 | |
| 467 | Block size alignment with mamba page sizes is handled later by |
| 468 | Platform.update_block_size_for_backend(), which runs after model |
| 469 | layers are constructed and the attention backend is known. |
| 470 | |
| 471 | Args: |
| 472 | vllm_config: vLLM Config |
| 473 | |
| 474 | """ |
| 475 | # Enable FULL_AND_PIECEWISE by default |
| 476 | MambaModelConfig.verify_and_update_config(vllm_config) |
| 477 | |
| 478 | |
| 479 | class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig): |
| 480 | @staticmethod |
| 481 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 482 | pooler_config = model_config.pooler_config |
| 483 | assert pooler_config is not None |
| 484 | if pooler_config.use_activation is None: |
| 485 | pooler_config.use_activation = False |
| 486 | |
| 487 | |
| 488 | class JinaEmbeddingsV5ModelConfig(VerifyAndUpdateConfig): |
| 489 | """Config handler for Jina Embeddings V5 embedding models.""" |
| 490 | |
| 491 | @staticmethod |
| 492 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 493 | """Enable the bidirectional encoder backbone for -nano checkpoints. |
| 494 | |
| 495 | The V5 family ships more than one backbone under a single |
| 496 | `architectures` entry: the `-small` variants are Qwen3 decoders, while |
| 497 | `-nano` is a bidirectional EuroBERT encoder. Upstream ships a separate |
| 498 | `configuration_*.py` per repository, so the config carries no backbone |
| 499 | field and the encoder variants are only identifiable by |
| 500 | `is_decoder=False`. For encoder checkpoints, set `is_causal=False` so the |
| 501 | Llama backbone uses EncoderOnlyAttention; `JinaEmbeddingsV5Model` then |
| 502 | dispatches to its encoder implementation. |
| 503 | """ |
| 504 | if getattr(model_config.hf_config, "is_decoder", True): |
| 505 | return |
| 506 | |
| 507 | model_config.hf_config.is_causal = False |
| 508 | |
| 509 | |
| 510 | class JinaForRankingConfig(VerifyAndUpdateConfig): |
| 511 | @staticmethod |
| 512 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 513 | model_config.hf_config.embedding_size = 512 |
| 514 | |
| 515 | |
| 516 | class JinaRobertaModelConfig(VerifyAndUpdateConfig): |
| 517 | @staticmethod |
| 518 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 519 | config = model_config.hf_config |
| 520 | |
| 521 | if config.position_embedding_type == "rotary": |
| 522 | assert config.__class__.__name__ == "XLMRobertaFlashConfig" |
| 523 | |
| 524 | head_dim = config.hidden_size // config.num_attention_heads |
| 525 | max_position = config.max_position_embeddings |
| 526 | # Jina-embeddings-v3 has max_position_embeddings=8194, which will cause |
| 527 | # out-of-bound index issue at RoPE for long prompts with torch.compile, |
| 528 | # because it can't be divided by triton num_warps(default=4 or 8). |
| 529 | # To deal with this, we increase max_position to multiple of n_warps, |
| 530 | # so that triton kernel won't hit out-of-bound index in RoPE cache. |
| 531 | if not model_config.enforce_eager: |
| 532 | max_position = round_up(max_position, 8) |
| 533 | |
| 534 | rotary_dim = getattr(config, "rotary_emb_dim", head_dim) |
| 535 | config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim |
| 536 | |
| 537 | config.rotary_kwargs = { |
| 538 | "head_size": head_dim, |
| 539 | "max_position": max_position, |
| 540 | "rope_parameters": config.rope_parameters, |
| 541 | } |
| 542 | |
| 543 | |
| 544 | class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig): |
| 545 | @staticmethod |
| 546 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 547 | config = model_config.hf_config |
| 548 | config.num_labels = 1 |
| 549 | config.get_text_config().num_labels = 1 |
| 550 | pooler_config = model_config.pooler_config |
| 551 | assert pooler_config is not None |
| 552 | if pooler_config.logit_mean is None: |
| 553 | pooler_config.logit_mean = 2.65 |
| 554 | |
| 555 | |
| 556 | class LlamaBidirectionalConfig(VerifyAndUpdateConfig): |
| 557 | @staticmethod |
| 558 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 559 | from vllm.config.pooler import SequencePoolingType |
| 560 | |
| 561 | hf_config = model_config.hf_config |
| 562 | hf_config.is_causal = False |
| 563 | |
| 564 | pooling_type_map: dict[str, SequencePoolingType] = { |
| 565 | "avg": "MEAN", |
| 566 | "cls": "CLS", |
| 567 | "last": "LAST", |
| 568 | } |
| 569 | |
| 570 | pooling = getattr(hf_config, "pooling", None) |
| 571 | if pooling is None: |
| 572 | raise ValueError( |
| 573 | "This bidirectional model requires a 'pooling' field in its HF " |
| 574 | "config (one of 'avg', 'cls', 'last'), but none was found. " |
| 575 | "Unlike the VL variants, no default is assumed here because " |
| 576 | "silently picking a pooling for an embedding model can produce " |
| 577 | "wrong embeddings." |
| 578 | ) |
| 579 | |
| 580 | pooling_type = pooling_type_map.get(pooling) |
| 581 | if pooling_type is None: |
| 582 | raise ValueError(f"pool_type {pooling!r} not supported") |
| 583 | |
| 584 | assert model_config.pooler_config is not None |
| 585 | model_config.pooler_config.seq_pooling_type = pooling_type |
| 586 | |
| 587 | |
| 588 | class LlamaNemotronVLConfig(VerifyAndUpdateConfig): |
| 589 | """Config handler for LlamaNemotronVL embedding models.""" |
| 590 | |
| 591 | @staticmethod |
| 592 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 593 | from vllm.config.pooler import SequencePoolingType |
| 594 | |
| 595 | hf_config = model_config.hf_config |
| 596 | |
| 597 | # Set bidirectional attention on the language model config |
| 598 | hf_config.is_causal = False |
| 599 | if hasattr(hf_config, "llm_config"): |
| 600 | hf_config.llm_config.is_causal = False |
| 601 | |
| 602 | if hasattr(hf_config, "vision_config"): |
| 603 | hf_config.patch_size = hf_config.vision_config.patch_size |
| 604 | |
| 605 | # Set up pooling type |
| 606 | pooling_type_map: dict[str, SequencePoolingType] = { |
| 607 | "avg": "MEAN", |
| 608 | "cls": "CLS", |
| 609 | "last": "LAST", |
| 610 | } |
| 611 | |
| 612 | # Get pooling type from config (check both top-level and llm_config) |
| 613 | pooling = getattr(hf_config, "pooling", None) |
| 614 | if pooling is None and hasattr(hf_config, "llm_config"): |
| 615 | pooling = getattr(hf_config.llm_config, "pooling", "avg") |
| 616 | |
| 617 | assert isinstance(pooling, str) |
| 618 | |
| 619 | pooling_type = pooling_type_map.get(pooling) |
| 620 | if pooling_type is None: |
| 621 | raise ValueError(f"pool_type {pooling!r} not supported") |
| 622 | |
| 623 | assert model_config.pooler_config is not None |
| 624 | model_config.pooler_config.seq_pooling_type = pooling_type |
| 625 | |
| 626 | |
| 627 | class MambaModelConfig(VerifyAndUpdateConfig): |
| 628 | @classmethod |
| 629 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 630 | """Enable FULL_AND_PIECEWISE cuda graph mode by default (required |
| 631 | to get good performance for mamba layers in V1). |
| 632 | |
| 633 | Args: |
| 634 | vllm_config: vLLM Config |
| 635 | |
| 636 | """ |
| 637 | model_config = vllm_config.model_config |
| 638 | cache_config = vllm_config.cache_config |
| 639 | |
| 640 | if cache_config.enable_prefix_caching: |
| 641 | if cache_config.mamba_cache_mode == "none": |
| 642 | cache_config.mamba_cache_mode = "align" |
| 643 | logger.info( |
| 644 | "Mamba cache mode is set to '%s' for %s by default " |
| 645 | "when prefix caching is enabled", |
| 646 | cache_config.mamba_cache_mode, |
| 647 | model_config.architecture, |
| 648 | ) |
| 649 | assert vllm_config.scheduler_config.enable_chunked_prefill, ( |
| 650 | "Chunked prefill is required for mamba cache mode 'align'." |
| 651 | ) |
| 652 | # By default, mamba block size will be set to max_model_len (see |
| 653 | # below). When enabling prefix caching, we align mamba block size |
| 654 | # to the block size as the basic granularity for prefix caching. |
| 655 | if cache_config.mamba_block_size is None: |
| 656 | cache_config.mamba_block_size = cache_config.block_size |
| 657 | else: |
| 658 | if cache_config.mamba_cache_mode != "none": |
| 659 | cache_config.mamba_cache_mode = "none" |
| 660 | logger.warning( |
| 661 | "Mamba cache mode is set to 'none' when prefix caching is disabled" |
| 662 | ) |
| 663 | if cache_config.mamba_block_size is None: |
| 664 | cache_config.mamba_block_size = model_config.max_model_len |
| 665 | |
| 666 | |
| 667 | class NemotronHForCausalLMConfig(VerifyAndUpdateConfig): |
| 668 | DEFAULT_MAMBA_SSM_CACHE_DTYPE: "MambaDType" = "float32" |
| 669 | """Only `float32` is known to have no accuracy issues by default.""" |
| 670 | |
| 671 | @classmethod |
| 672 | def update_mamba_ssm_cache_dtype( |
| 673 | cls, *, cache_config: "CacheConfig", hf_config: "PreTrainedConfig" |
| 674 | ) -> None: |
| 675 | """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto' |
| 676 | (or not explicitly set), to the value specified in the HF config, or to |
| 677 | `float32` if not specified. |
| 678 | """ |
| 679 | if cache_config.mamba_ssm_cache_dtype == "auto": |
| 680 | mamba_ssm_cache_dtype: MambaDType = getattr( |
| 681 | hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE |
| 682 | ) |
| 683 | logger.info( |
| 684 | "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model", |
| 685 | mamba_ssm_cache_dtype, |
| 686 | ) |
| 687 | cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype |
| 688 | |
| 689 | @classmethod |
| 690 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 691 | cls.update_mamba_ssm_cache_dtype( |
| 692 | cache_config=vllm_config.cache_config, |
| 693 | hf_config=vllm_config.model_config.hf_config, |
| 694 | ) |
| 695 | |
| 696 | |
| 697 | class NemotronHNanoVLV2Config(VerifyAndUpdateConfig): |
| 698 | @classmethod |
| 699 | def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: |
| 700 | NemotronHForCausalLMConfig.update_mamba_ssm_cache_dtype( |
| 701 | cache_config=vllm_config.cache_config, |
| 702 | hf_config=vllm_config.model_config.hf_config.text_config, |
| 703 | ) |
| 704 | |
| 705 | @staticmethod |
| 706 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 707 | mm_config = model_config.multimodal_config |
| 708 | if mm_config is not None: |
| 709 | video_kwargs = mm_config.media_io_kwargs.setdefault("video", {}) |
| 710 | video_kwargs.setdefault("video_backend", "nemotron_vl") |
| 711 | |
| 712 | |
| 713 | class NomicBertModelConfig(VerifyAndUpdateConfig): |
| 714 | @staticmethod |
| 715 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 716 | config = model_config.hf_config |
| 717 | |
| 718 | assert config.__class__.__name__ == "NomicBertConfig" |
| 719 | assert config.activation_function in ["swiglu", "gelu"] |
| 720 | config.position_embedding_type = getattr( |
| 721 | config, "position_embedding_type", "rope" |
| 722 | ) |
| 723 | |
| 724 | if config.activation_function == "swiglu": |
| 725 | config.hidden_act = "silu" |
| 726 | else: |
| 727 | config.hidden_act = config.activation_function |
| 728 | |
| 729 | assert config.mlp_fc1_bias == config.mlp_fc2_bias == config.qkv_proj_bias |
| 730 | config.bias = config.qkv_proj_bias |
| 731 | |
| 732 | assert config.rotary_emb_scale_base is None |
| 733 | assert not config.rotary_emb_interleaved |
| 734 | |
| 735 | config.layer_norm_eps = config.layer_norm_epsilon |
| 736 | config.intermediate_size = config.n_inner |
| 737 | config.hidden_size = config.n_embd |
| 738 | config.num_hidden_layers = config.n_layer |
| 739 | model_config.model_arch_config.hidden_size = config.hidden_size |
| 740 | model_config.model_arch_config.total_num_hidden_layers = ( |
| 741 | config.num_hidden_layers |
| 742 | ) |
| 743 | |
| 744 | head_dim = config.hidden_size // config.num_attention_heads |
| 745 | max_position_embeddings = getattr(config, "max_position_embeddings", 2048) |
| 746 | max_trained_positions = getattr( |
| 747 | config, "max_trained_positions", max_position_embeddings |
| 748 | ) |
| 749 | |
| 750 | rope_parameters = { |
| 751 | "max_trained_positions": max_trained_positions, |
| 752 | **(config.rope_parameters or {}), |
| 753 | } |
| 754 | |
| 755 | config.rotary_kwargs = { |
| 756 | "head_size": head_dim, |
| 757 | "max_position": model_config.max_model_len, |
| 758 | "rope_parameters": rope_parameters, |
| 759 | } |
| 760 | |
| 761 | |
| 762 | class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig): |
| 763 | @staticmethod |
| 764 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 765 | pooler_config = model_config.pooler_config |
| 766 | assert pooler_config is not None |
| 767 | |
| 768 | if pooler_config.step_tag_id is None: |
| 769 | pooler_config.step_tag_id = 151651 |
| 770 | |
| 771 | |
| 772 | class Qwen2ForRewardModelConfig(VerifyAndUpdateConfig): |
| 773 | @staticmethod |
| 774 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 775 | pooler_config = model_config.pooler_config |
| 776 | assert pooler_config is not None |
| 777 | |
| 778 | if pooler_config.use_activation is None: |
| 779 | pooler_config.use_activation = False |
| 780 | |
| 781 | |
| 782 | class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig): |
| 783 | @staticmethod |
| 784 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 785 | config = model_config.hf_config |
| 786 | |
| 787 | is_original_qwen3_reranker = getattr( |
| 788 | config, "is_original_qwen3_reranker", False |
| 789 | ) |
| 790 | |
| 791 | if not is_original_qwen3_reranker: |
| 792 | return |
| 793 | |
| 794 | tokens = getattr(config, "classifier_from_token", None) |
| 795 | assert tokens is not None and len(tokens) == 2, ( |
| 796 | "Try loading the original Qwen3 Reranker?, see: " |
| 797 | "https://github.com/vllm-project/vllm/tree/main/examples/pooling/score/qwen3_reranker_offline.py" |
| 798 | ) |
| 799 | text_config = config.get_text_config() |
| 800 | text_config.method = "from_2_way_softmax" |
| 801 | text_config.classifier_from_token = tokens |
| 802 | |
| 803 | |
| 804 | class Qwen3VLForSequenceClassificationConfig(Qwen3ForSequenceClassificationConfig): |
| 805 | pass |
| 806 | |
| 807 | |
| 808 | class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig): |
| 809 | @staticmethod |
| 810 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 811 | """Update mamba_ssm_cache_dtype for Qwen3.5 models when set to 'auto' |
| 812 | (or not explicitly set), to the value specified in the HF config's |
| 813 | mamba_ssm_dtype field. Warn if the user explicitly overrides it to a |
| 814 | different value. |
| 815 | """ |
| 816 | cache_config = vllm_config.cache_config |
| 817 | hf_text_config = vllm_config.model_config.hf_text_config |
| 818 | mamba_ssm_dtype = getattr(hf_text_config, "mamba_ssm_dtype", None) |
| 819 | if cache_config.mamba_ssm_cache_dtype == "auto": |
| 820 | if mamba_ssm_dtype is not None: |
| 821 | cache_config.mamba_ssm_cache_dtype = mamba_ssm_dtype |
| 822 | elif ( |
| 823 | mamba_ssm_dtype is not None |
| 824 | and cache_config.mamba_ssm_cache_dtype != mamba_ssm_dtype |
| 825 | ): |
| 826 | logger.warning( |
| 827 | "Qwen3.5 model specifies mamba_ssm_dtype='%s' in its config, " |
| 828 | "but --mamba-ssm-cache-dtype='%s' was passed. " |
| 829 | "Using the user-specified value.", |
| 830 | mamba_ssm_dtype, |
| 831 | cache_config.mamba_ssm_cache_dtype, |
| 832 | ) |
| 833 | |
| 834 | |
| 835 | class Qwen3_5ForCausalLMConfig(Qwen3_5ForConditionalGenerationConfig): |
| 836 | @staticmethod |
| 837 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 838 | Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) |
| 839 | |
| 840 | # Text-only Qwen3.5 models use one-dimensional positions. Remove the |
| 841 | # M-RoPE fields inherited from the multimodal configuration. |
| 842 | hf_text_config = vllm_config.model_config.hf_text_config |
| 843 | rope_parameters = getattr(hf_text_config, "rope_parameters", None) |
| 844 | if rope_parameters is not None: |
| 845 | rope_parameters.pop("mrope_section", None) |
| 846 | rope_parameters.pop("mrope_interleaved", None) |
| 847 | |
| 848 | |
| 849 | def _strip_qwen4_exp_mrope(model_config: "ModelConfig") -> None: |
| 850 | configs = { |
| 851 | id(config): config |
| 852 | for config in ( |
| 853 | getattr(model_config, "hf_config", None), |
| 854 | model_config.hf_text_config, |
| 855 | ) |
| 856 | if config is not None |
| 857 | } |
| 858 | for config in configs.values(): |
| 859 | rope_parameters = getattr(config, "rope_parameters", None) |
| 860 | if rope_parameters is not None: |
| 861 | rope_parameters.pop("mrope_section", None) |
| 862 | rope_parameters.pop("mrope_interleaved", None) |
| 863 | |
| 864 | |
| 865 | class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConfig): |
| 866 | """Apply the Qwen3.5 hybrid-cache contract to Qwen4Exp.""" |
| 867 | |
| 868 | @staticmethod |
| 869 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 870 | Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) |
| 871 | text_config = vllm_config.model_config.hf_text_config |
| 872 | if text_config.hc_count <= 1: |
| 873 | raise ValueError("Qwen4Exp requires hc_count > 1") |
| 874 | parallel_config = vllm_config.parallel_config |
| 875 | uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( |
| 876 | getattr(text_config, "indexer_n_heads", None) is not None |
| 877 | ) |
| 878 | if uses_ple_or_qsa and ( |
| 879 | parallel_config.enable_dbo or parallel_config.ubatch_size > 1 |
| 880 | ): |
| 881 | raise NotImplementedError( |
| 882 | "Qwen4Exp PLE/QSA does not support dual-batch overlap or microbatching" |
| 883 | ) |
| 884 | # Checked again in Qwen4ExpModelState; rejecting it here keeps the |
| 885 | # engine from loading weights first. |
| 886 | if text_config.ple_layer_ids and parallel_config.pipeline_parallel_size > 1: |
| 887 | raise NotImplementedError( |
| 888 | "Qwen4Exp N-gram PLE embedding requires pipeline_parallel_size=1 " |
| 889 | "because non-first pipeline ranks do not receive the raw input_ids " |
| 890 | "it needs. Please run with PP=1." |
| 891 | ) |
| 892 | multimodal_config = vllm_config.model_config.multimodal_config |
| 893 | if multimodal_config is not None and multimodal_config.language_model_only: |
| 894 | _strip_qwen4_exp_mrope(vllm_config.model_config) |
| 895 | spec_config = vllm_config.speculative_config |
| 896 | if spec_config is not None and spec_config.method not in { |
| 897 | "mtp", |
| 898 | "ngram", |
| 899 | "ngram_gpu", |
| 900 | }: |
| 901 | raise NotImplementedError( |
| 902 | "Qwen4Exp speculative decoding supports only its native MTP " |
| 903 | "checkpoint and linear n-gram proposers" |
| 904 | ) |
| 905 | |
| 906 | |
| 907 | class Qwen4ExpForCausalLMConfig(Qwen4ExpForConditionalGenerationConfig): |
| 908 | @staticmethod |
| 909 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 910 | Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) |
| 911 | |
| 912 | _strip_qwen4_exp_mrope(vllm_config.model_config) |
| 913 | |
| 914 | |
| 915 | class Qwen4ExpMTPConfig(Qwen4ExpForConditionalGenerationConfig): |
| 916 | """Preserve MRoPE for a VL target and use 1D RoPE for a text target.""" |
| 917 | |
| 918 | @staticmethod |
| 919 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 920 | Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) |
| 921 | if hasattr(vllm_config.model_config.hf_config, "vision_config"): |
| 922 | return |
| 923 | _strip_qwen4_exp_mrope(vllm_config.model_config) |
| 924 | |
| 925 | |
| 926 | class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig): |
| 927 | """Apply the attention contract declared by a ColQwen3.5 checkpoint.""" |
| 928 | |
| 929 | @staticmethod |
| 930 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 931 | configs = { |
| 932 | id(config): config |
| 933 | for config in ( |
| 934 | model_config.hf_config, |
| 935 | model_config.hf_text_config, |
| 936 | ) |
| 937 | } |
| 938 | declarations = [ |
| 939 | contract |
| 940 | for config in configs.values() |
| 941 | if (contract := getattr(config, "retrieval_attention_contract", None)) |
| 942 | is not None |
| 943 | ] |
| 944 | supported = {"causal", "bidirectional"} |
| 945 | if not declarations: |
| 946 | raise ValueError( |
| 947 | "ColQwen3.5 checkpoints must declare " |
| 948 | "retrieval_attention_contract as 'causal' or 'bidirectional'" |
| 949 | ) |
| 950 | if ( |
| 951 | any(not isinstance(contract, str) for contract in declarations) |
| 952 | or any(contract not in supported for contract in declarations) |
| 953 | or len(set(declarations)) != 1 |
| 954 | ): |
| 955 | raise ValueError( |
| 956 | "unsupported or conflicting ColQwen3.5 " |
| 957 | f"retrieval_attention_contract declarations: {declarations!r}" |
| 958 | ) |
| 959 | is_causal = declarations[0] == "causal" |
| 960 | for config in configs.values(): |
| 961 | config.is_causal = is_causal |
| 962 | |
| 963 | |
| 964 | class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): |
| 965 | @staticmethod |
| 966 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 967 | config = model_config.hf_config |
| 968 | |
| 969 | assert config.__class__.__name__ == "GteConfig" |
| 970 | assert config.hidden_act == "gelu" |
| 971 | |
| 972 | config.hidden_act = "geglu" |
| 973 | |
| 974 | head_dim = config.hidden_size // config.num_attention_heads |
| 975 | rotary_dim = getattr(config, "rotary_emb_dim", head_dim) |
| 976 | config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim |
| 977 | config.rotary_kwargs = { |
| 978 | "head_size": head_dim, |
| 979 | "max_position": config.max_position_embeddings, |
| 980 | "rope_parameters": config.rope_parameters, |
| 981 | } |
| 982 | |
| 983 | |
| 984 | class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig): |
| 985 | @staticmethod |
| 986 | def verify_and_update_model_config(model_config: "ModelConfig") -> None: |
| 987 | model_config.hf_config.is_causal = False |
| 988 | model_config.hf_config.embedding_size = model_config.hf_config.num_labels |
| 989 | |
| 990 | |
| 991 | class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig): |
| 992 | @staticmethod |
| 993 | def verify_and_update_config(vllm_config: "VllmConfig") -> None: |
| 994 | # LongCat-Flash-Lite's zero-expert MoE trips a data-dependent assert |
| 995 | # under torch.compile, and its n-gram inputs_embeds are only wired for |
| 996 | # FULL cudagraph capture (PIECEWISE prefill drops them). Default to |
| 997 | # no-compile + FULL cudagraph (prefill runs eager) unless the user |
| 998 | # configured compilation explicitly. |
| 999 | from vllm.config.compilation import CompilationMode, CUDAGraphMode |
| 1000 | |
| 1001 | compilation_config = vllm_config.compilation_config |
| 1002 | if compilation_config.mode is None: |
| 1003 | compilation_config.mode = CompilationMode.NONE |
| 1004 | if compilation_config.cudagraph_mode is None: |
| 1005 | compilation_config.cudagraph_mode = CUDAGraphMode.FULL |
| 1006 | |
| 1007 | |
| 1008 | MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { |
| 1009 | "ColBERTJinaRobertaModel": JinaRobertaModelConfig, |
| 1010 | "ColQwen3_5": ColQwen3_5Config, |
| 1011 | "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, |
| 1012 | "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig, |
| 1013 | "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig, |
| 1014 | "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, |
| 1015 | "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 |
| 1016 | "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 |
| 1017 | "FalconMambaForCausalLM": MambaModelConfig, |
| 1018 | "Gemma3TextModel": Gemma3TextModelConfig, |
| 1019 | "Gemma4ForCausalLM": Gemma4Config, |
| 1020 | "Gemma4ForConditionalGeneration": Gemma4Config, |
| 1021 | "Gemma4UnifiedForConditionalGeneration": Gemma4Config, |
| 1022 | "GlmMoeDsaForCausalLM": GlmMoeDsaForCausalLM, |
| 1023 | "GptOssForCausalLM": GptOssForCausalLMConfig, |
| 1024 | "LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig, |
| 1025 | "GteModel": SnowflakeGteNewModelConfig, |
| 1026 | "GteNewForSequenceClassification": GteNewModelConfig, |
| 1027 | "GteNewModel": GteNewModelConfig, |
| 1028 | "JambaForSequenceClassification": JambaForSequenceClassificationConfig, |
| 1029 | "JinaEmbeddingsV5Model": JinaEmbeddingsV5ModelConfig, |
| 1030 | "JinaForRanking": JinaForRankingConfig, |
| 1031 | "JinaVLForRanking": JinaVLForSequenceClassificationConfig, |
| 1032 | "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig, |
| 1033 | "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig, |
| 1034 | "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, |
| 1035 | "LlamaBidirectionalModel": LlamaBidirectionalConfig, |
| 1036 | "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig, |
| 1037 | "LlamaNemotronVLModel": LlamaNemotronVLConfig, |
| 1038 | "Mamba2ForCausalLM": MambaModelConfig, |
| 1039 | "MambaForCausalLM": MambaModelConfig, |
| 1040 | "NemotronHForCausalLM": NemotronHForCausalLMConfig, |
| 1041 | "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig, |
| 1042 | "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config, |
| 1043 | "NomicBertModel": NomicBertModelConfig, |
| 1044 | "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig, |
| 1045 | "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, |
| 1046 | "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig, |
| 1047 | "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, |
| 1048 | "Qwen3_5ForCausalLM": Qwen3_5ForCausalLMConfig, |
| 1049 | "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, |
| 1050 | "Qwen3_5MoeForCausalLM": Qwen3_5ForCausalLMConfig, |
| 1051 | "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, |
| 1052 | "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, |
| 1053 | "Qwen4ExpForConditionalGeneration": (Qwen4ExpForConditionalGenerationConfig), |
| 1054 | "Qwen4ExpMTP": Qwen4ExpMTPConfig, |
| 1055 | "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig, |
| 1056 | "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, |
| 1057 | "XLMRobertaModel": JinaRobertaModelConfig, |
| 1058 | } |