From cd075d1f64837117810a10d3c7ee1831344db3ad Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 11 Jun 2026 15:05:27 -0700 Subject: [PATCH] [RL] convert DeepSeek V4 APE layout through weight loader (#27307) --- .../srt/layers/attention/dsv4/compressor.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 73a58e52f..d3e87f906 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -28,7 +28,7 @@ from sglang.srt.mem_cache.deepseek_v4_compress_state import ( ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.models.deepseek_v2 import _is_hip -from sglang.srt.utils import add_prefix, get_bool_env_var +from sglang.srt.utils import add_prefix, get_bool_env_var, set_weight_attrs _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _tgemm = None @@ -369,6 +369,7 @@ class Compressor(nn.Module): self.ape = nn.Parameter( torch.empty(self.ratio, coff * self.head_dim, dtype=torch.float32) ) + set_weight_attrs(self.ape, {"weight_loader": self.load_ape_weight}) wkv_gate_dtype = torch.bfloat16 self.wkv_gate = ReplicatedLinear( self.dim, @@ -386,8 +387,7 @@ class Compressor(nn.Module): self.ape_converted = False - def apply_ape_hotfix(self): - assert not self.ape_converted + def _apply_ape_hotfix(self): self.ape_converted = True if self.overlap: @@ -395,6 +395,16 @@ class Compressor(nn.Module): ape = torch.cat([ape[0], ape[1]], dim=0) self.ape.data.copy_(ape.view(self.ratio, -1)) + def apply_ape_hotfix(self): + assert not self.ape_converted + self._apply_ape_hotfix() + + def load_ape_weight(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + assert param is self.ape + assert loaded_weight.shape == param.shape + param.data.copy_(loaded_weight) + self._apply_ape_hotfix() + def get_state_pool(self, attn_backend: AttentionBackend) -> CompressStatePool: token_to_kv_pool = attn_backend.token_to_kv_pool assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)