[misc] Unify MLA scaling init and remove dead buffer / scaling code (#33363)

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
+            # In transformers v5, rope_scaling is just rope_parameters.
-            self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
+            self._init_mla_scaling(self.hf_text_config.rope_scaling)
-            # 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
-                    )
         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)
+            self._init_mla_scaling(self.hf_config.rope_scaling)
-            if self.hf_config.rope_scaling:
-                self.scaling = compute_mla_mscale_scaling(
-                    self.hf_config.rope_scaling, self.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)
+            self._init_mla_scaling(getattr(tc, "rope_scaling", None))
-            if getattr(tc, "rope_scaling", None):
-                self.scaling = compute_mla_mscale_scaling(tc.rope_scaling, self.scaling)
         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._init_mla_scaling(self.hf_config.rope_scaling)
-            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
-                )
         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)
+            self._init_mla_scaling(self.hf_config.rope_scaling)
-            if self.hf_config.rope_scaling:
-                self.scaling = compute_mla_mscale_scaling(
-                    self.hf_config.rope_scaling, self.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

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

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

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

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
 
 

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.
+# state pool, so stale content is unsafe to execute. build_decode_registry
-_INDEX_SEMANTIC_BUFFERS = frozenset(
+# 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, share_input_buffer(name, buffer))
-                setattr(self, name, new_buffer)

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()