Capturing aux hidden states (Eagle3/DFlash) has a transient ~2x HBM spike: the model collects K separate [tokens, hidden] tensors and LogitsProcessor concatenates them into [tokens, K * hidden], holding both at once.
值得精读：AuxHiddenStatePacker 的设计模式可推广到其他需要累积张量并减少内存峰值的场景。关注 copies_on_append 属性的使用和兼容性机制。
未产生实质性代码审查讨论；合并者 hnyls2002 对原始提交进行了重构（提取独立模块、精简注释、增加断言），并通过 CI rerun 验证了测试通过。
aux_hidden_states.py 定义 AuxHiddenStatePacker 类和 AuxHiddenStates 等类型别名。Packer 通过 append 将每层捕获写入预分配 [tokens, K * hidden] 缓冲，finalize 返回 packed tensor。communicator.py 中的 prepare_attn_and_capture_last_layer_outputs 方法，参数类型从 Optional[List[torch.Tensor]] 改为 Optional[AuxHiddenStateAccumulator]，并利用 copies_on_append 属性避免不必要的克隆。logits_processor.py，导入 AuxHiddenStates 和 pack_aux_hidden_states，将 aux_hidden_states 参数类型改为 Optional[AuxHiddenStates]，并在 _get_pruned_states 中处理 packed tensor 和 list 两种形式。deepseek_v2.py，导入并使用 AuxHiddenStatePacker 替换列表，初始化时传入层数，最后调用 finalize 返回 packed buffer。| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
python/sglang/srt/layers/aux_hidden_states.py |
捕获管理 | added | 8.39 |
python/sglang/srt/layers/logits_processor.py |
Logits 处理 | modified | 6.67 |
python/sglang/srt/models/deepseek_v2.py |
DeepSeek V2 模型 | modified | 6.15 |
python/sglang/srt/layers/communicator.py |
通信层 | modified | 5.4 |
python/sglang/srt/layers/aux_hidden_states.py
core-logic
新增核心文件，定义了 AuxHiddenStatePacker 类和类型别名，是本次性能优化的核心实现。
"""Aux hidden states captured for Eagle3/DFlash draft models."""
from typing import List, Optional, Union
import torch
# 两种表示共存：已迁移到 AuxHiddenStatePacker 的模型传递一个 packed [tokens, K * hidden] tensor，
# 其余模型仍然传递一个 K 个 tensor 的列表。
AuxHiddenStates = Union[torch.Tensor, List[torch.Tensor]]
class AuxHiddenStatePacker:
"""替换模型收集 Eagle3/DFlash 捕获时使用的 `[]` 列表。
每个 `.append()` 直接写入预分配的 `[tokens, K * hidden]` 缓冲区，
避免列表路径在 `torch.cat` 时产生的瞬时 ~2x HBM。
假设所有捕获共享 leading shape 和 feature size。
"""
# `append` 会拷贝，因此生产者无需为后续 mutate 而 clone。
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:
# 第一次 append 时根据 hidden 的形状和总捕获次数分配完整缓冲区
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:
"""返回 packed 缓冲区；调用者应通过 `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
# 模型向下游捕获路径传递的类型：普通列表或就地写入的 packer
AuxHiddenStateAccumulator = Union[List[torch.Tensor], AuxHiddenStatePacker]
def pack_aux_hidden_states(aux_hidden_states: AuxHiddenStates) -> torch.Tensor:
"""统一将两种形式转换为 packed tensor（若已是 tensor 则直接返回）。"""
if isinstance(aux_hidden_states, torch.Tensor):
return aux_hidden_states
return torch.cat(aux_hidden_states, dim=-1)
通过 AuxHiddenStates Union 类型（torch.Tensor 或 List[torch.Tensor]）同时支持新 packed tensor 和旧 list 路径，降低迁移风险。
结论：采纳，未迁移模型无感知。 · 已解决
num_captures 与实际 append 次数不一致，finalize 断言会抛出异常，开发阶段可及早暴露错误。copies_on_append 属性可能被误用：若 accumulator 不拷贝，调用者必须保证后续不会修改已追加的张量。对 DeepSeek-V2 用户：减少推理过程中的峰值显存，降低 OOM 风险，尤其长序列场景。接口向后兼容，无需修改配置。对系统：捕获路径的代码可读性略有提升（类型标注更精确）。对团队：为 Eagle3/DFlash 的进一步优化奠定基础。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论