From b803edbc9fa5adbbf31136bf128d88533f7d39c2 Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:56:40 -0300 Subject: [PATCH] [Model] Streamline EmbeddingGemma 2 config resolution and harmonize Triton prefill attention (#60289) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins Co-authored-by: Matthew Bonanni --- .../pooling/test_embedding_gemma2.py | 131 +++++++----------- vllm/model_executor/models/config.py | 88 +----------- vllm/transformers_utils/config.py | 2 +- 3 files changed, 57 insertions(+), 164 deletions(-) diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 701d0914dd..0dcb72406d 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -573,9 +573,10 @@ def test_video_media_io_embedding_gemma2(): assert meta_conn.get("do_sample_frames") is False -def test_config_capping_and_explicit_limits(caplog): - """Verify EmbeddingGemma2ModelConfig capping and explicit user overrides.""" - import logging +def test_config_capping_and_explicit_limits(): + """Verify EmbeddingGemma2ModelConfig backend selection, default capping, + and preservation of explicit limits. + """ from types import SimpleNamespace from unittest.mock import patch @@ -585,19 +586,13 @@ def test_config_capping_and_explicit_limits(caplog): ) from vllm.v1.attention.backends.registry import AttentionBackendEnum - # (a) Defaults: uncapped max_model_len=262144 -> capped to 8192, - # and derived scheduler limits (>= 262144) lowered to 8192. + # (a) When attention backend is None, default to TRITON_ATTN, + # and default capping: max_model_len > 8192 without explicit limit capped to 8192 cfg_default = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=262144, - max_num_encoder_input_tokens=262144, - encoder_cache_size=262144, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=262144, + original_max_model_len=None, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): @@ -605,83 +600,61 @@ def test_config_capping_and_explicit_limits(caplog): assert cfg_default.attention_config.backend == AttentionBackendEnum.TRITON_ATTN assert cfg_default.model_config.max_model_len == 8192 - assert cfg_default.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_default.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_default.scheduler_config.encoder_cache_size == 8192 - # (b) Reachable explicit limits: explicit max_model_len=8192 and explicit - # max_num_batched_tokens=32768 stays 32768. - # Note: omitting max_model_len while passing max_num_batched_tokens < 262144 - # is rejected earlier by SchedulerConfig validation (32768 < 262144) before - # this hook runs. - cfg_explicit_batched = SimpleNamespace( + # (b) Default capping with original_max_model_len == -1 + cfg_default_minus1 = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=8192, original_max_model_len=8192 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=32768, - max_num_encoder_input_tokens=32768, - encoder_cache_size=32768, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=131072, + original_max_model_len=-1, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_batched) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_default_minus1) - assert cfg_explicit_batched.model_config.max_model_len == 8192 - assert cfg_explicit_batched.scheduler_config.max_num_batched_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.max_num_encoder_input_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.encoder_cache_size == 32768 + assert cfg_default_minus1.model_config.max_model_len == 8192 - # (c) Explicit value > uncapped_len (e.g. 300000) with uncapped model len is lowered - # to 8192 and logs a warning because it exceeds uncapped_len. - cfg_explicit_large = SimpleNamespace( + # (c) Explicit limit 4096 preserved + cfg_explicit_4k = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=300000, - max_num_encoder_input_tokens=300000, - encoder_cache_size=300000, - max_num_seqs=256, - verify_max_model_len=lambda x: None, - ), - ) - with ( - caplog.at_level(logging.WARNING), - patch.object(Gemma4Config, "verify_and_update_config"), - ): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_large) - - assert cfg_explicit_large.model_config.max_model_len == 8192 - assert cfg_explicit_large.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_explicit_large.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_explicit_large.scheduler_config.encoder_cache_size == 8192 - assert any( - "lowering explicitly set scheduler_config.max_num_batched_tokens" - " from 300000 to 8192" in rec.message - for rec in caplog.records - ) - - # (d) Explicit max_model_len=16384 is respected and not capped to 8192 - cfg_explicit_model_len = SimpleNamespace( - attention_config=SimpleNamespace(backend=None), - model_config=SimpleNamespace( - model="dummy", max_model_len=16384, original_max_model_len=16384 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=16384, - max_num_encoder_input_tokens=16384, - encoder_cache_size=16384, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=4096, + original_max_model_len=4096, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_model_len) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_4k) - assert cfg_explicit_model_len.model_config.max_model_len == 16384 - assert cfg_explicit_model_len.scheduler_config.max_num_batched_tokens == 16384 + assert cfg_explicit_4k.model_config.max_model_len == 4096 + + # (d) Explicit limit 16384 preserved and explicit backend preserved + cfg_explicit_16k = SimpleNamespace( + attention_config=SimpleNamespace(backend=AttentionBackendEnum.FLASH_ATTN), + model_config=SimpleNamespace( + max_model_len=16384, + original_max_model_len=16384, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_16k) + + assert cfg_explicit_16k.attention_config.backend == AttentionBackendEnum.FLASH_ATTN + assert cfg_explicit_16k.model_config.max_model_len == 16384 + + +def test_sentence_transformer_tokenizer_config_max_seq_length(tmp_path): + """Verify get_sentence_transformer_tokenizer_config accepts config + with only max_seq_length. + """ + import json + + from vllm.transformers_utils.config import ( + get_sentence_transformer_tokenizer_config, + ) + + # 1. Synthetic model directory with only max_seq_length in sentence_bert_config.json + cfg_file = tmp_path / "sentence_bert_config.json" + cfg_file.write_text(json.dumps({"max_seq_length": 8192})) + loaded = get_sentence_transformer_tokenizer_config(str(tmp_path)) + assert loaded is not None + assert loaded["max_seq_length"] == 8192 diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 1a1bf92061..7f6a3aef40 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -309,90 +309,10 @@ class EmbeddingGemma2ModelConfig(Gemma4Config): ) model_config = vllm_config.model_config - orig_len = getattr(model_config, "original_max_model_len", None) - if orig_len in (None, -1): - default_len = 8192 - # Check sentence_bert_config.json if available - try: - import json - from pathlib import Path - - cfg_path = Path(model_config.model) / "sentence_bert_config.json" - if cfg_path.exists(): - data = json.loads(cfg_path.read_text()) - if "max_seq_length" in data and isinstance( - data["max_seq_length"], int - ): - default_len = data["max_seq_length"] - except Exception: - pass - - if model_config.max_model_len > default_len: - uncapped_len = model_config.max_model_len - logger.info( - "EmbeddingGemma2: max_model_len not explicitly set; capping from " - "%d to %d.", - model_config.max_model_len, - default_len, - ) - model_config.max_model_len = default_len - scheduler_config = vllm_config.scheduler_config - if scheduler_config is not None: - new_batched_tokens = max( - model_config.max_model_len, scheduler_config.max_num_seqs - ) - if scheduler_config.max_num_batched_tokens >= uncapped_len: - if scheduler_config.max_num_batched_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - else: - logger.info( - "EmbeddingGemma2: lowering " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_batched_tokens = new_batched_tokens - if hasattr(vllm_config, "_set_compile_ranges"): - vllm_config._set_compile_ranges() - if scheduler_config.max_num_encoder_input_tokens >= uncapped_len: - if scheduler_config.max_num_encoder_input_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_encoder_input_tokens " - "from %d to %d to match capped max_model_len.", - scheduler_config.max_num_encoder_input_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_encoder_input_tokens = ( - new_batched_tokens - ) - if scheduler_config.encoder_cache_size >= uncapped_len: - if scheduler_config.encoder_cache_size > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.encoder_cache_size from %d to %d " - "to match capped max_model_len.", - scheduler_config.encoder_cache_size, - new_batched_tokens, - ) - scheduler_config.encoder_cache_size = new_batched_tokens - if hasattr(scheduler_config, "verify_max_model_len"): - scheduler_config.verify_max_model_len( - model_config.max_model_len - ) - elif model_config.max_model_len > 32768: - logger.warning( - "EmbeddingGemma2: max_model_len=%d is large; " - "consider --max-model-len 8192 to bound encoder memory.", - model_config.max_model_len, - ) + if model_config.max_model_len > 8192 and getattr( + model_config, "original_max_model_len", None + ) in (None, -1): + model_config.max_model_len = 8192 class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 4931204714..6a339297ee 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1139,7 +1139,7 @@ def get_sentence_transformer_tokenizer_config( logger.info("Found sentence-transformers tokenize configuration.") - if all(k in encoder_dict for k in ("max_seq_length", "do_lower_case")): + if "max_seq_length" in encoder_dict: return encoder_dict return None