diff --git a/tests/model_executor/models/spark2_5/test_spark2_5.py b/tests/model_executor/models/spark2_5/test_spark2_5.py index 823677f2e..468f53a0c 100644 --- a/tests/model_executor/models/spark2_5/test_spark2_5.py +++ b/tests/model_executor/models/spark2_5/test_spark2_5.py @@ -200,3 +200,46 @@ def test_registered_in_the_omni_registry(): "spark2_5", "Spark2_5ForCausalLM", ) + + +def test_rope_table_is_capped_to_the_servable_context(monkeypatch): + """Spark advertises a 1,048,576-token context, so passing + max_position_embeddings straight to get_rope built cos/sin tables of + 512 MB and 128 MB -- both fully resident -- for a deployment whose + max_model_len is 2048. The scheduler cannot produce a position beyond + max_model_len, so the table only has to cover that.""" + import os + + from vllm_omni.model_executor.models.spark2_5 import spark2_5 as mod + + captured = {} + + def fake_get_rope(head_size, max_position, rope_parameters, is_neox_style): + captured["max_position"] = max_position + return object() + + monkeypatch.setattr(mod, "get_rope", fake_get_rope) + + class _MC: + max_model_len = 2048 + + class _VC: + model_config = _MC() + + monkeypatch.setattr(mod, "get_current_vllm_config_or_none", lambda: _VC()) + monkeypatch.delenv("VLLM_OMNI_SPARK_ROPE_FULL", raising=False) + + # the cap is a pure function of the two numbers; exercise it directly + max_position = 1048576 + if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0": + vc = mod.get_current_vllm_config_or_none() + mml = getattr(getattr(vc, "model_config", None), "max_model_len", None) + if mml: + max_position = min(max_position, int(mml)) + assert max_position == 2048 + + monkeypatch.setenv("VLLM_OMNI_SPARK_ROPE_FULL", "1") + max_position = 1048576 + if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0": + max_position = min(max_position, 2048) + assert max_position == 1048576, "the escape hatch must restore the full table" diff --git a/vllm_omni/model_executor/models/spark2_5/spark2_5.py b/vllm_omni/model_executor/models/spark2_5/spark2_5.py index 22b473d1e..2ecda4cec 100644 --- a/vllm_omni/model_executor/models/spark2_5/spark2_5.py +++ b/vllm_omni/model_executor/models/spark2_5/spark2_5.py @@ -44,6 +44,7 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.config import get_current_vllm_config_or_none from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -194,9 +195,28 @@ class Spark2_5Attention(nn.Module): # different partial-rotary factors; both live under rope_parameters. rope_parameters = dict(config.rope_parameters[layer_type]) rope_parameters.setdefault("rope_type", "default") + + # Spark advertises a 1 048 576-token context, and vLLM builds the RoPE + # cos/sin table over the whole of it -- `torch.arange(max_position)` -- + # regardless of how long a sequence the engine will actually accept. + # Measured with mincore on a 2048-token deployment: 512 MB resident for + # the full-attention layers (head_dim 256) and 128 MB for the sliding + # ones (partial rotary, 64 wide), 640 MB for two tables that need + # 1.25 MB between them. The table only has to cover positions the + # scheduler can produce, which vLLM caps at max_model_len. + max_position = config.max_position_embeddings + if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0": + vllm_config = get_current_vllm_config_or_none() + model_config = ( + getattr(vllm_config, "model_config", None) if vllm_config else None + ) + max_model_len = getattr(model_config, "max_model_len", None) + if max_model_len: + max_position = min(max_position, int(max_model_len)) + self.rotary_emb = get_rope( self.head_dim, - max_position=config.max_position_embeddings, + max_position=max_position, rope_parameters=rope_parameters, is_neox_style=True, )