From 99df920cdbba20f033ba58765f5da04f59efe4fe Mon Sep 17 00:00:00 2001 From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Date: Thu, 19 Feb 2026 19:32:32 -0600 Subject: [PATCH] Register tensors with symmetric memory for qwen (#18643) --- python/sglang/srt/models/qwen2_moe.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 5e235cfa1..bbb883a2d 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -324,7 +324,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module): final_hidden_states = self._forward_router_experts(hidden_states) if shared_output is not None: - final_hidden_states = final_hidden_states + shared_output + # In-place add is required to keep final_hidden_states in the + # symmetric memory pool (when --enable-symm-mem is used). + # An out-of-place add would allocate a new tensor outside symm + # memory, breaking subsequent symmetric collective operations. + final_hidden_states += shared_output if self.tp_size > 1 and not use_reduce_scatter: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)