diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0bb6149b4..e0cb3f46f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -312,6 +312,8 @@ class Envs: SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False) # Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2 SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM = EnvBool(False) + # Delay all-gather after qlora for better performance for Deepseek v3.2 + SGLANG_USE_AG_AFTER_QLORA = EnvBool(False) # Quantization SGLANG_INT4_WEIGHT = EnvBool(False) 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 cc50cb345..0fd8d9788 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 @@ -4,21 +4,24 @@ from typing import TYPE_CHECKING import torch import torch_npu +from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( NPUFusedMLAPreprocess, is_fia_nz, is_mla_preprocess_enabled, ) +from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full from sglang.srt.layers.attention.nsa.utils import ( cp_split_and_rebuild_position, nsa_use_prefill_cp, ) -from sglang.srt.layers.communicator import get_attn_tp_context +from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.utils import BumpAllocator +_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() # region MHA @@ -28,6 +31,7 @@ def forward_mha_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", + layer_scatter_modes, ): if m.q_lora_rank is not None: q, latent_cache = ( @@ -55,6 +59,13 @@ def forward_mha_prepare_npu( else: q = m.q_a_layernorm(q) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + q = scattered_to_tp_attn_full(q, forward_batch) + latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch) q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) else: @@ -142,6 +153,7 @@ def forward_mla_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", + layer_scatter_modes, ): if is_mla_preprocess_enabled(): if not hasattr(m, "mla_preprocess"): @@ -184,6 +196,13 @@ def forward_mla_prepare_npu( k_nope = latent_cache[..., : m.kv_lora_rank] q = m.q_a_layernorm(q) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + q = scattered_to_tp_attn_full(q, forward_batch) + latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch) k_nope = m.kv_a_layernorm(k_nope) # q_lora needed by indexer @@ -285,6 +304,7 @@ def forward_dsa_prepare_npu( hidden_states: torch.Tensor, forward_batch: "ForwardBatch", zero_allocator: "BumpAllocator", + layer_scatter_modes, ): dynamic_scale = None if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode(): @@ -313,7 +333,13 @@ def forward_dsa_prepare_npu( # overlap qk norm q = m.q_a_layernorm(q) - + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + q = scattered_to_tp_attn_full(q, forward_batch) + latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch) q_lora = q.clone() # required for topk_indices q_event = None @@ -353,7 +379,13 @@ def forward_dsa_prepare_npu( ) topk_indices = m.indexer( - hidden_states, q_lora, positions, forward_batch, m.layer_id, dynamic_scale + hidden_states, + q_lora, + positions, + forward_batch, + m.layer_id, + layer_scatter_modes, + dynamic_scale, ) return ( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 323f77c67..f2e659332 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -12,6 +12,7 @@ from sglang.jit_kernel.fused_store_index_cache import ( fused_store_index_k_cache, ) from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.layers.utils import MultiPlatformOp @@ -43,6 +44,7 @@ from sglang.srt.layers.attention.nsa.utils import ( is_nsa_enable_prefill_cp, is_nsa_prefill_cp_in_seq_split, ) +from sglang.srt.layers.communicator import ScatterMode from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.rotary_embedding import get_rope_wrapper @@ -50,6 +52,7 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import get_global_server_args +_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool @@ -1203,6 +1206,7 @@ class Indexer(MultiPlatformOp): positions: torch.Tensor, forward_batch: ForwardBatch, layer_id: int, + layer_scatter_modes=None, dynamic_scale: torch.Tensor = None, ) -> torch.Tensor: if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None: @@ -1223,7 +1227,7 @@ class Indexer(MultiPlatformOp): cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) - bs = x.shape[0] + bs = q_lora.shape[0] if self.alt_stream is not None: self.alt_stream.wait_stream(torch.npu.current_stream()) with torch.npu.stream(self.alt_stream): @@ -1276,6 +1280,12 @@ class Indexer(MultiPlatformOp): k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] k = self.k_norm(k_proj) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + k = scattered_to_tp_attn_full(k, forward_batch) k_pe, k_nope = torch.split( k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], @@ -1356,7 +1366,12 @@ class Indexer(MultiPlatformOp): torch.npu.current_stream().wait_event(q_rope_event) if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): torch.npu.current_stream().wait_event(weights_event) - + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + weights = scattered_to_tp_attn_full(weights, forward_batch) block_table = forward_batch.attn_backend.forward_metadata.block_tables if ( is_prefill @@ -1450,3 +1465,19 @@ class Indexer(MultiPlatformOp): sparse_mode=3, ) return topk_indices_prev[0], topk_indices_next[0] + + +def scattered_to_tp_attn_full( + hidden_states: torch.Tensor, + forward_batch, +) -> torch.Tensor: + hidden_states, local_hidden_states = ( + torch.empty( + (forward_batch.input_ids.shape[0], hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ), + hidden_states, + ) + attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous()) + return hidden_states diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 57c18c413..03e5516c2 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -29,6 +29,7 @@ from sglang.srt.distributed import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) +from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa.utils import ( is_nsa_enable_prefill_cp, nsa_use_prefill_cp, @@ -74,6 +75,7 @@ _is_sm100_supported = _is_cuda and is_sm100_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() _is_gfx95_supported = is_gfx95_supported() _is_npu = is_npu() +_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() if _use_aiter and _is_gfx95_supported: from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant @@ -189,7 +191,7 @@ class AttnTpContext: def init_context(self, q_lora_rank, is_nsa): self.allow_input_scattered = ( get_global_server_args().enable_attn_tp_input_scattered - and _is_cuda + and (_is_cuda or _is_npu) and q_lora_rank is not None and not is_nsa and get_tensor_model_parallel_world_size() > 1 @@ -693,6 +695,8 @@ class CommunicateSimpleFn: if (input_mode == ScatterMode.SCATTERED) and ( output_mode == ScatterMode.TP_ATTN_FULL ): + if _use_ag_after_qlora: + return CommunicateSimpleFn._trivial return CommunicateSimpleFn._scattered_to_tp_attn_full raise NotImplementedError(f"{input_mode=} {output_mode=}") diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index eb36952ca..04e97d952 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1301,6 +1301,7 @@ class DeepseekV2AttentionMLA( hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, + layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, ): s = self.forward_prepare( @@ -1308,6 +1309,7 @@ class DeepseekV2AttentionMLA( hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, + layer_scatter_modes=layer_scatter_modes, llama_4_scaling=llama_4_scaling, ) return self.forward_core(s) @@ -1318,6 +1320,7 @@ class DeepseekV2AttentionMLA( hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, + layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, ): if self.attn_mha.kv_b_proj is None: @@ -1370,15 +1373,30 @@ class DeepseekV2AttentionMLA( ) elif attn_forward_method == AttnForwardMethod.MHA_NPU: inner_state = forward_mha_prepare_npu( - self, positions, hidden_states, forward_batch, zero_allocator + self, + positions, + hidden_states, + forward_batch, + zero_allocator, + layer_scatter_modes, ) elif attn_forward_method == AttnForwardMethod.MLA_NPU: inner_state = forward_mla_prepare_npu( - self, positions, hidden_states, forward_batch, zero_allocator + self, + positions, + hidden_states, + forward_batch, + zero_allocator, + layer_scatter_modes, ) elif attn_forward_method == AttnForwardMethod.DSA_NPU: inner_state = forward_dsa_prepare_npu( - self, positions, hidden_states, forward_batch, zero_allocator + self, + positions, + hidden_states, + forward_batch, + zero_allocator, + layer_scatter_modes, ) else: raise NotImplementedError @@ -1505,6 +1523,10 @@ class DeepseekV2DecoderLayer(nn.Module): prefix=add_prefix("self_attn", prefix), alt_stream=alt_stream, ) + if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get(): + raise ValueError( + "SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank" + ) self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False) @@ -1627,6 +1649,7 @@ class DeepseekV2DecoderLayer(nn.Module): forward_batch=forward_batch, zero_allocator=zero_allocator, llama_4_scaling=llama_4_scaling, + layer_scatter_modes=self.layer_scatter_modes, ) hidden_states, residual = self.layer_communicator.prepare_mlp(