[NPU]glm5.2 fp8 memory opt (#38807)

python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py

@@ -10,6 +10,7 @@ from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_token_to_kv_pool,
 )
+from sglang.srt.runtime_context import get_disagg
 from sglang.srt.utils import get_bool_env_var
 
 if TYPE_CHECKING:
@@ -328,6 +329,9 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
             qkv_weight[:, self.q_lora_rank :].contiguous()
         )
 
+        if get_disagg().disaggregation_mode != "null":
+            qkv_weight.data.untyped_storage().resize_(0)
+
     def get_sin_cos(self, positions):
         cos_sin = self.rotary_emb.cos_sin_cache[positions]
         cos, sin = cos_sin.chunk(2, dim=-1)
@@ -590,6 +594,15 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
             dequant_q_norm,
         )
 
+    def uses_mlaprolog(self) -> bool:
+        _is_arch35_dsa = (
+            self.is_npu_arch35 and get_token_to_kv_pool().index_head_dim is not None
+        )
+        return _is_arch35_dsa or (
+            hasattr(self.quant_config, "ignore")
+            and any(re.fullmatch(r".*kv_b_proj", l) for l in self.quant_config.ignore)
+        )
+
     def forward(self, positions, hidden_states, forward_batch, zero_allocator):
         # assert self.quant_config and self.quant_config.get_name() == "modelslim"
         # route by `qkv_a_proj` quant type as MTP layers can be unquantized
@@ -602,10 +615,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
             self.is_npu_arch35 and get_token_to_kv_pool().index_head_dim is not None
         )
         # with the mlaprolog enabled, the kv_b_proj layers are unquantized
-        _is_mlaprolog = _is_arch35_dsa or (
+        _is_mlaprolog = self.uses_mlaprolog()
-            hasattr(self.quant_config, "ignore")
-            and any(re.fullmatch(r".*kv_b_proj", l) for l in self.quant_config.ignore)
-        )
         if _is_w8a8 and not _is_arch35_dsa:
             return self.forward_mlapo(
                 positions, hidden_states, forward_batch, zero_allocator

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -1,4 +1,3 @@
-import re
 from typing import TYPE_CHECKING, Optional
 
 import torch
@@ -18,6 +17,7 @@ from sglang.srt.layers.attention.dsa.utils import (
 )
 from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.runtime_context import get_disagg
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -598,14 +598,14 @@ def npu_mla_preprocess(
             m.v_head_dim,
             m.quant_config,
         )
+        if (
+            get_disagg().disaggregation_mode == "decode"
+            and m.mla_preprocess.uses_mlaprolog()
+            and m.w_kc is not None
+        ):
+            m.w_kc.untyped_storage().resize_(0)
     # mlaprolog does not require additional calculation of q_lora
-    _is_mlaprolog = (
+    if m.mla_preprocess.uses_mlaprolog():
-        _is_npu_arch35 and get_token_to_kv_pool().index_head_dim is not None
-    ) or (
-        hasattr(m.quant_config, "ignore")
-        and any(re.fullmatch(r".*kv_b_proj", l) for l in m.quant_config.ignore)
-    )
-    if _is_mlaprolog:
         (
             q_pe,
             k_pe,