minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] Optimize TP LMHead with All-to-All (#32313)

Browse Source
This commit is contained in:
Siyuan Chen
2026-08-17 19:55:27 -07:00
committed by GitHub
parent d6c837489a
commit fcdaaf8a5d
7 changed files with 226 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+41 -1
View File
@@ -2404,14 +2404,54 @@ def _data_parallelism_defaults(view: Any) -> dict:
return {}
@register_post_process
def _tp_lm_head_all_to_all_default(view: Any) -> dict:
"""Enable the TP LM-head all-to-all path only for pure-DP decode nodes.
Prefill-only and colocated nodes keep the feature disabled by default: the
LM-head weight layout is fixed at load time, so enabling the TP path would
also move their long prefills away from the communication-free DP LM head.
An explicit CLI value always wins.
"""
if view.enable_tp_lm_head_all_to_all is not None:
return {}
enable = (
view.disaggregation_mode == "decode"
and view.enable_dp_attention
and view.dp_size > 1
and view.tp_size == view.dp_size
and view.attn_cp_size == 1
and not view.enable_dp_lm_head
)
return {"enable_tp_lm_head_all_to_all": enable}
@register_post_process
def _dp_lm_head_validation(view: Any) -> dict:
"""Read-only validation pass: dp-attention is a prerequisite for the
dp LM head. Reads the mid-resolution values through the view."""
dp LM head and the TP LM-head all-to-all path. Reads the mid-resolution
values through the view."""
if view.enable_dp_lm_head:
assert (
view.enable_dp_attention
), "Please enable dp attention when setting enable_dp_lm_head. "
if view.enable_tp_lm_head_all_to_all:
assert view.enable_dp_attention, (
"Please enable dp attention when setting " "enable_tp_lm_head_all_to_all."
)
assert not view.enable_dp_lm_head, (
"--enable-tp-lm-head-all-to-all uses a TP-sharded LM head and is "
"incompatible with --enable-dp-lm-head."
)
assert view.tp_size == view.dp_size, (
"--enable-tp-lm-head-all-to-all currently requires tp_size == "
f"dp_size, got tp_size={view.tp_size}, dp_size={view.dp_size}."
)
assert view.attn_cp_size == 1, (
"--enable-tp-lm-head-all-to-all currently requires "
f"attn_cp_size == 1, got {view.attn_cp_size}."
)
return {}
python/sglang/srt/distributed/bootstrap.py
+53
View File
@@ -45,6 +45,14 @@ logger = logging.getLogger(__name__)
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu_arm64 = is_host_cpu_arm64()
# A representative per-peer payload for materializing the PyNCCL P2P
# connections used by TP LM-head all-to-all. The input/output tensors are
# temporary; NCCL owns the transport resources retained after the warmup.
# In dsv4-pro, assume bs per dp is 120 and the vocab_size is 129280.
# Therefore, the chunk size that each peer sends is 120*129280/8=1.849MB.
# The total warmup bytes per peer should be 1.849*2 = 4MB
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20
class TorchDistributedResult(msgspec.Struct, frozen=True, kw_only=True):
tp_group: object
@@ -110,6 +118,17 @@ def init_torch_distributed(
tp_size=ps.tp_size, pp_size=ps.pp_size, moe_ep_size=ps.moe_ep_size
)
# CUDA graph capture enables the PyNCCL communicator for TP LM-head
# all-to-all. Exercise that exact send/recv path before measuring
# pre_model_load_memory so its persistent transport allocations are
# included in later KV-cache sizing instead of appearing during capture.
if (
device == "cuda"
and get_parallel().enable_tp_lm_head_all_to_all
and ps.tp_size > 1
):
_prewarm_tp_lm_head_all_to_all()
pre_model_load_memory = get_available_gpu_memory(
device,
ps.gpu_id,
@@ -271,6 +290,40 @@ def _prewarm_nccl(*, tp_size: int, pp_size: int, moe_ep_size: int) -> None:
)
def _prewarm_tp_lm_head_all_to_all() -> None:
"""Materialize PyNCCL P2P resources before model-memory accounting."""
warmup_start = time.perf_counter()
tp_group = get_tp_group()
pynccl_comm = tp_group.pynccl_comm
if pynccl_comm is None or not pynccl_comm.available:
raise RuntimeError(
"--enable-tp-lm-head-all-to-all requires an available PyNCCL "
"communicator for CUDA graph capture."
)
numel = tp_group.world_size * _TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER
warmup_input = torch.empty(numel, dtype=torch.uint8, device=tp_group.device)
warmup_output = torch.empty_like(warmup_input)
# PyNCCL is disabled outside graph-capture contexts by default. Enable it
# explicitly so eager startup does not fall back to ProcessGroupNCCL and
# miss the P2P resources required by the captured all-to-all.
with pynccl_comm.change_state(enable=True):
pynccl_comm.all_to_all_single(warmup_output, warmup_input)
current_platform.synchronize()
del warmup_input, warmup_output
current_platform.empty_cache()
warmup_elapsed = time.perf_counter() - warmup_start
logger.info(
"TP LM-head PyNCCL all-to-all warmup completed in %.3fs "
"(tp_size=%d, bytes_per_peer=%d)",
warmup_elapsed,
tp_group.world_size,
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER,
)
def _check_tp_memory_balance(
*, pre_model_load_memory: float, local_gpu_memory: float
) -> None:
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
+4 -2
View File
@@ -176,7 +176,9 @@ def _compile_deep_gemm_one_type_all(
m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0)))
# Here the precompilation is only run on the first rank, so gpu_id should be 0
memory_budget = get_available_gpu_memory(device="cuda", gpu_id=0)
memory_budget = get_available_gpu_memory(
device="cuda", gpu_id=torch.cuda.current_device()
)
# If the memory budget is less memory requirement, we need to reduce max_m to avoid out of memory, which might further cause hanging during warmup
max_m = max(m_list)
@@ -193,7 +195,7 @@ def _compile_deep_gemm_one_type_all(
kernel_type, max_m=max_m, n=n, k=k, num_groups=num_groups
)
> memory_budget
and max_m > 4096
and max_m > 2048
):
max_m = max_m // 2
logger.warning(
python/sglang/srt/layers/logits_processor.py
+80 -3
View File
@@ -24,6 +24,7 @@ from torch import nn
from sglang.kernels.ops.activation.softcap import (
softcap_inplace_logits as fused_softcap,
)
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.layers.aux_hidden_states import (
AuxHiddenStates,
@@ -292,6 +293,7 @@ class LogitsProcessor(nn.Module):
self.vocab_size = config.vocab_size
self.logit_scale = logit_scale
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
if self.use_attn_tp_group:
self.attn_tp_size = get_parallel().attn_tp_size
@@ -666,15 +668,22 @@ class LogitsProcessor(nn.Module):
if self.logit_scale is not None:
logits.mul_(self.logit_scale)
used_tp_lm_head_all_to_all = False
if self.do_tensor_parallel_all_gather:
if self.use_attn_tp_group:
logits = self._gather_attn_tp_logits(logits)
elif self._can_use_tp_lm_head_all_to_all(
logits, local_hidden_states, lm_head, logits_metadata
):
logits = self._tp_lm_head_all_to_all(logits)
used_tp_lm_head_all_to_all = True
else:
logits = self._logits_gatherer(logits)
logits = self._scatter_dp_attn_logits(
logits, local_hidden_states, logits_metadata
)
if not used_tp_lm_head_all_to_all:
logits = self._scatter_dp_attn_logits(
logits, local_hidden_states, logits_metadata
)
logits = self._copy_logits_to_buffer(
logits, logits_metadata, use_buffer=use_logits_buffer
@@ -793,6 +802,54 @@ class LogitsProcessor(nn.Module):
)
return global_logits
def _can_use_tp_lm_head_all_to_all(
self,
logits: torch.Tensor,
local_hidden_states: torch.Tensor,
lm_head: VocabParallelEmbedding,
logits_metadata: LogitsMetadata,
) -> bool:
if not self.use_tp_lm_head_all_to_all:
return False
tp_size = get_parallel().tp_size
base_lm_head = getattr(lm_head, "base_layer", lm_head)
if getattr(base_lm_head, "tp_size", None) != tp_size:
# Tied embeddings may be replicated across DP ranks (tp_size=1),
# even though the logits processor runs in a larger global TP
# group. Such logits are full-vocabulary rather than TP shards and
# therefore do not satisfy the all-to-all layout contract.
return False
# Every participant must make the same collective choice. Decode CUDA
# graphs omit CPU counts and fill every GPU count with the same padded
# bucket size. Eager batches carry the same global CPU count list on
# every rank, so they are also safe when all entries are equal.
global_counts_cpu = logits_metadata.global_num_tokens_for_logprob_cpu
is_equal_padded_graph_layout = global_counts_cpu is None and (
logits_metadata.global_num_tokens_for_logprob_gpu is not None
)
is_equal_eager_layout = (
global_counts_cpu is not None
and len(global_counts_cpu) == tp_size
and len(global_counts_cpu) > 0
and all(count == global_counts_cpu[0] for count in global_counts_cpu)
)
if not (is_equal_padded_graph_layout or is_equal_eager_layout):
return False
local_rows = local_hidden_states.shape[0]
return local_rows > 0 and logits.shape[0] == local_rows * tp_size
def _tp_lm_head_all_to_all(self, logits: torch.Tensor) -> torch.Tensor:
"""Exchange only the row block owned by each destination DP rank."""
logits = logits.contiguous()
all_to_all_output = torch.empty_like(logits)
get_tp_group().all_to_all_single(all_to_all_output.view(-1), logits.view(-1))
return _reassemble_tp_lm_head_all_to_all_output(
all_to_all_output, get_parallel().tp_size
)
def _scatter_dp_attn_logits(
self,
logits: torch.Tensor,
@@ -951,6 +1008,26 @@ class LogitsProcessor(nn.Module):
)
def _reassemble_tp_lm_head_all_to_all_output(
all_to_all_output: torch.Tensor, tp_size: int
) -> torch.Tensor:
"""Convert source-major all-to-all output to row-major full-vocab logits.
Each source TP rank contributes ``[local_rows, vocab_shard]`` for this
destination DP rank. ``all_to_all_single`` concatenates those contributions
along dim 0, while the sampler expects the vocab shards concatenated along
dim 1.
"""
assert all_to_all_output.shape[0] % tp_size == 0
local_rows = all_to_all_output.shape[0] // tp_size
vocab_shard = all_to_all_output.shape[1]
return (
all_to_all_output.view(tp_size, local_rows, vocab_shard)
.permute(1, 0, 2)
.reshape(local_rows, tp_size * vocab_shard)
)
def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool:
return all(hasattr(lm_head, attr_name) for attr_name in attr_names)
python/sglang/srt/server_args.py
+23 -4
View File
@@ -1166,6 +1166,22 @@ class ServerArgs:
),
NS("parallel"),
] = False
enable_tp_lm_head_all_to_all: A[
Optional[bool],
Arg(
help="Use all-to-all instead of TP all-gather followed by DP scatter "
"for the TP-sharded LM head under DP attention. By default this is "
"enabled only on decode-only PD nodes with pure DP attention "
"(tp_size == dp_size > 1 and attn_cp_size == 1), and disabled on "
"prefill-only and colocated nodes. Pass "
"--no-enable-tp-lm-head-all-to-all to opt out. The path is "
"incompatible with --enable-dp-lm-head; batches without an equal "
"padded row count fall back to the existing all-gather path.",
action=argparse.BooleanOptionalAction,
resolvable=True,
),
NS("parallel"),
] = None
enable_attn_tp_input_scattered: A[
bool,
"Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent.",
@@ -6692,11 +6708,14 @@ class ServerArgs:
prefill_cfg.max_bs
)
# The dp-lm-head validation moved to the resolution pipeline
# (arg_groups/overrides.py: _dp_lm_head_validation), invoked here at
# its legacy slot.
from sglang.srt.arg_groups.overrides import _dp_lm_head_validation
# Resolve the phase-aware TP LM-head default before validating the
# resulting DP/TP LM-head configuration.
from sglang.srt.arg_groups.overrides import (
_dp_lm_head_validation,
_tp_lm_head_all_to_all_default,
)
run_post_process_pass(self, _tp_lm_head_all_to_all_default)
run_post_process_pass(self, _dp_lm_head_validation)
def _handle_moe_kernel_config(self):