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

Fix world-size-one aliasing in MLP batch sync (#34997)

Browse Source 
Co-authored-by: wangwenchen0407 <wangwenchen@meta.com>
This commit is contained in:
Lianmin Zheng
2026-08-17 00:04:09 -07:00
committed by GitHub
co-authored by wangwenchen0407
parent f7a404e9c3
commit 12a455a910
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/dp_attn.py
+8 -4
View File
@@ -137,10 +137,14 @@ class MLPSyncBatchInfo:
local_info_tensor = self._get_local_tensor(device=device) local_info_tensor = self._get_local_tensor(device=device)
fallback_tensor = self._get_fallback_tensor(device=device) fallback_tensor = self._get_fallback_tensor(device=device)
info_width = local_info_tensor.numel() info_width = local_info_tensor.numel()
# Inactive max_world_size slots must decode as IDLE. # Inactive max_world_size slots must decode as IDLE. repeat() (not
global_info_tensor = fallback_tensor.expand( # expand().contiguous()) so the buffer never aliases fallback_tensor:
self.dp_size, self.tp_size * self.cp_size, info_width # at world size 1 the expanded view is already contiguous, contiguous()
).contiguous() # is a no-op, and the masked fallback writes below would then read and
# write the same storage.
global_info_tensor = fallback_tensor.repeat(
self.dp_size, self.tp_size * self.cp_size, 1
)
if use_all_reduce: if use_all_reduce:
# Admission can expose different WORLD sizes; use fixed global slots. # Admission can expose different WORLD sizes; use fixed global slots.