diff --git a/python/sglang/srt/layers/aux_hidden_states.py b/python/sglang/srt/layers/aux_hidden_states.py new file mode 100644 index 000000000..ac4aff91b --- /dev/null +++ b/python/sglang/srt/layers/aux_hidden_states.py @@ -0,0 +1,58 @@ +"""Aux hidden states captured for Eagle3/DFlash draft models.""" + +from typing import List, Optional, Union + +import torch + +# Two representations coexist: models migrated to AuxHiddenStatePacker pass one +# packed [tokens, K * hidden] tensor, the rest still pass a list of K tensors. +AuxHiddenStates = Union[torch.Tensor, List[torch.Tensor]] + + +class AuxHiddenStatePacker: + """Drop-in for the ``[]`` a model collects Eagle3/DFlash captures into. + + Each ``.append()`` writes into one preallocated ``[tokens, K * hidden]`` + buffer, avoiding the list path's transient ~2x HBM at ``torch.cat``. + Assumes all captures share leading shape and feature size. + """ + + # ``append`` copies, so producers need not clone a tensor they later mutate. + copies_on_append = True + + def __init__(self, num_captures: int) -> None: + self._num_captures = int(num_captures) + self._buffer: Optional[torch.Tensor] = None + self._feature_size: Optional[int] = None + self._idx = 0 + + def append(self, hidden: torch.Tensor) -> None: + feature_size = int(hidden.shape[-1]) + if self._buffer is None: + self._feature_size = feature_size + self._buffer = hidden.new_empty( + (*hidden.shape[:-1], feature_size * self._num_captures) + ) + start = self._idx * self._feature_size + self._buffer[..., start : start + self._feature_size].copy_(hidden) + self._idx += 1 + + def __len__(self) -> int: + return self._idx + + def finalize(self) -> torch.Tensor: + """Return the packed buffer; callers guard the empty case on ``len()``.""" + assert ( + self._buffer is not None and self._idx == self._num_captures + ), f"captured {self._idx} of {self._num_captures} aux hidden states" + return self._buffer + + +# What a model hands down the capture path: a plain list, or a packer writing in place. +AuxHiddenStateAccumulator = Union[List[torch.Tensor], AuxHiddenStatePacker] + + +def pack_aux_hidden_states(aux_hidden_states: AuxHiddenStates) -> torch.Tensor: + if isinstance(aux_hidden_states, torch.Tensor): + return aux_hidden_states + return torch.cat(aux_hidden_states, dim=-1) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index f2b810286..92ab37d80 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -16,7 +16,7 @@ from contextlib import contextmanager from dataclasses import dataclass from enum import Enum, auto from functools import partial -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable, Dict, Optional, Tuple, Union import torch @@ -35,6 +35,7 @@ from sglang.srt.layers.attention.dsa.utils import ( dsa_use_prefill_cp, is_dsa_enable_prefill_cp, ) +from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator from sglang.srt.layers.dp_attention import ( attn_tp_all_gather_into_tensor, attn_tp_reduce_scatter_tensor, @@ -506,7 +507,7 @@ class LayerCommunicator: hidden_states: torch.Tensor, residual: torch.Tensor, forward_batch: ForwardBatch, - captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, + captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None, post_residual_addition: Optional[torch.Tensor] = None, quant_format: str = "", ): @@ -525,6 +526,8 @@ class LayerCommunicator: ) if ( gathered_last_layer_output is residual + # An accumulator that copies on append already holds a snapshot. + and not getattr(captured_last_layer_outputs, "copies_on_append", False) and not self._post_attn_residual_is_read_only(residual) ): gathered_last_layer_output = residual.clone() diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index cf55b97cf..641644c7d 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -25,6 +25,10 @@ from sglang.kernels.ops.activation.softcap import ( softcap_inplace_logits as fused_softcap, ) from sglang.srt.distributed.device_communicators import triton_symm_mem_ag +from sglang.srt.layers.aux_hidden_states import ( + AuxHiddenStates, + pack_aux_hidden_states, +) from sglang.srt.layers.dp_attention import ( DpPaddingMode, attn_tp_all_gather, @@ -389,7 +393,7 @@ class LogitsProcessor(nn.Module): hidden_states, lm_head: VocabParallelEmbedding, logits_metadata: Union[LogitsMetadata, ForwardBatch], - aux_hidden_states: Optional[torch.Tensor] = None, + aux_hidden_states: Optional[AuxHiddenStates] = None, hidden_states_before_norm: Optional[torch.Tensor] = None, ) -> LogitsProcessorOutput: # Extract MIS indices before ForwardBatch → LogitsMetadata conversion @@ -485,7 +489,7 @@ class LogitsProcessor(nn.Module): self, hidden_states: torch.Tensor, hidden_states_before_norm: Optional[torch.Tensor], - aux_hidden_states: Optional[torch.Tensor], + aux_hidden_states: Optional[AuxHiddenStates], logits_metadata: LogitsMetadata, ): pruned_states_before_norm: Optional[torch.Tensor] = None @@ -505,7 +509,11 @@ class LogitsProcessor(nn.Module): pruned_states = hidden_states pruned_states_before_norm = hidden_states_before_norm if aux_hidden_states is not None: - aux_pruned_states = [hidden for hidden in aux_hidden_states] + aux_pruned_states = ( + aux_hidden_states + if isinstance(aux_hidden_states, torch.Tensor) + else [hidden for hidden in aux_hidden_states] + ) sample_indices = None input_logprob_indices = None @@ -519,7 +527,11 @@ class LogitsProcessor(nn.Module): if hidden_states_before_norm is not None: pruned_states_before_norm = hidden_states_before_norm[last_index] if aux_hidden_states is not None: - aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states] + aux_pruned_states = ( + aux_hidden_states[last_index] + if isinstance(aux_hidden_states, torch.Tensor) + else [hidden[last_index] for hidden in aux_hidden_states] + ) sample_indices = None input_logprob_indices = None else: @@ -551,11 +563,14 @@ class LogitsProcessor(nn.Module): input_logprob_indices_pt = 0 input_logprob_indices = [] pt, pruned_states_list, pruned_states_before_norm_list = 0, [], [] - aux_pruned_states_lists = ( - [[] for _ in aux_hidden_states] - if aux_hidden_states is not None - else None - ) + is_packed_aux_hidden_states = isinstance(aux_hidden_states, torch.Tensor) + aux_pruned_states_lists = None + if aux_hidden_states is not None: + aux_pruned_states_lists = ( + [] + if is_packed_aux_hidden_states + else [[] for _ in aux_hidden_states] + ) for idx, (extend_logprob_start_len, extend_len) in enumerate( zip( @@ -581,10 +596,15 @@ class LogitsProcessor(nn.Module): hidden_states_before_norm[pt + start_len : pt + extend_len] ) if aux_pruned_states_lists is not None: - for j, hidden in enumerate(aux_hidden_states): - aux_pruned_states_lists[j].append( - hidden[pt + start_len : pt + extend_len] + if is_packed_aux_hidden_states: + aux_pruned_states_lists.append( + aux_hidden_states[pt + start_len : pt + extend_len] ) + else: + for j, hidden in enumerate(aux_hidden_states): + aux_pruned_states_lists[j].append( + hidden[pt + start_len : pt + extend_len] + ) # Map each token to its sequence index, for chunked computation # of input logprobs token_to_seq_idx.extend([idx] * (extend_len - start_len)) @@ -603,7 +623,11 @@ class LogitsProcessor(nn.Module): if hidden_states_before_norm is not None: pruned_states_before_norm = torch.cat(pruned_states_before_norm_list) if aux_pruned_states_lists is not None: - aux_pruned_states = [torch.cat(lst) for lst in aux_pruned_states_lists] + aux_pruned_states = ( + torch.cat(aux_pruned_states_lists) + if is_packed_aux_hidden_states + else [torch.cat(lst) for lst in aux_pruned_states_lists] + ) # Build the index tensors via pinned host memory + non-blocking H2D # so the small copy doesn't drain the stream. @@ -631,10 +655,10 @@ class LogitsProcessor(nn.Module): self, hidden_states: torch.Tensor, hidden_states_before_norm: Optional[torch.Tensor], - aux_hidden_states: Optional[List[torch.Tensor]], + aux_hidden_states: Optional[AuxHiddenStates], pruned_states: torch.Tensor, pruned_states_before_norm: Optional[torch.Tensor], - aux_pruned_states: Optional[List[torch.Tensor]], + aux_pruned_states: Optional[AuxHiddenStates], sample_indices: Optional[torch.Tensor], logits_metadata: LogitsMetadata, ) -> Optional[torch.Tensor]: @@ -643,8 +667,7 @@ class LogitsProcessor(nn.Module): if logits_metadata.capture_hidden_mode.need_capture(): if logits_metadata.capture_hidden_mode.is_full(): if aux_hidden_states is not None: - aux_hidden_states = torch.cat(aux_hidden_states, dim=-1) - hidden_states_to_store = aux_hidden_states + hidden_states_to_store = pack_aux_hidden_states(aux_hidden_states) else: hidden_states_to_store = hidden_states hidden_states_to_store_before_norm = hidden_states_before_norm @@ -652,7 +675,8 @@ class LogitsProcessor(nn.Module): # Get the last token hidden states. If sample_indices is None, # pruned states only contain the last tokens already. if aux_hidden_states is not None: - aux_pruned_states = torch.cat(aux_pruned_states, dim=-1) + assert aux_pruned_states is not None + aux_pruned_states = pack_aux_hidden_states(aux_pruned_states) hidden_states_to_store = ( aux_pruned_states[sample_indices] if sample_indices is not None diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index bc177ef24..98fc64785 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -67,6 +67,10 @@ from sglang.srt.layers.attention.dsa.utils import ( dsa_use_prefill_cp, is_dsa_enable_prefill_cp, ) +from sglang.srt.layers.aux_hidden_states import ( + AuxHiddenStateAccumulator, + AuxHiddenStatePacker, +) from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, @@ -2234,7 +2238,7 @@ 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, + captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None, next_full_attention_layer_id: Optional[int] = None, ) -> torch.Tensor: hidden_states_orig = hidden_states @@ -2618,7 +2622,8 @@ class DeepseekV2Model(nn.Module): normal_end_layer = self.first_k_dense_replace elif self.first_k_dense_replace < normal_start_layer: normal_end_layer = normal_start_layer = 0 - aux_hidden_states = [] + # Append-compatible, so the shared capture path below is unchanged. + aux_hidden_states = AuxHiddenStatePacker(len(self.layers_to_capture)) if self.pp_group.is_first_rank: topk_indices = None for i in range(normal_start_layer, normal_end_layer): @@ -2704,7 +2709,7 @@ class DeepseekV2Model(nn.Module): ) if len(aux_hidden_states) == 0: return hidden_states - return hidden_states, aux_hidden_states + return hidden_states, aux_hidden_states.finalize() class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):