diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 15b84ed23..3d80a871c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -855,21 +855,8 @@ class ModelConfig: if is_deepseek_dsa(self.hf_text_config) else None ) - # Handle rope scaling - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - # in transformers v5, rope_scaling is just rope_parameters for backward compatibility - rope_scaling = self.hf_text_config.rope_scaling - if rope_scaling: - # v5 uses "rope_type", v4 uses "type" - rope_type = ( - rope_scaling.get("rope_type") - or rope_scaling.get("type") - or "default" - ) - if rope_type != "default": - self.scaling = compute_mla_mscale_scaling( - rope_scaling, self.scaling - ) + # In transformers v5, rope_scaling is just rope_parameters. + self._init_mla_scaling(self.hf_text_config.rope_scaling) elif ( "DeepseekV4ForCausalLM" in self.hf_config.architectures or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures @@ -883,11 +870,7 @@ class ModelConfig: self.index_head_dim = self.hf_config.index_head_dim self.compress_ratios = self.hf_config.compress_ratios self.attention_arch = AttentionArch.MHA - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if self.hf_config.rope_scaling: - self.scaling = compute_mla_mscale_scaling( - self.hf_config.rope_scaling, self.scaling - ) + self._init_mla_scaling(self.hf_config.rope_scaling) elif "Glm4MoeForCausalLMNextN" in self.hf_config.architectures: if self.head_dim is None: self.head_dim = ( @@ -930,9 +913,7 @@ class ModelConfig: self.qk_rope_head_dim = tc.qk_rope_head_dim self.v_head_dim = tc.v_head_dim self.qk_nope_head_dim = tc.qk_nope_head_dim - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if getattr(tc, "rope_scaling", None): - self.scaling = compute_mla_mscale_scaling(tc.rope_scaling, self.scaling) + self._init_mla_scaling(getattr(tc, "rope_scaling", None)) elif ( "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures or "BailingMoeForCausalLMNextN" in self.hf_config.architectures @@ -943,12 +924,7 @@ class ModelConfig: self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.v_head_dim = self.hf_config.v_head_dim - # Handle rope scaling with yarn - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if self.hf_config.rope_scaling: - self.scaling = compute_mla_mscale_scaling( - self.hf_config.rope_scaling, self.scaling - ) + self._init_mla_scaling(self.hf_config.rope_scaling) elif "SarvamMLAForCausalLM" in self.hf_config.architectures: self.head_dim = ( self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim @@ -958,11 +934,7 @@ class ModelConfig: self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim self.v_head_dim = self.hf_config.v_head_dim - self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - if self.hf_config.rope_scaling: - self.scaling = compute_mla_mscale_scaling( - self.hf_config.rope_scaling, self.scaling - ) + self._init_mla_scaling(self.hf_config.rope_scaling) else: if ( "MistralModel" in self.hf_config.architectures @@ -1039,6 +1011,12 @@ class ModelConfig: if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": self.vocab_size = self.hf_text_config.vision_vocab_size + def _init_mla_scaling(self, rope_scaling: Optional[dict]) -> None: + """Base MLA attention scale from the head dims, then the rope mscale.""" + self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) + if rope_scaling: + self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling) + def get_total_num_attention_heads(self) -> int: return self.num_attention_heads @@ -1964,7 +1942,9 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models. Transformers v5 also exposes the default RoPE parameters through ``rope_scaling``. Those parameters do not request any scaling. + Warns if 'factor' is missing from a scaling request (common in v5 configs). """ + # v5 uses "rope_type", v4 uses "type" rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type") if rope_type == "default": return base_scaling diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py index 0bd908a8c..410019ce2 100644 --- a/python/sglang/srt/layers/attention/cutlass_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutlass_mla_backend.py @@ -77,7 +77,6 @@ class CutlassMLABackend(FlashInferMLAAttnBackend): self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim self.v_head_dim = model_runner.model_config.v_head_dim - self.scaling = model_runner.model_config.scaling self.data_type = model_runner.kv_cache_dtype self.q_data_type = model_runner.dtype self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 65106cbf5..92fa2bf29 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -85,7 +85,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend): self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim self.v_head_dim = model_runner.model_config.v_head_dim - self.scaling = model_runner.model_config.scaling self.data_type = model_runner.kv_cache_dtype self.q_data_type = model_runner.dtype self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 2fbafa3e5..e16fbad31 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -223,7 +223,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): # Runtime parameters self.backend = backend - self.scaling = config.scaling self.data_type = model_runner.kv_cache_dtype self.q_data_type = model_runner.dtype self.page_size = model_runner.page_size diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 338dae157..5e60753b4 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -32,7 +32,10 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple import torch -from sglang.srt.model_executor.input_buffers import share_input_buffer +from sglang.srt.model_executor.input_buffers import ( + INDEX_SEMANTIC_BUFFERS, + share_input_buffer, +) if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -775,6 +778,22 @@ def build_decode_registry( bind=canary, ) + # ZERO covers replay; a mid-serving recapture is covered by + # ForwardInputBuffers.reset_index_buffers, which keys off + # INDEX_SEMANTIC_BUFFERS. Diverge and one of the two skips a buffer. + registered = set(reg.slot_names()) + zero_slots = { + name + for name in registered + if reg.get_slot(name).padding_policy is PaddingPolicy.ZERO + } + expected_zero = INDEX_SEMANTIC_BUFFERS & registered + assert zero_slots == expected_zero, ( + "ZERO-policy slots and INDEX_SEMANTIC_BUFFERS disagree: " + f"zero_but_unlisted={sorted(zero_slots - expected_zero)}, " + f"listed_but_not_zero={sorted(expected_zero - zero_slots)}" + ) + return reg diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index 1a6702126..15d74385f 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -41,34 +41,10 @@ def share_input_buffer(name: str, new_buffer: torch.Tensor) -> torch.Tensor: return canonical.as_strided(new_buffer.size(), new_buffer.stride()) -def share_input_buffers_in(obj) -> None: - """Pool every tensor buffer on ``obj`` (dataclass / ``SimpleNamespace``) - through the process-wide pool, in place. No-op on NPU; recurses into dict / - dataclass buffer fields (``pp_proxy_tensors`` / ``ngram_embedding_info``).""" - if is_npu(): - return - - for name, buffer in list(vars(obj).items()): - if buffer is None: - continue - if dataclasses.is_dataclass(buffer): - buffer = vars(buffer) - if isinstance(buffer, dict): - for sub_name, sub_buffer in buffer.items(): - assert isinstance( - sub_buffer, torch.Tensor - ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}." - buffer[sub_name] = share_input_buffer(f"{name}.{sub_name}", sub_buffer) - else: - assert isinstance( - buffer, torch.Tensor - ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}." - setattr(obj, name, share_input_buffer(name, buffer)) - - # Values that index the rope table, the KV pool, req_to_token, or the mamba -# state pool, so stale content is unsafe to execute. -_INDEX_SEMANTIC_BUFFERS = frozenset( +# state pool, so stale content is unsafe to execute. build_decode_registry +# asserts its ZERO-policy slots against this set. +INDEX_SEMANTIC_BUFFERS = frozenset( { "positions", "mrope_positions", @@ -83,13 +59,10 @@ _INDEX_SEMANTIC_BUFFERS = frozenset( @dataclass class ForwardInputBuffers: - def _share_one_buffer(self, name: str, new_buffer: torch.Tensor) -> torch.Tensor: - return share_input_buffer(name, new_buffer) - def reset_index_buffers(self) -> None: """Zero the index-semantic buffers this set declares.""" for f in fields(self): - if f.name not in _INDEX_SEMANTIC_BUFFERS: + if f.name not in INDEX_SEMANTIC_BUFFERS: continue buffer = getattr(self, f.name) if buffer is not None: @@ -115,13 +88,11 @@ class ForwardInputBuffers: assert isinstance( sub_buffer, torch.Tensor ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}." - new_buffer = self._share_one_buffer( + buffer[sub_name] = share_input_buffer( f"{name}.{sub_name}", sub_buffer ) - buffer[sub_name] = new_buffer else: assert isinstance( buffer, torch.Tensor ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}." - new_buffer = self._share_one_buffer(name, buffer) - setattr(self, name, new_buffer) + setattr(self, name, share_input_buffer(name, buffer)) diff --git a/test/registered/unit/configs/test_model_config_scaling.py b/test/registered/unit/configs/test_model_config_scaling.py index 4d3c0fa7d..f458b47fd 100644 --- a/test/registered/unit/configs/test_model_config_scaling.py +++ b/test/registered/unit/configs/test_model_config_scaling.py @@ -1,13 +1,22 @@ import math import unittest -from sglang.srt.configs.model_config import compute_mla_mscale_scaling +from sglang.srt.configs.model_config import ModelConfig, compute_mla_mscale_scaling from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=1, suite="base-a-test-cpu") +def _mla_scaling(rope_scaling) -> tuple[float, float]: + """Run ModelConfig._init_mla_scaling on fixed head dims -> (base, result).""" + config = ModelConfig.__new__(ModelConfig) + config.qk_nope_head_dim = 128 + config.qk_rope_head_dim = 64 + config._init_mla_scaling(rope_scaling) + return 1 / math.sqrt(192), config.scaling + + class TestMlaMscaleScaling(CustomTestCase): def test_ignores_transformers_v5_default_rope_parameters(self): base_scaling = 1 / math.sqrt(72) @@ -65,5 +74,24 @@ class TestMlaMscaleScaling(CustomTestCase): ) +class TestInitMlaScaling(CustomTestCase): + """ModelConfig._init_mla_scaling must not re-add a "default" fallback of its + own: DeepseekV2AttentionMLA stamps rope_type="deepseek_yarn" on any non-empty + rope_scaling, so a dict without rope_type/type still gets a yarn rope and its + mscale belongs in the scale FlashInferMLA reads as sm_scale.""" + + def test_applies_mscale_without_rope_type(self): + base, scaling = _mla_scaling({"factor": 40, "mscale_all_dim": 1}) + self.assertGreater(scaling, base) + + def test_ignores_default_rope_type(self): + base, scaling = _mla_scaling({"rope_type": "default", "factor": 40}) + self.assertEqual(scaling, base) + + def test_no_rope_scaling_keeps_base(self): + base, scaling = _mla_scaling(None) + self.assertEqual(scaling, base) + + if __name__ == "__main__": unittest.main()