diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 6a516be2b144..4035f3f440ea 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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 {} diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 4db027f79095..c69628391e5a 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -45,6 +45,14 @@ _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: diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 815c7074096e..518adf56532a 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -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( diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 7cea446d9cb8..2487072b63b2 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -24,6 +24,7 @@ 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 @@ def __init__( 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 @@ def _get_logits( 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 @@ def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor: ) 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 @@ def compute_logprobs_for_multi_item_scoring( ) +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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8596fd959b8b..11d96366c99b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.", @@ -6679,11 +6695,14 @@ def _handle_data_parallelism(self): 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): diff --git a/test/registered/mock_model/test_e2e_tp.py b/test/registered/mock_model/test_e2e_tp.py index 4e177222c18f..415b39726e15 100644 --- a/test/registered/mock_model/test_e2e_tp.py +++ b/test/registered/mock_model/test_e2e_tp.py @@ -4,7 +4,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.mock_model.utils import run_mock_model_bench_serving -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import CustomTestCase, is_in_amd_ci register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large") register_amd_ci(est_time=167, stage="extra-a", runner_config="2-gpu-large-amd") @@ -16,6 +16,29 @@ def test_tp_no_canary_violation(self) -> None: extra_server_args=["--tp", "2", "--mem-fraction-static", "0.88"], ) + @unittest.skipIf(is_in_amd_ci(), "PyNccl CUDA graph smoke test requires CUDA.") + def test_tp_lm_head_all_to_all_cuda_graph(self) -> None: + """Smoke-test the PyNccl all-to-all path used during graph capture.""" + run_mock_model_bench_serving( + extra_server_args=[ + "--tp", + "2", + "--dp", + "2", + "--enable-dp-attention", + "--enable-tp-lm-head-all-to-all", + "--cuda-graph-max-bs-decode", + "4", + "--mem-fraction-static", + "0.88", + "--attention-backend", + "triton", + ], + num_prompts=2, + random_input_len=32, + random_output_len=2, + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 73c502b62047..a7d502bd269e 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -69,6 +69,7 @@ def test_server_args_whitelist_is_exactly_the_migrated_fields(self): "quantization", "enable_dp_attention", "enable_dp_lm_head", + "enable_tp_lm_head_all_to_all", "moe_a2a_backend", "ep_size", "moe_dense_tp_size",