Pack aux hidden states into a preallocated buffer (#28956)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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)

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()

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,10 +563,13 @@ class LogitsProcessor(nn.Module):
             input_logprob_indices_pt = 0
             input_logprob_indices = []
             pt, pruned_states_list, pruned_states_before_norm_list = 0, [], []
+            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 = (
-                [[] for _ in aux_hidden_states]
+                    []
-                if aux_hidden_states is not None
+                    if is_packed_aux_hidden_states
-                else None
+                    else [[] for _ in aux_hidden_states]
                 )
 
             for idx, (extend_logprob_start_len, extend_len) in enumerate(
@@ -581,6 +596,11 @@ class LogitsProcessor(nn.Module):
                         hidden_states_before_norm[pt + start_len : pt + extend_len]
                     )
                 if aux_pruned_states_lists is not None:
+                    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]
@@ -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 = pack_aux_hidden_states(aux_hidden_states)
-                    hidden_states_to_store = 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

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):