diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 318ee9703..36cdb23c3 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -498,11 +498,13 @@ class LayerCommunicator: forward_batch: ForwardBatch, captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, post_residual_addition: Optional[torch.Tensor] = None, + quant_format: str = "", ): hidden_states, residual = self.prepare_attn( hidden_states, residual, forward_batch, + quant_format=quant_format, post_residual_addition=post_residual_addition, ) if captured_last_layer_outputs is not None: @@ -511,12 +513,39 @@ class LayerCommunicator: forward_batch=forward_batch, context=self._context, ) - if gathered_last_layer_output is residual: - # Clone to avoid modifying the original residual by Custom RMSNorm inplace operation + if ( + gathered_last_layer_output is residual + and not self._post_attn_residual_is_read_only(residual) + ): gathered_last_layer_output = residual.clone() captured_last_layer_outputs.append(gathered_last_layer_output) return hidden_states, residual + def _post_attn_residual_is_read_only(self, residual: torch.Tensor) -> bool: + """True if ``prepare_mlp``'s post-attention RMSNorm leaves ``residual`` + untouched, so Eagle3 aux capture can keep its reference and skip the clone. + + Only the flashinfer all-reduce-fusion path writes a fresh ``residual_out`` + (see ``flashinfer_allreduce_residual_rmsnorm``); the aiter fused kernel and + every plain norm fold into ``residual`` in place. That path is reachable + only from the ``_gather_*`` communicate-fns, and only when they fall past + their input-scattered branch. + """ + norm_fn = getattr( + self._communicate_with_all_reduce_and_layer_norm_fn, + "func", + self._communicate_with_all_reduce_and_layer_norm_fn, + ) + uses_gather_norm = norm_fn in ( + CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, + CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual_moe, + ) + return ( + uses_gather_norm + and not get_attn_tp_context().input_scattered + and apply_flashinfer_allreduce_fusion(residual.shape[0]) + ) + def prepare_attn( self, hidden_states: torch.Tensor, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c6208ff1d..fa60ae4ff 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -75,7 +75,6 @@ from sglang.srt.layers.communicator_dsa_cp import DSACPLayerCommunicator from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, - get_attention_tp_group, get_attention_tp_rank, get_attention_tp_size, ) @@ -2130,13 +2129,17 @@ class DeepseekV2DecoderLayer(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, llama_4_scaling: Optional[torch.Tensor] = None, prev_topk_indices: Optional[torch.Tensor] = None, + captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, ) -> torch.Tensor: hidden_states_orig = hidden_states - hidden_states, residual = self.layer_communicator.prepare_attn( - hidden_states, - residual, - forward_batch, - getattr(self, "_gfx95_quant_format", ""), + hidden_states, residual = ( + self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( + hidden_states, + residual, + forward_batch, + captured_last_layer_outputs=captured_last_layer_outputs, + quant_format=getattr(self, "_gfx95_quant_format", ""), + ) ) hidden_states = self.self_attn( @@ -2483,14 +2486,6 @@ class DeepseekV2Model(nn.Module): else get_global_expert_distribution_recorder().with_current_layer(i) ) with ctx: - if i in self.layers_to_capture: - if self.enable_a2a_moe and i > self.first_k_dense_replace: - aux_hidden_state = get_attention_tp_group().all_gather( - hidden_states + residual, dim=0 - ) - aux_hidden_states.append(aux_hidden_state) - else: - aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] hidden_states, residual, topk_indices = layer( positions, @@ -2501,6 +2496,9 @@ class DeepseekV2Model(nn.Module): gemm_output_zero_allocator, llama_4_scaling, prev_topk_indices=topk_indices, + captured_last_layer_outputs=( + aux_hidden_states if i in self.layers_to_capture else None + ), ) if normal_end_layer != self.end_layer: