@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f4b76056ee
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 09:45 UTC

Improve enable chunked_prefill & prefix_caching logic. (#26623)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

wang.yuqi committed Nov 28, 2025 at 14:05 UTC f4b76056ee5c3a3f917527da5be3786e1b8530c6
11 files changed +456 -133
tests/models/language/pooling/test_auto_prefix_cache_support.py
+1 -3
@@ -105,8 +105,6 @@ def test_embed_models(
105 def test_non_causal_models(
106 hf_runner, vllm_runner, example_prompts, model: str, dtype: str
107 ) -> None:
108 - with vllm_runner(
109 - model, max_model_len=512, dtype=dtype, enable_prefix_caching=True
110 - ) as vllm_model:
108 + with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:
109 cache_config = vllm_model.llm.llm_engine.cache_config
110 assert not cache_config.enable_prefix_caching
tests/test_config.py
+239 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +import logging
4 import os
5 from dataclasses import MISSING, Field, asdict, dataclass, field
6 from unittest.mock import patch
@@ -602,6 +602,244 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files):
602 assert os.path.exists(config2.tokenizer) and os.path.isdir(config2.tokenizer)
603
604
605 +@pytest.mark.parametrize(
606 + ("model_id", "expected_attn_type", "expected_result", "reason"),
607 + [
608 + # pooling models
609 + (
610 + "jason9693/Qwen2.5-1.5B-apeach",
611 + "decoder",
612 + True,
613 + "Pooling models with causal attn and last pooling support chunked prefill.",
614 + ),
615 + (
616 + "Qwen/Qwen3-Embedding-0.6B",
617 + "decoder",
618 + True,
619 + "Pooling models with causal attn and last pooling support chunked prefill.",
620 + ),
621 + (
622 + "Qwen/Qwen2.5-Math-PRM-7B",
623 + "decoder",
624 + False,
625 + "Pooling models with step pooling does not support chunked prefill.",
626 + ),
627 + (
628 + "internlm/internlm2-1_8b-reward",
629 + "decoder",
630 + False,
631 + "Pooling models with all pooling does not support chunked prefill.",
632 + ),
633 + (
634 + "BAAI/bge-base-en",
635 + "encoder_only",
636 + False,
637 + "Pooling models with bidirectional attn does not support chunked prefill.",
638 + ),
639 + (
640 + "boltuix/NeuroBERT-NER",
641 + "encoder_only",
642 + False,
643 + "Pooling models with bidirectional attn does not support chunked prefill.",
644 + ),
645 + (
646 + "papluca/xlm-roberta-base-language-detection",
647 + "encoder_only",
648 + False,
649 + "Pooling models with bidirectional attn does not support chunked prefill.",
650 + ),
651 + (
652 + "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
653 + "encoder_only",
654 + False,
655 + "Pooling models with bidirectional attn does not support chunked prefill.",
656 + ),
657 + (
658 + "intfloat/e5-small",
659 + "encoder_only",
660 + False,
661 + "Pooling models with bidirectional attn does not support chunked prefill.",
662 + ),
663 + # multimodal models
664 + (
665 + "openai/clip-vit-base-patch32",
666 + "decoder",
667 + True,
668 + "Pooling models with causal attn and last pooling support chunked prefill.",
669 + ),
670 + (
671 + "google/siglip-base-patch16-224",
672 + "encoder_only",
673 + False,
674 + "Pooling models with bidirectional attn does not support chunked prefill.",
675 + ),
676 + # generate models
677 + (
678 + "Qwen/Qwen3-0.6B",
679 + "decoder",
680 + True,
681 + "Generative models support chunked prefill.",
682 + ),
683 + (
684 + "Qwen/Qwen3-Next-80B-A3B-Instruct",
685 + "hybrid",
686 + True,
687 + "Generative models support chunked prefill.",
688 + ),
689 + (
690 + "ibm-granite/granite-4.0-h-small",
691 + "hybrid",
692 + True,
693 + "Generative models support chunked prefill.",
694 + ),
695 + (
696 + "state-spaces/mamba-130m-hf",
697 + "attention_free",
698 + True,
699 + "Generative models support chunked prefill.",
700 + ),
701 + # encoder_decoder models
702 + (
703 + "openai/whisper-small",
704 + "encoder_decoder",
705 + False,
706 + "Encoder decoder models does not support chunked prefill.",
707 + ),
708 + ],
709 +)
710 +def test_is_chunked_prefill_supported(
711 + model_id: str,
712 + expected_attn_type: str,
713 + expected_result: bool,
714 + reason: str,
715 + caplog_vllm,
716 +):
717 + model_config = ModelConfig(model_id, trust_remote_code=True)
718 + assert model_config.attn_type == expected_attn_type
719 + with caplog_vllm.at_level(level=logging.DEBUG):
720 + assert model_config.is_chunked_prefill_supported == expected_result
721 + assert reason in caplog_vllm.text
722 +
723 +
724 +@pytest.mark.parametrize(
725 + ("model_id", "expected_attn_type", "expected_result", "reason"),
726 + [
727 + # pooling models
728 + (
729 + "jason9693/Qwen2.5-1.5B-apeach",
730 + "decoder",
731 + True,
732 + "Pooling models with causal attn and last pooling support prefix caching.",
733 + ),
734 + (
735 + "Qwen/Qwen3-Embedding-0.6B",
736 + "decoder",
737 + True,
738 + "Pooling models with causal attn and last pooling support prefix caching.",
739 + ),
740 + (
741 + "Qwen/Qwen2.5-Math-PRM-7B",
742 + "decoder",
743 + False,
744 + "Pooling models with step pooling does not support prefix caching.",
745 + ),
746 + (
747 + "internlm/internlm2-1_8b-reward",
748 + "decoder",
749 + False,
750 + "Pooling models with all pooling does not support prefix caching.",
751 + ),
752 + (
753 + "BAAI/bge-base-en",
754 + "encoder_only",
755 + False,
756 + "Pooling models with bidirectional attn does not support prefix caching.",
757 + ),
758 + (
759 + "boltuix/NeuroBERT-NER",
760 + "encoder_only",
761 + False,
762 + "Pooling models with bidirectional attn does not support prefix caching.",
763 + ),
764 + (
765 + "papluca/xlm-roberta-base-language-detection",
766 + "encoder_only",
767 + False,
768 + "Pooling models with bidirectional attn does not support prefix caching.",
769 + ),
770 + (
771 + "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
772 + "encoder_only",
773 + False,
774 + "Pooling models with bidirectional attn does not support prefix caching.",
775 + ),
776 + (
777 + "intfloat/e5-small",
778 + "encoder_only",
779 + False,
780 + "Pooling models with bidirectional attn does not support prefix caching.",
781 + ),
782 + # multimodal models
783 + (
784 + "openai/clip-vit-base-patch32",
785 + "decoder",
786 + True,
787 + "Pooling models with causal attn and last pooling support prefix caching.",
788 + ),
789 + (
790 + "google/siglip-base-patch16-224",
791 + "encoder_only",
792 + False,
793 + "Pooling models with bidirectional attn does not support prefix caching.",
794 + ),
795 + # generate models
796 + (
797 + "Qwen/Qwen3-0.6B",
798 + "decoder",
799 + True,
800 + "Generative models support prefix caching.",
801 + ),
802 + (
803 + "Qwen/Qwen3-Next-80B-A3B-Instruct",
804 + "hybrid",
805 + False,
806 + "Hybrid models does not support prefix caching since the feature is still experimental.", # noqa: E501
807 + ),
808 + (
809 + "ibm-granite/granite-4.0-h-small",
810 + "hybrid",
811 + False,
812 + "Hybrid models does not support prefix caching since the feature is still experimental.", # noqa: E501
813 + ),
814 + (
815 + "state-spaces/mamba-130m-hf",
816 + "attention_free",
817 + False,
818 + "Attention free models does not support prefix caching since the feature is still experimental.", # noqa: E501
819 + ),
820 + # encoder_decoder models
821 + (
822 + "openai/whisper-small",
823 + "encoder_decoder",
824 + False,
825 + "Encoder decoder models does not support prefix caching.",
826 + ),
827 + ],
828 +)
829 +def test_is_prefix_caching_supported(
830 + model_id: str,
831 + expected_attn_type: str,
832 + expected_result: bool,
833 + reason: str,
834 + caplog_vllm,
835 +):
836 + model_config = ModelConfig(model_id, trust_remote_code=True)
837 + assert model_config.attn_type == expected_attn_type
838 + with caplog_vllm.at_level(level=logging.DEBUG):
839 + assert model_config.is_prefix_caching_supported == expected_result
840 + assert reason in caplog_vllm.text
841 +
842 +
843 @pytest.mark.parametrize(
844 ("backend", "custom_ops", "expected"),
845 [
vllm/config/model.py
+109
@@ -107,6 +107,10 @@ _RUNNER_CONVERTS: dict[RunnerType, list[ConvertType]] = {
107 "draft": [],
108 }
109
110 +AttnTypeStr = Literal[
111 + "decoder", "encoder", "encoder_only", "encoder_decoder", "attention_free", "hybrid"
112 +]
113 +
114
115 @config
116 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
@@ -1752,6 +1756,111 @@ class ModelConfig:
1756 logger.info("Using max model len %s", max_model_len)
1757 return max_model_len
1758
1759 + @property
1760 + def attn_type(self) -> AttnTypeStr:
1761 + if self.pooler_config is not None:
1762 + pooling_type = self._model_info.default_pooling_type.lower()
1763 + if pooling_type == "cls":
1764 + return "encoder_only"
1765 + else:
1766 + is_causal = getattr(self.hf_config, "is_causal", True)
1767 + return "encoder_only" if not is_causal else self._model_info.attn_type
1768 + elif self.is_hybrid:
1769 + return "hybrid"
1770 + elif self.is_attention_free:
1771 + return "attention_free"
1772 + elif self.is_encoder_decoder:
1773 + return "encoder_decoder"
1774 + else:
1775 + return "decoder"
1776 +
1777 + @property
1778 + def is_chunked_prefill_supported(self) -> bool:
1779 + attn_type = self.attn_type
1780 + if self.pooler_config is not None:
1781 + # for pooling models
1782 + if attn_type == "encoder_only":
1783 + logger.debug(
1784 + "Pooling models with bidirectional attn does not support "
1785 + "chunked prefill."
1786 + )
1787 + return False
1788 + elif attn_type == "decoder":
1789 + pooling_type = self.pooler_config.pooling_type.lower()
1790 + if pooling_type in ["all", "mean", "step", "cls"]:
1791 + logger.debug(
1792 + "Pooling models with %s pooling does not "
1793 + "support chunked prefill.",
1794 + pooling_type,
1795 + )
1796 + return False
1797 + else:
1798 + # pooling_type == "last"
1799 + logger.debug(
1800 + "Pooling models with causal attn and last pooling support "
1801 + "chunked prefill."
1802 + )
1803 + return True
1804 + # vllm currently does not have pooling models using hybrid,
1805 + # attention_free or encoder_decoder attn types.
1806 + return attn_type != "encoder_decoder"
1807 + else:
1808 + if attn_type == "encoder_decoder":
1809 + logger.debug("Encoder decoder models does not support chunked prefill.")
1810 + return False
1811 + logger.debug("Generative models support chunked prefill.")
1812 + return True
1813 +
1814 + @property
1815 + def is_prefix_caching_supported(self) -> bool:
1816 + attn_type = self.attn_type
1817 + if self.pooler_config is not None:
1818 + # for pooling models
1819 + if attn_type == "encoder_only":
1820 + logger.debug(
1821 + "Pooling models with bidirectional attn does not "
1822 + "support prefix caching."
1823 + )
1824 + return False
1825 + elif attn_type == "decoder":
1826 + pooling_type = self.pooler_config.pooling_type.lower()
1827 + if pooling_type in ["all", "mean", "step", "cls"]:
1828 + logger.debug(
1829 + "Pooling models with %s pooling does not "
1830 + "support prefix caching.",
1831 + pooling_type,
1832 + )
1833 + return False
1834 + else:
1835 + # pooling_type == "last"
1836 + logger.debug(
1837 + "Pooling models with causal attn and last pooling support "
1838 + "prefix caching."
1839 + )
1840 + return True
1841 + # vllm currently does not have pooling models using hybrid,
1842 + # attention_free or encoder_decoder attn types.
1843 + return False
1844 + else:
1845 + if attn_type == "hybrid":
1846 + logger.debug(
1847 + "Hybrid models does not support prefix caching since the feature "
1848 + "is still experimental."
1849 + )
1850 + return False
1851 + elif attn_type == "attention_free":
1852 + logger.debug(
1853 + "Attention free models does not support prefix caching since the "
1854 + "feature is still experimental."
1855 + )
1856 + return False
1857 + elif attn_type == "encoder_decoder":
1858 + logger.debug("Encoder decoder models does not support prefix caching.")
1859 + return False
1860 + else: # attn_type == "decoder"
1861 + logger.debug("Generative models support prefix caching.")
1862 + return True
1863 +
1864 def is_model_moe(
1865 self,
1866 ) -> bool:
vllm/config/pooler.py
+4 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any
4 +from typing import Any, Literal
5
6 from pydantic.dataclasses import dataclass
7
@@ -11,13 +11,15 @@ from vllm.utils.hashing import safe_hash
11
12 logger = init_logger(__name__)
13
14 +PoolingTypeStr = Literal["LAST", "ALL", "CLS", "STEP", "MEAN"]
15 +
16
17 @config
18 @dataclass
19 class PoolerConfig:
20 """Controls the behavior of output pooling in pooling models."""
21
20 - pooling_type: str | None = None
22 + pooling_type: PoolingTypeStr | None = None
23 """
24 The pooling method of the pooling model. This should be a key in
25 [`vllm.model_executor.layers.pooler.PoolingType`][].
vllm/config/vllm.py
+19 -57
@@ -721,65 +721,27 @@ class VllmConfig:
721 "correctness and to realize prefill savings. "
722 )
723
724 - disable_chunked_prefill_reasons: list[str] = []
724 + if self.model_config and self.model_config.is_encoder_decoder:
725 + from vllm.multimodal import MULTIMODAL_REGISTRY
726
726 - if self.model_config:
727 - if self.model_config.pooler_config:
728 - pooling_type = self.model_config.pooler_config.pooling_type
729 - if pooling_type is None or pooling_type.lower() != "last":
730 - disable_chunked_prefill_reasons.append(
731 - 'Only "last" pooling supports chunked '
732 - "prefill and prefix caching; disabling both."
733 - )
734 - if not getattr(self.model_config.hf_config, "is_causal", True):
735 - disable_chunked_prefill_reasons.append(
736 - "Only models using causal attention support chunked "
737 - "prefill and prefix caching; disabling both."
738 - )
739 - elif self.model_config.is_encoder_decoder:
740 - from vllm.multimodal import MULTIMODAL_REGISTRY
741 -
742 - self.scheduler_config.max_num_encoder_input_tokens = (
743 - MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(self.model_config)
744 - )
745 - logger.debug(
746 - "Encoder-decoder model detected: setting "
747 - "`max_num_encoder_input_tokens` to encoder length (%s)",
748 - self.scheduler_config.max_num_encoder_input_tokens,
749 - )
750 - if (
751 - self.model_config.architecture == "WhisperForConditionalGeneration"
752 - and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"
753 - ):
754 - logger.warning(
755 - "Whisper is known to have issues with "
756 - "forked workers. If startup is hanging, "
757 - "try setting 'VLLM_WORKER_MULTIPROC_METHOD' "
758 - "to 'spawn'."
759 - )
760 -
761 - # Final off-switch for CP/APC:
762 - # Disable for (a) collected blockers, (b) encoder–decoder, or
763 - # (c) explicit CP=False when APC wasn't requested.
764 - # Do NOT disable merely because the resolved CP flag is False.
765 - apc_requested = (
766 - self.cache_config is not None and self.cache_config.enable_prefix_caching
767 - )
768 - if (
769 - disable_chunked_prefill_reasons
770 - or (self.model_config is not None and self.model_config.is_encoder_decoder)
771 - or (
772 - self.scheduler_config.enable_chunked_prefill is False
773 - and not apc_requested
727 + self.scheduler_config.max_num_encoder_input_tokens = (
728 + MULTIMODAL_REGISTRY.get_encdec_max_encoder_len(self.model_config)
729 )
775 - ):
776 - for reason in disable_chunked_prefill_reasons:
777 - logger.info(reason)
778 - self.scheduler_config.enable_chunked_prefill = False
779 - self.scheduler_config.long_prefill_token_threshold = 0
780 -
781 - if self.cache_config is not None:
782 - self.cache_config.enable_prefix_caching = False
730 + logger.debug(
731 + "Encoder-decoder model detected: setting "
732 + "`max_num_encoder_input_tokens` to encoder length (%s)",
733 + self.scheduler_config.max_num_encoder_input_tokens,
734 + )
735 + if (
736 + self.model_config.architecture == "WhisperForConditionalGeneration"
737 + and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"
738 + ):
739 + logger.warning(
740 + "Whisper is known to have issues with "
741 + "forked workers. If startup is hanging, "
742 + "try setting 'VLLM_WORKER_MULTIPROC_METHOD' "
743 + "to 'spawn'."
744 + )
745
746 if (
747 self.kv_events_config is not None
vllm/engine/arg_utils.py
+31 -59
@@ -1349,30 +1349,10 @@ class EngineArgs:
1349 self.tokenizer = model_config.tokenizer
1350
1351 self._check_feature_supported(model_config)
1352 -
1353 - # Set default arguments for V1 Engine.
1354 - self._set_default_args(usage_context, model_config)
1355 - # Disable chunked prefill and prefix caching for:
1356 - # POWER (ppc64le)/s390x/RISCV CPUs in V1
1357 - if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
1358 - CpuArchEnum.POWERPC,
1359 - CpuArchEnum.S390X,
1360 - CpuArchEnum.RISCV,
1361 - ):
1362 - logger.info(
1363 - "Chunked prefill is not supported for ARM and POWER, "
1364 - "S390X and RISC-V CPUs; "
1365 - "disabling it for V1 backend."
1366 - )
1367 - self.enable_chunked_prefill = False
1368 - logger.info(
1369 - "Prefix caching is not supported for ARM and POWER, "
1370 - "S390X and RISC-V CPUs; "
1371 - "disabling it for V1 backend."
1372 - )
1373 - self.enable_prefix_caching = False
1374 -
1375 - assert self.enable_chunked_prefill is not None
1352 + self._set_default_chunked_prefill_and_prefix_caching_args(model_config)
1353 + self._set_default_max_num_seqs_and_batched_tokens_args(
1354 + usage_context, model_config
1355 + )
1356
1357 sliding_window: int | None = None
1358 if not is_interleaved(model_config.hf_text_config):
@@ -1805,34 +1785,6 @@ class EngineArgs:
1785 )
1786 _raise_unsupported_error(feature_name=name)
1787
1808 - @classmethod
1809 - def get_chunked_prefill_prefix_caching_defaults(
1810 - cls,
1811 - model_config: ModelConfig,
1812 - ) -> tuple[bool, bool]:
1813 - if model_config.runner_type != "pooling":
1814 - default_chunked_prefill = True
1815 -
1816 - # Disable prefix caching default for hybrid models and mamba-only
1817 - # models since the feature is still experimental.
1818 - default_prefix_caching = not (
1819 - model_config.is_hybrid or model_config.is_attention_free
1820 - )
1821 - else:
1822 - assert model_config.pooler_config is not None
1823 -
1824 - pooling_type = model_config.pooler_config.pooling_type
1825 - incremental_prefill_supported = (
1826 - pooling_type is not None
1827 - and pooling_type.lower() == "last"
1828 - and getattr(model_config.hf_config, "is_causal", True)
1829 - )
1830 -
1831 - default_chunked_prefill = incremental_prefill_supported
1832 - default_prefix_caching = incremental_prefill_supported
1833 -
1834 - return default_chunked_prefill, default_prefix_caching
1835 -
1788 @classmethod
1789 def get_batch_defaults(
1790 cls,
@@ -1916,14 +1868,11 @@ class EngineArgs:
1868
1869 return default_max_num_batched_tokens, default_max_num_seqs
1870
1919 - def _set_default_args(
1920 - self, usage_context: UsageContext, model_config: ModelConfig
1871 + def _set_default_chunked_prefill_and_prefix_caching_args(
1872 + self, model_config: ModelConfig
1873 ) -> None:
1922 - """Set Default Arguments for V1 Engine."""
1923 - (
1924 - default_chunked_prefill,
1925 - default_prefix_caching,
1926 - ) = self.get_chunked_prefill_prefix_caching_defaults(model_config)
1874 + default_chunked_prefill = model_config.is_chunked_prefill_supported
1875 + default_prefix_caching = model_config.is_prefix_caching_supported
1876
1877 if self.prefill_context_parallel_size > 1:
1878 default_chunked_prefill = False
@@ -1984,6 +1933,29 @@ class EngineArgs:
1933 scope="local",
1934 )
1935
1936 + # Disable chunked prefill and prefix caching for:
1937 + # POWER (ppc64le)/s390x/RISCV CPUs in V1
1938 + if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
1939 + CpuArchEnum.POWERPC,
1940 + CpuArchEnum.S390X,
1941 + CpuArchEnum.RISCV,
1942 + ):
1943 + logger.info(
1944 + "Chunked prefill is not supported for ARM and POWER, "
1945 + "S390X and RISC-V CPUs; "
1946 + "disabling it for V1 backend."
1947 + )
1948 + self.enable_chunked_prefill = False
1949 + logger.info(
1950 + "Prefix caching is not supported for ARM and POWER, "
1951 + "S390X and RISC-V CPUs; "
1952 + "disabling it for V1 backend."
1953 + )
1954 + self.enable_prefix_caching = False
1955 +
1956 + def _set_default_max_num_seqs_and_batched_tokens_args(
1957 + self, usage_context: UsageContext, model_config: ModelConfig
1958 + ):
1959 world_size = self.pipeline_parallel_size * self.tensor_parallel_size
1960 (
1961 default_max_num_batched_tokens,
vllm/model_executor/models/bert.py
+2 -2
@@ -32,7 +32,7 @@ from vllm.tasks import PoolingTask
32 from vllm.v1.pool.metadata import PoolingMetadata
33
34 from .interfaces import SupportsCrossEncoding, SupportsQuant
35 -from .interfaces_base import default_pooling_type
35 +from .interfaces_base import attn_type, default_pooling_type
36 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
37
38
@@ -432,7 +432,6 @@ class BertModel(nn.Module, SupportsQuant):
432 return loaded_params
433
434
435 -@default_pooling_type("ALL")
435 class BertPoolingModel(BertModel):
436 is_pooling_model = True
437
@@ -864,6 +863,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
863 )
864
865
866 +@attn_type("encoder_only")
867 @default_pooling_type("ALL")
868 class BertForTokenClassification(nn.Module):
869 is_pooling_model = True
vllm/model_executor/models/interfaces_base.py
+32 -3
@@ -19,10 +19,14 @@ from vllm.utils.func_utils import supports_kw
19
20 if TYPE_CHECKING:
21 from vllm.config import VllmConfig
22 + from vllm.config.model import AttnTypeStr
23 + from vllm.config.pooler import PoolingTypeStr
24 from vllm.model_executor.layers.pooler import Pooler
25 else:
26 VllmConfig = Any
27 Pooler = Any
28 + PoolingTypeStr = Any
29 + AttnTypeStr = Any
30
31 logger = init_logger(__name__)
32
@@ -165,7 +169,7 @@ class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
169 MRO of your model class.
170 """
171
168 - default_pooling_type: ClassVar[str] = "LAST"
172 + default_pooling_type: ClassVar[PoolingTypeStr] = "LAST"
173 """
174 Indicates the [vllm.config.pooler.PoolerConfig.pooling_type][]
175 to use by default.
@@ -175,6 +179,17 @@ class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
179 decorator to conveniently set this field.
180 """
181
182 + attn_type: ClassVar[AttnTypeStr] = "decoder"
183 + """
184 + Indicates the
185 + [vllm.config.model.ModelConfig.attn_type][]
186 + to use by default.
187 +
188 + You can use the
189 + [vllm.model_executor.models.interfaces_base.attn_type][]
190 + decorator to conveniently set this field.
191 + """
192 +
193 pooler: Pooler
194 """The pooler is only called on TP rank 0."""
195
@@ -199,7 +214,7 @@ def is_pooling_model(
214 _T = TypeVar("_T", bound=type[nn.Module])
215
216
202 -def default_pooling_type(pooling_type: str):
217 +def default_pooling_type(pooling_type: PoolingTypeStr):
218 """Decorator to set `VllmModelForPooling.default_pooling_type`."""
219
220 def func(model: _T) -> _T:
@@ -209,5 +224,19 @@ def default_pooling_type(pooling_type: str):
224 return func
225
226
212 -def get_default_pooling_type(model: type[object] | object) -> str:
227 +def get_default_pooling_type(model: type[object] | object) -> PoolingTypeStr:
228 return getattr(model, "default_pooling_type", "LAST")
229 +
230 +
231 +def attn_type(attn_type: AttnTypeStr):
232 + """Decorator to set `VllmModelForPooling.attn_type`."""
233 +
234 + def func(model: _T) -> _T:
235 + model.attn_type = attn_type # type: ignore
236 + return model
237 +
238 + return func
239 +
240 +
241 +def get_attn_type(model: type[object] | object) -> AttnTypeStr:
242 + return getattr(model, "attn_type", "decoder")
vllm/model_executor/models/modernbert.py
+2 -1
@@ -28,7 +28,7 @@ from vllm.tasks import PoolingTask
28 from vllm.v1.pool.metadata import PoolingMetadata
29
30 from .interfaces import SupportsCrossEncoding
31 -from .interfaces_base import default_pooling_type
31 +from .interfaces_base import attn_type, default_pooling_type
32 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
33
34
@@ -396,6 +396,7 @@ class ModernBertPredictionHead(nn.Module):
396 return self.norm(self.act(self.dense(hidden_states)))
397
398
399 +@attn_type("encoder_only")
400 @default_pooling_type("ALL")
401 class ModernBertForTokenClassification(nn.Module):
402 is_pooling_model = True
vllm/model_executor/models/registry.py
+13 -2
@@ -17,7 +17,7 @@ from collections.abc import Callable, Set
17 from dataclasses import asdict, dataclass, field
18 from functools import lru_cache
19 from pathlib import Path
20 -from typing import TypeVar
20 +from typing import TYPE_CHECKING, Any, TypeVar
21
22 import torch.nn as nn
23 import transformers
@@ -33,6 +33,14 @@ from vllm.logging_utils import logtime
33 from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module
34 from vllm.utils.hashing import safe_hash
35
36 +if TYPE_CHECKING:
37 + from vllm.config.model import AttnTypeStr
38 + from vllm.config.pooler import PoolingTypeStr
39 +else:
40 + AttnTypeStr = Any
41 + PoolingTypeStr = Any
42 +
43 +
44 from .interfaces import (
45 has_inner_state,
46 has_noops,
@@ -47,6 +55,7 @@ from .interfaces import (
55 supports_transcription,
56 )
57 from .interfaces_base import (
58 + get_attn_type,
59 get_default_pooling_type,
60 is_pooling_model,
61 is_text_generation_model,
@@ -509,7 +518,8 @@ class _ModelInfo:
518 architecture: str
519 is_text_generation_model: bool
520 is_pooling_model: bool
512 - default_pooling_type: str
521 + attn_type: AttnTypeStr
522 + default_pooling_type: PoolingTypeStr
523 supports_cross_encoding: bool
524 supports_multimodal: bool
525 supports_multimodal_raw_input_only: bool
@@ -530,6 +540,7 @@ class _ModelInfo:
540 is_text_generation_model=is_text_generation_model(model),
541 is_pooling_model=is_pooling_model(model),
542 default_pooling_type=get_default_pooling_type(model),
543 + attn_type=get_attn_type(model),
544 supports_cross_encoding=supports_cross_encoding(model),
545 supports_multimodal=supports_multimodal(model),
546 supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
vllm/v1/engine/core.py
+4 -3
@@ -119,11 +119,12 @@ class EngineCore:
119 # Setup scheduler.
120 Scheduler = vllm_config.scheduler_config.get_scheduler_cls()
121
122 - if len(kv_cache_config.kv_cache_groups) == 0:
122 + if len(kv_cache_config.kv_cache_groups) == 0: # noqa: SIM102
123 # Encoder models without KV cache don't support
124 # chunked prefill. But do SSM models?
125 - logger.info("Disabling chunked prefill for model without KVCache")
126 - vllm_config.scheduler_config.enable_chunked_prefill = False
125 + if vllm_config.scheduler_config.enable_chunked_prefill:
126 + logger.warning("Disabling chunked prefill for model without KVCache")
127 + vllm_config.scheduler_config.enable_chunked_prefill = False
128
129 scheduler_block_size = (
130 vllm_config.cache_config.block_size