执行摘要
- 一句话：LayerNorm 序列并行：prefill 激活内存与延迟双降
- 推荐动作：值得精读，尤其三个设计决策：(1) CUDA graph 下捕获区内写入的 Python 标志在回放时过期这一问题的识别与
runs_sp 规避方案；(2) LayerCommunicator 构造全 SCATTERED sibling 的委托模式，避免在每个模型里打点；(3) fused matmul+collective 的可用性探测与普通 collective fallback 策略，兼顾性能与跨 torch 版本兼容性。对后续做任何序列切分、通信融合类特性都有直接参考价值。
功能与动机
长上下文 prefill 的激活显存与延迟随序列长度线性增长。PR body 指出，纯 TP 下 row-parallel 的 all_reduce 代数上等价于 reduce_scatter + all_gather，两种形式移动相同字节数——拆开后 LayerNorm/residual 只需在 1/tp 的 token 上计算，既降低激活内存又无额外通信开销。设计目标明确为 off by default and byte-identical when off，避免影响现有用户；实测 2×H100 上 GSM8K 0.904 vs 0.906（fp 重排噪声），prefill 0.28173s → 0.27180s（-3.5%）；2×B200 上 0.14388s → 0.13589s（-5.6%）；早期 8×H100（Qwen3-32B、32K）版本为 -16.5% 延迟、-27.7% 激活内存。
实现拆解
-
新增 SP 核心模块 python/sglang/srt/layers/layernorm_sp.py（+245 行）：集中承载白名单 SP_SUPPORTED_ARCHITECTURES（当前仅 Qwen3ForCausalLM）、initialize_layernorm_sp（分布式初始化后按 worker 物化 flags.sp.enabled）、边界 collective sp_entry_scatter / sp_exit_gather、出口判定 maybe_exit_gather，以及基于 torch symmetric memory 的 fused matmul+collective 快路径 column_parallel_g_matmul / row_parallel_gbar_matmul（带普通 collective 回退）。关键设计：runs_sp(forward_mode) 只依赖配置与 forward 模式，专供 CUDA graph 捕获区域外使用；_SPForwardState.num_tokens 使用实例属性而非 ForwardFlags 整数槽，避免 torch.compile 按序列长度重新编译。
-
接入前向执行链路：layers/communicator.py 的 LayerCommunicator 在 SP 启用时构造一个全 SCATTERED 的 _sp_variant sibling，并在 prepare_attn 中于首个 decoder layer（residual is None）判定 EXTEND 写入 sp_active、执行入口 scatter，随后 prepare_attn / prepare_mlp / postprocess_layer 全部委托给 sibling，使 norm/residual 在本地 shard 上计算、不触发额外 collective；layers/linear.py 的 ColumnParallelLinear 与 RowParallelLinear 在 sp_active 且 tp_size > 1 时分别进入 g（all-gather+matmul）与 g-bar（matmul+reduce-scatter）分支；layers/logits_processor.py 在 LM head 之前调用 maybe_exit_gather 撤销切分——判定用 runs_sp 而非 sp_active，因为捕获区域内写入的标志在 CUDA graph 回放时不会重执行。
-
运行时状态与配置入口：runtime_context.py 新增 SpFlags（挂到 Flags 根下）与 ForwardFlags.sp_active 槽，并把它加入 _GRAPH_VISIBLE 保证图捕获期间线性层可读；server_args.py 新增 --enable-layernorm-sp（parallel 命名空间）；distributed/bootstrap.py 在 initialize_dp_attention 旁边调用 initialize_layernorm_sp。
-
配置校验：新增 arg_groups/layernorm_sp_hook.py 并注册进解析管道（pipeline.py），validate_layernorm_sp 依次拒绝：白名单外架构、tp_size <= 1、--enable-dp-attention、投机解码；另有静默门控：仅 EXTEND 生效；fused 快路径要求未量化、无 bias 的 bf16/fp16 线性层，否则回退普通 collective。
-
测试配套：CPU 单测 test/registered/unit/layers/test_layernorm_sp.py（白名单、prefill-only、四项校验、runs_sp 忽略 sp_active 的回归、sp_active 槽注册检查）；e2e test/registered/models_e2e/test_layernorm_sp.py 以 Qwen3-8B + --tp 2 --enable-layernorm-sp 跑 GSM8K（阈值 0.85），分别注册为 CPU 与 2-gpu-large CI 阶段。
关键文件：
python/sglang/srt/layers/layernorm_sp.py（模块 序列并行；类别 source；类型 core-logic；符号 initialize_layernorm_sp, layernorm_sp_enabled, runs_sp, _SPForwardState）: 新增 SP 核心模块：白名单、启用判定、entry/exit collective、fused matmul 快路径全部集中于此，是 models 零侵入解耦的关键。python/sglang/srt/arg_groups/layernorm_sp_hook.py（模块 配置校验；类别 source；类型 configuration；符号 handle_layernorm_sp, validate_layernorm_sp）: 启动期校验 hook：在解析管道中拒绝不支持的架构、tp_size<=1、dp-attention 与投机解码，保证不兼容配置 fail loud。python/sglang/srt/layers/communicator.py（模块 层通信；类别 source；类型 core-logic）: LayerCommunicator 新增 _sp_variant 全 SCATTERED sibling 并委托 prepare_attn / prepare_mlp / postprocess_layer，是 SP 区域在解码器内运转的枢纽。python/sglang/srt/layers/linear.py（模块 线性层；类别 source；类型 dependency-wiring）: Row / ColumnParallelLinear 在 sp_active 时自动参与 g / g-bar，无需在模型代码里逐个标记，是零侵入解耦的关键。python/sglang/srt/runtime_context.py（模块 运行时状态；类别 source；类型 core-logic；符号 SpFlags）: 新增 SpFlags 与 ForwardFlags.sp_active 槽（含 _GRAPH_VISIBLE 注册），是 SP 运行期状态的基础。python/sglang/srt/layers/logits_processor.py（模块 输出头；类别 source；类型 dependency-wiring）: LM head 前的 exit gather 接入点，调用 maybe_exit_gather 撤销序列切分，模型无关。python/sglang/srt/server_args.py（模块 参数配置；类别 source；类型 configuration）: 新增 --enable-layernorm-sp 用户入口参数（parallel 命名空间），是功能对外开关。python/sglang/srt/distributed/bootstrap.py（模块 分布式启动；类别 source；类型 entrypoint）: 分布式初始化后调用 initialize_layernorm_sp，物化 flags.sp.enabled，与 initialize_dp_attention 并列。python/sglang/srt/arg_groups/pipeline.py（模块 解析管线；类别 source；类型 entrypoint）: 把 layernorm_sp_hook 注册进解析管道，使校验成为 resolution pipeline 的一部分。test/registered/unit/layers/test_layernorm_sp.py（模块 单元测试；类别 test；类型 test-coverage；符号 _initialize, TestLayerNormSPGating, tearDown, test_initialize_enables_only_for_allowlisted_arch）: CPU 单测覆盖白名单、prefill-only、四项校验与 CUDA graph 相关回归，锁死 runs_sp 语义。test/registered/models_e2e/test_layernorm_sp.py（模块 端到端测试；类别 test；类型 test-coverage；符号 TestLayerNormSPAccuracy, setUpClass, tearDownClass）: e2e 精度测试：Qwen3-8B tp=2 GSM8K 阈值 0.85，验证 SP 不回归模型输出。
关键符号：initialize_layernorm_sp, layernorm_sp_enabled, runs_sp, set_sp_num_tokens, sp_num_tokens, sp_entry_scatter, sp_exit_gather, maybe_exit_gather, sp_fused_matmul_eligible, column_parallel_g_matmul, row_parallel_gbar_matmul, handle_layernorm_sp, validate_layernorm_sp
关键源码片段
python/sglang/srt/layers/layernorm_sp.py
新增 SP 核心模块：白名单、启用判定、entry/exit collective、fused matmul 快路径全部集中于此，是 models 零侵入解耦的关键。
# Megatron 风格 LayerNorm 序列并行（arXiv:2205.05198）的边界与门控核心。
# SP 只对白名单架构生效；机制是通用的，后续模型只需扩展这一行并做硬件验证。
SP_SUPPORTED_ARCHITECTURES = frozenset({'Qwen3ForCausalLM'})
def runs_sp(forward_mode) -> bool:
'''判断本次 forward 是否运行 SP：模型启用且仅限 prefill（EXTEND）。
CUDA graph 捕获区域之外的代码必须用它，不能读 ``sp_active``——
捕获区域内 Python 写入在 graph replay 时不会重执行，flag 是过期的。
'''
from sglang.srt.model_executor.forward_batch_info import ForwardMode
return layernorm_sp_enabled() and forward_mode == ForwardMode.EXTEND
class _SPForwardState:
'''保存当前 SP forward 的真实（未 padding）token 数。
特意用实例属性而非 ForwardFlags 整数槽：该值在 torch.compile 追踪的
线性层代码内被读取，dynamo 对属性来源 int 按 automatic-dynamic 处理，
而 dict 槽 int 会随序列长度触发重新编译。
'''
num_tokens: int = 0
def sp_entry_scatter(hidden_states: torch.Tensor) -> torch.Tensor:
'''把复制的 [M, h] hidden states 沿 token 维切分。
M 会 pad 到 tp_size 整数倍，padding 行由出口侧丢弃；输入在 TP group 内
各 rank 持有完整副本，所以这里只是本地切片，无需通信。
'''
num_tokens = hidden_states.shape[0]
set_sp_num_tokens(num_tokens) # 记录真实 token 数，供 g 侧裁剪
tp_group = get_tp_group()
tp_size = tp_group.world_size
if tp_size == 1:
return hidden_states
padded = ceil_align(num_tokens, tp_size)
if padded != num_tokens:
hidden_states = torch.nn.functional.pad(
hidden_states, (0, 0, 0, padded - num_tokens)
)
return hidden_states.tensor_split(tp_size)[tp_group.rank_in_group].contiguous()
def sp_exit_gather(hidden_states: torch.Tensor, num_tokens: int) -> torch.Tensor:
'''g：沿 dim 0 all-gather 各 rank 的 shard，再 narrow 到真实 token 数。'''
tp_group = get_tp_group()
tp_size = tp_group.world_size
if tp_size == 1:
return hidden_states[:num_tokens]
output = hidden_states.new_empty(
(hidden_states.shape[0] * tp_size, *hidden_states.shape[1:])
)
tp_group.all_gather_into_tensor(output, hidden_states.contiguous())
return output[:num_tokens]
def maybe_exit_gather(*, hidden_states, hidden_states_before_norm, input_ids, forward_mode):
'''LM head 之前撤销 sequence sharding 并离开 SP 区域。
token 数取自 input_ids、判定用 runs_sp，而不是读 sp_active：CUDA graph
回放时捕获区内写入不会重执行，读 flag 会把序列分片的 hidden states
喂给 LM head 造成错误输出。
'''
if not runs_sp(forward_mode) or input_ids is None:
return hidden_states, hidden_states_before_norm
num_tokens = input_ids.shape[0]
hidden_states = sp_exit_gather(hidden_states, num_tokens=num_tokens)
if hidden_states_before_norm is not None:
hidden_states_before_norm = sp_exit_gather(
hidden_states_before_norm, num_tokens=num_tokens
)
get_forward().set('sp_active', False)
return hidden_states, hidden_states_before_norm
python/sglang/srt/layers/communicator.py
LayerCommunicator 新增 _sp_variant 全 SCATTERED sibling 并委托 prepare_attn / prepare_mlp / postprocess_layer，是 SP 区域在解码器内运转的枢纽。
class LayerCommunicator:
def __init__(self, ..., _is_sp_variant: bool = False):
# SP 启用时，norm/residual 在序列 shard 上运行且不需要 collective，
# 因此构造一个全 SCATTERED 的 sibling 通信器并委托给它；
# _is_sp_variant 防止 sibling 递归再套一层。
self._sp_variant: Optional[LayerCommunicator] = None
if not _is_sp_variant and layernorm_sp.layernorm_sp_enabled():
self._sp_variant = LayerCommunicator(
layer_scatter_modes=LayerScatterModes(
layer_input_mode=ScatterMode.SCATTERED,
attn_mode=ScatterMode.SCATTERED,
mlp_mode=ScatterMode.SCATTERED,
middle_residual_mode=ScatterMode.SCATTERED,
layer_output_mode=ScatterMode.SCATTERED,
),
input_layernorm=input_layernorm,
post_attention_layernorm=post_attention_layernorm,
allow_reduce_scatter=allow_reduce_scatter,
is_last_layer=is_last_layer,
qkv_latent_func=qkv_latent_func,
force_layernorm_before_dp_gather=force_layernorm_before_dp_gather,
enable_fused_ar_quant=enable_fused_ar_quant,
fused_ar_quant_keep_bf16=fused_ar_quant_keep_bf16,
_is_sp_variant=True,
)
def prepare_attn(self, hidden_states, residual, forward_batch, ...):
# residual 为 None 标记首个 decoder layer，SP 区域在此打开；
# 每个 forward 都重新判定，避免中途崩溃把状态泄漏到下一轮。
if self._sp_variant is not None:
if residual is None:
get_forward().set(
'sp_active', forward_batch.forward_mode == ForwardMode.EXTEND
)
if get_forward().sp_active:
# 入口 scatter：此后各层看到的 token 数都是本 rank 的 1/tp 分片。
hidden_states = layernorm_sp.sp_entry_scatter(hidden_states)
if get_forward().sp_active:
# 委托给全 SCATTERED sibling：norm/residual 在本地 shard 上
# 计算，不触发任何 gather/scatter collective。
return self._sp_variant.prepare_attn(
hidden_states, residual, forward_batch, quant_format, post_residual_addition
)
# 非 SP：原有 input_scattered 判定与 norm 逻辑继续执行。
python/sglang/srt/layers/linear.py
Row / ColumnParallelLinear 在 sp_active 时自动参与 g / g-bar，无需在模型代码里逐个标记，是零侵入解耦的关键。
class ColumnParallelLinear(LinearBase):
def forward(self, input_):
bias = self.bias if not self.skip_bias_add else None
# Megatron SP 的 g：输入是本 rank 的 [M_pad/tp, K] 序列分片，
# 先 all-gather 回完整序列再 matmul。参与层（qkv / gate_up）的
# gather_output=False，因此没有输出侧 all-gather 需要处理。
if get_forward().sp_active and self.tp_size > 1:
output = layernorm_sp.column_parallel_g_matmul(self, input_, bias)
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
# 非 SP 普通路径：量化方法 matmul + 可选输出 all-gather。
output_parallel = self.quant_method.apply(self, input_, bias)
if self.gather_output:
output = tensor_model_parallel_all_gather(output_parallel)
else:
output = output_parallel
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
class RowParallelLinear(LinearBase):
def forward(self, input_, skip_all_reduce=False, ...):
# bias 只在 rank 0 参与 GEMM，避免 TP>1 时被重复累加。
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
# Megatron SP 的 g-bar：沿 token 维 reduce-scatter 替代 all-reduce，
# 输出保持分片供下一段 SP LayerNorm 使用。o_proj / down 的
# reduce_results=False，因此 SP 下由线性层自己承担归约。
if (
get_forward().sp_active
and self.tp_size > 1
and not skip_all_reduce
and output_tensor is None
):
output = layernorm_sp.row_parallel_gbar_matmul(self, input_parallel, bias_)
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
# 非 SP 路径：原 reduce / all-reduce 逻辑继续执行。
评论区精华
review 的核心交锋集中在架构收敛、模型解耦与测试配套：
风险与影响
- 风险：
- CUDA graph 状态一致性：
sp_active 在捕获区内写入、回放时过期，已用 runs_sp 规避并有回归测试兜底；但 _SPForwardState.num_tokens 是模块级全局实例，若 forward 中途异常退出，可能把上一轮 token 数泄漏到下一 forward（代码通过每 forward 重写缓解该风险）。
- 依赖 torch 内部 API：fused 快路径探测
torch.ops.symm_mem.fused_matmul_reduce_scatter / fused_all_gather_matmul，torch 升级可能改变这些内部符号；有探测+fallback 缓解，但 fallback 会丢失 B200 上的主要收益。
- 覆盖范围窄：白名单仅
Qwen3ForCausalLM，e2e 只跑 Qwen3-8B / tp=2；padding 非对齐、更大 TP、B200 fused path 的正确性主要依赖手动验证。
- NVLink 前置未强制：PR body 明说 "NVLink/NVSwitch is required in practice but not yet enforced"，非 NVLink 环境开启可能无收益甚至负收益。
- 默认路径影响小：默认关闭且分支由
sp_active 门控，普通用户无行为变化；关闭时字节一致。
- 影响：用户侧：默认关闭，行为零变化；开启后仅影响 Qwen3 dense + tp>1 的 prefill，长上下文收益显著（激活内存降 27.7% @ 8×H100），精度无回归。系统侧：核心执行路径（linear / communicator / logits_processor）新增 SP 分支，但由 sp_active 门控，默认路径计算开销几乎为零。团队侧：建立了可扩展的 SP 接入模板——新模型家族只需 allowlist 一行 + 硬件验证；PR body 已列出 B300、DP-attention 感知 SP、投机解码、chunked/split prefill 等后续方向。
- 风险标记：CUDA graph 标志过期风险, 依赖 torch 内部 symm_mem API, 仅 Qwen3 白名单, NVLink 前置未强制, 异常路径状态泄漏
关联脉络
- PR #36933 [2/N][Mixed] Mixed chunk prefill with spec enabled: 同为调度/前向路径改造：SP 的 prefill-only 判定基于 forward_mode == EXTEND，mixed chunk prefill 扩展了 forward 模式组合，二者在 runs_sp 判定与 EXTEND 语义上存在潜在交互，需要后续回归验证。
参与讨论