chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: Support mrope_section with rope_type: "yarn" (#13313)

Browse Source 
Signed-off-by: Raayan Dhar raayan.dhar@gmail.com <raayan.dhar@gmail.com>
Signed-off-by: raayandhar <raayan.dhar@gmail.com>
This commit is contained in:
Raayan Dhar
2026-02-14 06:51:44 -08:00
committed by GitHub
parent f51e9d9ca1
commit 92cdd398cd
2 changed files with 128 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -2710,6 +2710,83 @@ class MRotaryEmbedding(RotaryEmbedding):
return llm_pos_ids
# Adapted from https://github.com/vllm-project/vllm/blob/3779eb8c81449b924a23457fc77e45a0e6171178/vllm/model_executor/layers/rotary_embedding.py#L554
class YaRNScalingMRotaryEmbedding(MRotaryEmbedding):
"""MRoPE-enabled rotary embedding with YaRN context scaling."""
def __init__(
self,
head_size: int,
rotary_dim: int,
max_position_embeddings: int,
base: int,
is_neox_style: bool,
scaling_factor: float,
dtype: torch.dtype,
*,
mrope_section: Optional[List[int]] = None,
mrope_interleaved: bool = False,
extrapolation_factor: float = 1,
attn_factor: float = 1,
beta_fast: int = 32,
beta_slow: int = 1,
truncate: bool = True,
) -> None:
self.scaling_factor = scaling_factor
self.extrapolation_factor = extrapolation_factor
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
self.truncate = truncate
self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor)
super().__init__(
head_size,
rotary_dim,
max_position_embeddings,
base,
is_neox_style,
dtype,
mrope_section=mrope_section,
mrope_interleaved=mrope_interleaved,
)
def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
pos_freqs = self.base ** (
torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
)
inv_freq_extrapolation = 1.0 / pos_freqs
inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
low, high = _yarn_find_correction_range(
self.beta_fast,
self.beta_slow,
self.rotary_dim,
self.base,
self.max_position_embeddings,
self.truncate,
)
inv_freq_mask = (
1
- _yarn_linear_ramp_mask(low, high, self.rotary_dim // 2, dtype=torch.float)
) * self.extrapolation_factor
inv_freq = (
inv_freq_interpolation * (1 - inv_freq_mask)
+ inv_freq_extrapolation * inv_freq_mask
)
return inv_freq
def _compute_cos_sin_cache(self) -> torch.Tensor:
inv_freq = self._compute_inv_freq(self.scaling_factor)
t = torch.arange(
self.max_position_embeddings * self.scaling_factor, dtype=torch.float32
)
freqs = torch.einsum("i,j -> ij", t, inv_freq)
cos = freqs.cos() * self.mscale
sin = freqs.sin() * self.mscale
cache = torch.cat((cos, sin), dim=-1)
return cache
class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
"""3D rotary positional embedding. [h w h w h w h w... t t t...]"""
@@ -3162,16 +3239,30 @@ def get_rope(
in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow")
}
extra_kwargs["truncate"] = rope_scaling.get("truncate", True)
rotary_emb = YaRNScalingRotaryEmbedding(
head_size,
rotary_dim,
original_max_position,
base,
is_neox_style,
scaling_factor,
dtype,
**extra_kwargs,
)
if "mrope_section" in rope_scaling:
rotary_emb = YaRNScalingMRotaryEmbedding(
head_size,
rotary_dim,
original_max_position,
base,
is_neox_style,
scaling_factor,
dtype,
mrope_section=rope_scaling["mrope_section"],
mrope_interleaved=rope_scaling.get("mrope_interleaved", False),
**extra_kwargs,
)
else:
rotary_emb = YaRNScalingRotaryEmbedding(
head_size,
rotary_dim,
original_max_position,
base,
is_neox_style,
scaling_factor,
dtype,
**extra_kwargs,
)
elif scaling_type == "deepseek_yarn":
scaling_factor = rope_scaling["factor"]
original_max_position = rope_scaling["original_max_position_embeddings"]

29
test/registered/rotary/test_mrope.py
View File

@@ -74,6 +74,22 @@ MODELS_TO_TEST = [
num_tokens_list = [11, 8192]
def create_yarn_rope_scaling(original_config, scaling_factor=2.0):
yarn_config = {
"rope_type": "yarn",
"factor": scaling_factor,
"original_max_position_embeddings": original_config.max_position_embeddings,
}
if hasattr(original_config, "rope_scaling") and original_config.rope_scaling:
if "mrope_section" in original_config.rope_scaling:
yarn_config["mrope_section"] = original_config.rope_scaling["mrope_section"]
if "mrope_interleaved" in original_config.rope_scaling:
yarn_config["mrope_interleaved"] = original_config.rope_scaling[
"mrope_interleaved"
]
return yarn_config
@pytest.mark.skipif(not (_is_cuda or _is_hip), reason="Skipping CUDA/ROCm only tests.")
@pytest.mark.parametrize(
"model_info, model_name",
@@ -85,12 +101,16 @@ num_tokens_list = [11, 8192]
@pytest.mark.parametrize("tp_size", [1, 2])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize("num_tokens", num_tokens_list)
@pytest.mark.parametrize(
"rope_scaling_type", ["default", "yarn"], ids=["mrope_default", "mrope_yarn"]
)
def test_mrope(
model_name: str,
model_info: MRoPETestInfo,
tp_size: int,
dtype: torch.dtype,
num_tokens: int,
rope_scaling_type: str,
):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
@@ -117,13 +137,18 @@ def test_mrope(
partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
rotary_dim = int(head_dim * partial_rotary_factor)
if rope_scaling_type == "yarn":
rope_scaling_config = create_yarn_rope_scaling(config, scaling_factor=2.0)
else:
rope_scaling_config = config.rope_scaling
mrope_helper_class = get_rope(
head_size=head_dim,
rotary_dim=rotary_dim,
max_position=max_position,
base=rope_theta,
is_neox_style=is_neox_style,
rope_scaling=config.rope_scaling,
rope_scaling=rope_scaling_config,
dtype=dtype,
).to(device=device)
@@ -133,7 +158,7 @@ def test_mrope(
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
)
query_native, key_native = mrope_helper_class._forward_native(
query_native, key_native = mrope_helper_class.forward_native(
positions,
query.clone(),
key.clone(),