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

[Diffusion] Fix Ulysses/Ring process group construction under TP to enable correct Wan2.2 tensor parallelism (#16532)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-01-07 20:52:36 +08:00
committed by GitHub
parent 62d0280f62
commit 32a6540afc
8 changed files with 144 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py
View File

@@ -28,6 +28,8 @@ from cache_dit import (
from cache_dit.caching.block_adapters import BlockAdapterRegister
from cache_dit.parallelism import ParallelismBackend, ParallelismConfig
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_group
_original_similarity = None
@@ -53,7 +55,8 @@ def _patch_cache_dit_similarity():
sp_group = getattr(self, "_sglang_sp_group", None)
tp_group = getattr(self, "_sglang_tp_group", None)
target_group = sp_group or tp_group
tp_sp_group = getattr(self, "_sglang_tp_sp_group", None)
target_group = tp_sp_group or sp_group or tp_group
if target_group is None:
return _original_similarity(
@@ -316,6 +319,14 @@ def enable_cache_on_transformer(
if context_manager is not None:
context_manager._sglang_sp_group = sp_group
context_manager._sglang_tp_group = tp_group
# In mixed TP + SP (Ulysses/Ring) mode, cache-dit decisions must be consistent
# across the full TP×SP model-parallel slice. Prefer using SGLang's DIT group
# as a conservative superset group; fallback to None.
tp_sp_group = None
if sp_group is not None and tp_group is not None:
tp_sp_group = get_dit_group()
context_manager._sglang_tp_sp_group = tp_sp_group
return transformer
@@ -488,5 +499,12 @@ def enable_cache_on_dual_transformer(
if context_manager is not None:
context_manager._sglang_sp_group = sp_group
context_manager._sglang_tp_group = tp_group
tp_sp_group = None
if sp_group is not None and tp_group is not None:
try:
tp_sp_group = get_dit_group()
except Exception:
tp_sp_group = None
context_manager._sglang_tp_sp_group = tp_sp_group
return transformer, transformer_2