wehub-resource-sync
869 lines
37 KiB
Python
869 lines
37 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.math_utils import round_up
|
|
|
|
if TYPE_CHECKING:
|
|
from transformers import PretrainedConfig
|
|
|
|
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
|
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class VerifyAndUpdateConfig:
|
|
@staticmethod
|
|
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
|
return
|
|
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
return
|
|
|
|
|
|
class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
|
|
@classmethod
|
|
def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
|
|
hf_config = vllm_config.model_config.hf_config
|
|
|
|
# Mirror the check in vllm/model_executor/models/deepseek_v2.py
|
|
is_v32 = hasattr(hf_config, "index_topk")
|
|
assert is_v32
|
|
|
|
cache_config = vllm_config.cache_config
|
|
if cache_config.cache_dtype == "bfloat16":
|
|
cache_config.cache_dtype = "auto"
|
|
logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
|
|
|
|
|
|
class Ernie4_5_VLMoeForConditionalGenerationConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
|
# Ernie4.5-VL conditionally executes text/vision MoE branches, so
|
|
# fast_moe_cold_start can silently produce incorrect execution order.
|
|
vllm_config.compilation_config.fast_moe_cold_start = False
|
|
|
|
|
|
class Gemma3TextModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
hf_config = model_config.hf_config
|
|
hf_config.is_causal = not hf_config.use_bidirectional_attention
|
|
|
|
|
|
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
|
|
|
|
|
|
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
|
|
(head_dim) vs full attention (global_head_dim) 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.
|
|
"""
|
|
hf_text_config = vllm_config.model_config.hf_text_config
|
|
head_dim = getattr(hf_text_config, "head_dim", None)
|
|
global_head_dim = getattr(hf_text_config, "global_head_dim", None)
|
|
|
|
if head_dim is None or global_head_dim is None or head_dim == global_head_dim:
|
|
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_dim, global_head_dim)
|
|
|
|
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 "
|
|
"(head_dim=%d, global_head_dim=%d). Using FA4 for "
|
|
"all layers to avoid mixed FA3/FA4 penalty.",
|
|
head_dim,
|
|
global_head_dim,
|
|
)
|
|
elif vllm_config.attention_config.backend is None:
|
|
vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
|
|
logger.info(
|
|
"Gemma4 model has heterogeneous head dimensions "
|
|
"(head_dim=%d, global_head_dim=%d). FA4 not available, "
|
|
"forcing TRITON_ATTN backend.",
|
|
head_dim,
|
|
global_head_dim,
|
|
)
|
|
|
|
|
|
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 DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
quant_config = getattr(model_config.hf_config, "quantization_config", None)
|
|
if quant_config is not None and quant_config.get("quant_method") == "fp8":
|
|
model_type = getattr(model_config.hf_config, "model_type", None)
|
|
if model_type == "deepseek_v4":
|
|
model_config.hf_config.quantization_config["quant_method"] = (
|
|
"deepseek_v4_fp8"
|
|
)
|
|
|
|
hf_text_quant_config = getattr(
|
|
model_config.hf_text_config, "quantization_config", None
|
|
)
|
|
if (
|
|
hf_text_quant_config is not None
|
|
and hf_text_quant_config.get("quant_method") == "fp8"
|
|
):
|
|
model_type = getattr(model_config.hf_text_config, "model_type", None)
|
|
if model_type == "deepseek_v4":
|
|
model_config.hf_text_config.quantization_config["quant_method"] = (
|
|
"deepseek_v4_fp8"
|
|
)
|
|
|
|
|
|
class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
quant_config = getattr(model_config.hf_config, "quantization_config", None)
|
|
if quant_config is not None and quant_config.get("quant_method") == "mxfp4":
|
|
model_config.hf_config.quantization_config["quant_method"] = "gpt_oss_mxfp4"
|
|
|
|
hf_text_quant_config = getattr(
|
|
model_config.hf_text_config, "quantization_config", None
|
|
)
|
|
if (
|
|
hf_text_quant_config is not None
|
|
and hf_text_quant_config.get("quant_method") == "mxfp4"
|
|
):
|
|
model_config.hf_text_config.quantization_config["quant_method"] = (
|
|
"gpt_oss_mxfp4"
|
|
)
|
|
|
|
@staticmethod
|
|
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
|
structured_outputs_config = vllm_config.structured_outputs_config
|
|
if structured_outputs_config.reasoning_parser == "":
|
|
structured_outputs_config.reasoning_parser = "openai_gptoss"
|
|
|
|
# Increase the max capture size from 512 to 1024 for performance.
|
|
# NOTE(woosuk): This will increase the number of CUDA graphs
|
|
# from 67 to 83.
|
|
compilation_config = vllm_config.compilation_config
|
|
# Only override when the user has not set either of
|
|
# cudagraph_capture_sizes or max_cudagraph_capture_size.
|
|
if (
|
|
compilation_config.cudagraph_capture_sizes is None
|
|
and compilation_config.max_cudagraph_capture_size is None
|
|
):
|
|
compilation_config.max_cudagraph_capture_size = 1024
|
|
logger.info(
|
|
"Overriding max cuda graph capture size to %d for performance.", 1024
|
|
)
|
|
|
|
|
|
class GteNewModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
|
|
assert config.__class__.__name__ == "NewConfig"
|
|
assert config.hidden_act == "gelu"
|
|
|
|
config.hidden_act = "geglu"
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
rotary_dim = getattr(config, "rotary_emb_dim", head_dim)
|
|
config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
|
|
config.rotary_kwargs = {
|
|
"head_size": head_dim,
|
|
"max_position": config.max_position_embeddings,
|
|
"rope_parameters": config.rope_parameters,
|
|
}
|
|
|
|
|
|
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
|
|
"""
|
|
cache_config = vllm_config.cache_config
|
|
|
|
# Disable calculate_kv_scales for hybrid models: uninitialized
|
|
# recurrent state corrupts scales during the calibration pass.
|
|
# See issue: https://github.com/vllm-project/vllm/issues/37554
|
|
|
|
if cache_config.calculate_kv_scales:
|
|
logger.warning(
|
|
"Disabling calculate_kv_scales for hybrid model '%s'. "
|
|
"Hybrid models with recurrent layers (GDN, Mamba, SSM) "
|
|
"produce unreliable KV cache scales during the "
|
|
"calibration pass because recurrent state is "
|
|
"uninitialized. Using default scale of 1.0 instead.",
|
|
vllm_config.model_config.model,
|
|
)
|
|
cache_config.calculate_kv_scales = False
|
|
|
|
# Enable FULL_AND_PIECEWISE by default
|
|
MambaModelConfig.verify_and_update_config(vllm_config)
|
|
|
|
|
|
class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
pooler_config = model_config.pooler_config
|
|
if pooler_config.use_activation is None:
|
|
pooler_config.use_activation = False
|
|
|
|
|
|
class JinaForRankingConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
model_config.hf_config.embedding_size = 512
|
|
|
|
|
|
class JinaRobertaModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
|
|
if config.position_embedding_type == "rotary":
|
|
assert config.__class__.__name__ == "XLMRobertaFlashConfig"
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
max_position = config.max_position_embeddings
|
|
# Jina-embeddings-v3 has max_position_embeddings=8194, which will cause
|
|
# out-of-bound index issue at RoPE for long prompts with torch.compile,
|
|
# because it can't be divided by triton num_warps(default=4 or 8).
|
|
# To deal with this, we increase max_position to multiple of n_warps,
|
|
# so that triton kernel won't hit out-of-bound index in RoPE cache.
|
|
if not model_config.enforce_eager:
|
|
max_position = round_up(max_position, 8)
|
|
|
|
rotary_dim = getattr(config, "rotary_emb_dim", head_dim)
|
|
config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
|
|
|
|
config.rotary_kwargs = {
|
|
"head_size": head_dim,
|
|
"max_position": max_position,
|
|
"rope_parameters": config.rope_parameters,
|
|
}
|
|
|
|
|
|
class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
config.num_labels = 1
|
|
pooler_config = model_config.pooler_config
|
|
if pooler_config.logit_mean is None:
|
|
pooler_config.logit_mean = 2.65
|
|
|
|
|
|
class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
from vllm.config.pooler import SequencePoolingType
|
|
|
|
hf_config = model_config.hf_config
|
|
hf_config.is_causal = False
|
|
|
|
pooling_type_map: dict[str, SequencePoolingType] = {
|
|
"avg": "MEAN",
|
|
"cls": "CLS",
|
|
"last": "LAST",
|
|
}
|
|
|
|
pooling_type = pooling_type_map.get(hf_config.pooling)
|
|
if pooling_type is None:
|
|
raise ValueError(f"pool_type {hf_config.pooling!r} not supported")
|
|
|
|
model_config.pooler_config.seq_pooling_type = pooling_type
|
|
|
|
|
|
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")
|
|
|
|
pooling_type = pooling_type_map.get(pooling)
|
|
if pooling_type is None:
|
|
raise ValueError(f"pool_type {pooling!r} not supported")
|
|
|
|
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 = (
|
|
"all" if model_config.supports_mamba_prefix_caching else "align"
|
|
)
|
|
logger.warning(
|
|
"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'."
|
|
)
|
|
logger.info(
|
|
"Warning: Prefix caching in Mamba cache '%s' "
|
|
"mode is currently enabled. "
|
|
"Its support for Mamba layers is experimental. "
|
|
"Please report any issues you may observe.",
|
|
cache_config.mamba_cache_mode,
|
|
)
|
|
# 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 = "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 = 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 NemotronHNanoVLV2Config(VerifyAndUpdateConfig):
|
|
@classmethod
|
|
def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
|
|
NemotronHForCausalLMConfig.update_mamba_ssm_cache_dtype(
|
|
cache_config=vllm_config.cache_config,
|
|
hf_config=vllm_config.model_config.hf_config.text_config,
|
|
)
|
|
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
mm_config = model_config.multimodal_config
|
|
if mm_config is not None:
|
|
video_kwargs = mm_config.media_io_kwargs.setdefault("video", {})
|
|
video_kwargs.setdefault("video_backend", "nemotron_vl")
|
|
|
|
|
|
class NomicBertModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
|
|
assert config.__class__.__name__ == "NomicBertConfig"
|
|
assert config.activation_function in ["swiglu", "gelu"]
|
|
config.position_embedding_type = getattr(
|
|
config, "position_embedding_type", "rope"
|
|
)
|
|
|
|
if config.activation_function == "swiglu":
|
|
config.hidden_act = "silu"
|
|
else:
|
|
config.hidden_act = config.activation_function
|
|
|
|
assert config.mlp_fc1_bias == config.mlp_fc2_bias == config.qkv_proj_bias
|
|
config.bias = config.qkv_proj_bias
|
|
|
|
assert config.rotary_emb_scale_base is None
|
|
assert not config.rotary_emb_interleaved
|
|
|
|
config.layer_norm_eps = config.layer_norm_epsilon
|
|
config.intermediate_size = config.n_inner
|
|
config.hidden_size = config.n_embd
|
|
config.num_hidden_layers = config.n_layer
|
|
model_config.model_arch_config.hidden_size = config.hidden_size
|
|
model_config.model_arch_config.total_num_hidden_layers = (
|
|
config.num_hidden_layers
|
|
)
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
max_position_embeddings = getattr(config, "max_position_embeddings", 2048)
|
|
max_trained_positions = getattr(
|
|
config, "max_trained_positions", max_position_embeddings
|
|
)
|
|
|
|
rope_parameters = {
|
|
"max_trained_positions": max_trained_positions,
|
|
**(config.rope_parameters or {}),
|
|
}
|
|
|
|
config.rotary_kwargs = {
|
|
"head_size": head_dim,
|
|
"max_position": model_config.max_model_len,
|
|
"rope_parameters": rope_parameters,
|
|
}
|
|
|
|
|
|
class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
pooler_config = model_config.pooler_config
|
|
|
|
if pooler_config.step_tag_id is None:
|
|
pooler_config.step_tag_id = 151651
|
|
|
|
|
|
class Qwen2ForRewardModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
pooler_config = model_config.pooler_config
|
|
|
|
if pooler_config.use_activation is None:
|
|
pooler_config.use_activation = False
|
|
|
|
|
|
class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
|
|
is_original_qwen3_reranker = getattr(
|
|
config, "is_original_qwen3_reranker", False
|
|
)
|
|
|
|
if not is_original_qwen3_reranker:
|
|
return
|
|
|
|
tokens = getattr(config, "classifier_from_token", None)
|
|
assert tokens is not None and len(tokens) == 2, (
|
|
"Try loading the original Qwen3 Reranker?, see: "
|
|
"https://github.com/vllm-project/vllm/tree/main/examples/pooling/score/qwen3_reranker_offline.py"
|
|
)
|
|
text_config = config.get_text_config()
|
|
text_config.method = "from_2_way_softmax"
|
|
text_config.classifier_from_token = tokens
|
|
|
|
|
|
class Qwen3VLForSequenceClassificationConfig(Qwen3ForSequenceClassificationConfig):
|
|
pass
|
|
|
|
|
|
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 ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
|
|
"""ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache
|
|
handling and additionally serves BIDIRECTIONAL attention: ColPali-style
|
|
document/query encoding attends over the whole sequence, not causally. Set
|
|
is_causal=False so Qwen3NextAttention builds its full_attention layers with
|
|
AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are
|
|
unaffected). Generation arches keep the parent (causal) and are untouched.
|
|
"""
|
|
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
model_config.hf_config.is_causal = False
|
|
|
|
|
|
class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
config = model_config.hf_config
|
|
|
|
assert config.__class__.__name__ == "GteConfig"
|
|
assert config.hidden_act == "gelu"
|
|
|
|
config.hidden_act = "geglu"
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
rotary_dim = getattr(config, "rotary_emb_dim", head_dim)
|
|
config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
|
|
config.rotary_kwargs = {
|
|
"head_size": head_dim,
|
|
"max_position": config.max_position_embeddings,
|
|
"rope_parameters": config.rope_parameters,
|
|
}
|
|
|
|
|
|
class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
|
|
model_config.hf_config.is_causal = False
|
|
model_config.hf_config.embedding_size = model_config.hf_config.num_labels
|
|
|
|
|
|
class LongcatFlashNgramForCausalLMConfig(VerifyAndUpdateConfig):
|
|
@staticmethod
|
|
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
|
|
# LongCat-Flash-Lite's zero-expert MoE trips a data-dependent assert
|
|
# under torch.compile, and its n-gram inputs_embeds are only wired for
|
|
# FULL cudagraph capture (PIECEWISE prefill drops them). Default to
|
|
# no-compile + FULL cudagraph (prefill runs eager) unless the user
|
|
# configured compilation explicitly.
|
|
from vllm.config.compilation import CompilationMode, CUDAGraphMode
|
|
|
|
compilation_config = vllm_config.compilation_config
|
|
if compilation_config.mode is None:
|
|
compilation_config.mode = CompilationMode.NONE
|
|
if compilation_config.cudagraph_mode is None:
|
|
compilation_config.cudagraph_mode = CUDAGraphMode.FULL
|
|
|
|
|
|
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
|
|
"ColBERTJinaRobertaModel": JinaRobertaModelConfig,
|
|
"ColQwen3_5": ColQwen3_5Config,
|
|
"DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
|
|
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
|
|
"DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501
|
|
"Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
|
|
"FalconMambaForCausalLM": MambaModelConfig,
|
|
"Gemma3TextModel": Gemma3TextModelConfig,
|
|
"Gemma4ForCausalLM": Gemma4Config,
|
|
"Gemma4ForConditionalGeneration": Gemma4Config,
|
|
"Gemma4UnifiedForConditionalGeneration": Gemma4Config,
|
|
"GptOssForCausalLM": GptOssForCausalLMConfig,
|
|
"LongcatFlashNgramForCausalLM": LongcatFlashNgramForCausalLMConfig,
|
|
"GteModel": SnowflakeGteNewModelConfig,
|
|
"GteNewForSequenceClassification": GteNewModelConfig,
|
|
"GteNewModel": GteNewModelConfig,
|
|
"JambaForSequenceClassification": JambaForSequenceClassificationConfig,
|
|
"JinaForRanking": JinaForRankingConfig,
|
|
"JinaVLForRanking": JinaVLForSequenceClassificationConfig,
|
|
"LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
|
|
"LlamaBidirectionalModel": LlamaBidirectionalConfig,
|
|
"LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
|
|
"LlamaNemotronVLModel": LlamaNemotronVLConfig,
|
|
"Mamba2ForCausalLM": MambaModelConfig,
|
|
"MambaForCausalLM": MambaModelConfig,
|
|
"NemotronHForCausalLM": NemotronHForCausalLMConfig,
|
|
"NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
|
|
"NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config,
|
|
"NomicBertModel": NomicBertModelConfig,
|
|
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
|
|
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
|
|
"Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
|
|
"Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
|
|
"Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
|
|
"Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
|
|
"UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig,
|
|
"VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
|
|
"XLMRobertaModel": JinaRobertaModelConfig,
|
|
}
|