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

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 {}
 
 

python/sglang/srt/distributed/bootstrap.py

@@ -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

@@ -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

@@ -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,12 +668,19 @@ 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)
 
+        if not used_tp_lm_head_all_to_all:
             logits = self._scatter_dp_attn_logits(
                 logits, local_hidden_states, logits_metadata
             )
@@ -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

@@ -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
+        # Resolve the phase-aware TP LM-head default before validating the
-        # (arg_groups/overrides.py: _dp_lm_head_validation), invoked here at
+        # resulting DP/TP LM-head configuration.
-        # its legacy slot.
+        from sglang.srt.arg_groups.overrides import (
-        from sglang.srt.arg_groups.overrides import _dp_lm_head_validation
+            _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):

test/registered/mock_model/test_e2e_tp.py

@@ -4,7 +4,7 @@ import unittest
 
 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 @@ class TestE2ETensorParallel(CustomTestCase):
             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()

test/registered/unit/test_model_overrides.py

@@ -69,6 +69,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
                     "quantization",
                     "enable_dp_attention",
                     "enable_dp_lm_head",
+                    "enable_tp_lm_head_all_to_all",
                     "moe_a2a_backend",
                     "ep_size",
                     "moe_dense_tp_size",