iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model] Streamline EmbeddingGemma 2 config resolution and harmonize Triton prefill attention (#60289)

Browse Source 
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
This commit is contained in:
Luciano Martins
2026-10-07 14:56:40 -07:00
committed by GitHub
co-authored by Luciano Martins Matthew Bonanni
parent 947dd62c07
commit b803edbc9f
3 changed files with 57 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/multimodal/pooling/test_embedding_gemma2.py
+52 -79
View File
@@ -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
vllm/model_executor/models/config.py
+4 -84
View File
@@ -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):
vllm/transformers_utils/config.py
+1 -1
View File
@@ -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