From 315306d8a9f93ed24e3a6d532b3c1da41e83ca9d Mon Sep 17 00:00:00 2001 From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Date: Wed, 4 Feb 2026 05:58:28 -0600 Subject: [PATCH] Make sure we always disable symm memory without dp padding (#18129) --- python/sglang/srt/distributed/parallel_state.py | 6 +++++- python/sglang/srt/layers/dp_attention.py | 2 +- python/sglang/srt/layers/vocab_parallel_embedding.py | 5 ++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 3070178b6..d45c87d72 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -320,9 +320,11 @@ class GroupCoordinator: from sglang.srt.distributed.device_communicators.torch_symm_mem import ( TorchSymmMemCommunicator, ) + from sglang.srt.layers.dp_attention import is_allocation_symmetric self.is_symmetric_memory_enabled = is_symmetric_memory_enabled self.use_symmetric_memory = use_symmetric_memory + self.is_allocation_symmetric = is_allocation_symmetric if is_hip(): from sglang.srt.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, @@ -807,7 +809,9 @@ class GroupCoordinator: # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 output_size = (input_size[0] * world_size,) + input_size[1:] # Allocate output tensor. - with self.use_symmetric_memory(self): + with self.use_symmetric_memory( + self, disabled=not self.is_allocation_symmetric() + ): output_tensor = torch.empty( output_size, dtype=input_.dtype, device=input_.device ) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 0b5e2765d..a215913e6 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -114,7 +114,7 @@ class _DpGatheredBufferWrapper: @classmethod def get_global_dp_buffer(cls) -> torch.Tensor: - with use_symmetric_memory(get_tp_group()): + with use_symmetric_memory(get_tp_group(), disabled=not cls._dp_max_padding): buffer = torch.empty( (cls._global_dp_buffer_len, cls._hidden_size), dtype=cls._dtype, diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 1df12f66e..63b376c24 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -23,6 +23,7 @@ from sglang.srt.layers.dp_attention import ( attn_tp_all_reduce, get_attention_tp_rank, get_attention_tp_size, + is_allocation_symmetric, ) from sglang.srt.layers.parameter import BasevLLMParameter from sglang.srt.layers.quantization.base_config import ( @@ -482,7 +483,9 @@ class VocabParallelEmbedding(torch.nn.Module): masked_input = input_ # Get the embeddings. - with use_symmetric_memory(get_tp_group(), disabled=not self.enable_tp): + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): output_parallel = self.quant_method.embedding(self, masked_input.long()) if self.tp_size > 1: