spark2_5: size the RoPE table to the servable context, not the advertised one

618 MB of resident memory and 649 MB of peak, for one line -- the largest
single memory item found in this study, and unlike the others it costs nothing.

Spark advertises a 1,048,576-token context. This file passed
max_position=config.max_position_embeddings straight into get_rope, and vLLM
builds its cos/sin table with torch.arange(max_position), so a 2048-token
deployment carried tables for a million positions. Because the sliding and full
layers are trained with different theta and different partial rotary factors
there are two of them:

    layers.0 (full, head_dim 256)      [1048576, 256]  512 MB
    layers.3 (sliding, partial 0.25)   [1048576,  64]  128 MB

Both fully resident -- confirmed with mincore(2), not inferred. At
max_model_len 2048 they need 1.25 MB between them.

Found by memory_ledger.py, which attributes memory instead of subtracting
toward a remainder: 640 MB had been sitting in an unnamed "other" bucket across
114 tensors, and naming the largest storages showed what it was. After the cap
that bucket is 1.3 MB and unique tensor storage falls 638.8 MB, which agrees
independently with the end-to-end measurement.

vLLM's scheduler cannot produce a position beyond max_model_len, so covering
that is sufficient. Measured with VLLM_OMNI_SPARK_ROPE_FULL toggling the two in
identical code, 3 interleaved passes per arm on node-local cores:

    1M tables   settled 3891 MB   peak 4019 MB   init 13.5 s
    capped      settled 3273 MB   peak 3370 MB   init 13.7 s

The escape hatch restores the old behaviour and is covered by a test, as is the
cap itself.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude Opus 5 committed 2026-09-14 08:40:44 +08:00
1 parent bc865385a8
commit df3bbfe675
2 files changed
+64 -1

No files matched your search

@@ -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"
@@ -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,
)