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

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-05 15:59:33 -07:00
committed by GitHub
parent 990a446773
commit c0ef548eef
7 changed files with 69 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+14 -34
View File
@@ -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
python/sglang/srt/layers/attention/cutlass_mla_backend.py
-1
View File
@@ -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
-1
View File
@@ -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
-1
View File
@@ -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
+20 -1
View File
@@ -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
+6 -35
View File
@@ -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))