diff --git a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py index 47a863117..9d0f5fc94 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py +++ b/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 ( - hasattr(self.quant_config, "ignore") - and any(re.fullmatch(r".*kv_b_proj", l) for l in self.quant_config.ignore) - ) + _is_mlaprolog = self.uses_mlaprolog() if _is_w8a8 and not _is_arch35_dsa: return self.forward_mlapo( positions, hidden_states, forward_batch, zero_allocator diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index addcb91e5..375550947 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/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 = ( - _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: + if m.mla_preprocess.uses_mlaprolog(): ( q_pe, k_pe,