[RL] convert DeepSeek V4 APE layout through weight loader (#27307)

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)