diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 619d32b11..72cb802d8 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -406,9 +406,11 @@ def initialize_model_parallel( try: from .yunchang import PROCESS_GROUP as _YC_PROCESS_GROUP - from .yunchang import set_seq_parallel_pg as _set_seq_parallel_pg + from .yunchang import ( + set_seq_parallel_pg_by_sp_groups as _set_seq_parallel_pg_by_sp_groups, + ) except ImportError: - _set_seq_parallel_pg = None + _set_seq_parallel_pg_by_sp_groups = None class _DummyProcessGroup: ULYSSES_PG = torch.distributed.group.WORLD @@ -416,11 +418,15 @@ def initialize_model_parallel( PROCESS_GROUP = _DummyProcessGroup() else: - _set_seq_parallel_pg( + # Build yunchang SP sub-groups based on the true SP groups. This is + # critical when TP>1, because SP groups may be strided in global ranks + # (e.g., tp-sp order). + sp_groups = rank_generator.get_ranks("sp") + _set_seq_parallel_pg_by_sp_groups( sp_ulysses_degree=ulysses_degree, sp_ring_degree=ring_degree, - rank=get_world_group().rank_in_group, - world_size=dit_parallel_size, + rank=get_world_group().rank, + sp_groups=sp_groups, ) PROCESS_GROUP = _YC_PROCESS_GROUP diff --git a/python/sglang/multimodal_gen/runtime/distributed/yunchang.py b/python/sglang/multimodal_gen/runtime/distributed/yunchang.py index e77c3a43b..241e57e8c 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/yunchang.py +++ b/python/sglang/multimodal_gen/runtime/distributed/yunchang.py @@ -82,3 +82,72 @@ def set_seq_parallel_pg( PROCESS_GROUP.ULYSSES_PG = ulyssess_pg PROCESS_GROUP.RING_PG = ring_pg + + +def set_seq_parallel_pg_by_sp_groups( + sp_ulysses_degree, + sp_ring_degree, + rank: int, + sp_groups: list[list[int]], + use_ulysses_low: bool = True, +): + """Create Ulysses/Ring process groups inside each SP group. + + This is required when TP>1, because SP groups are not necessarily made of + consecutive global ranks (e.g., tp-sp order makes SP ranks strided). + + Args: + sp_ulysses_degree: ulysses degree inside SP. + sp_ring_degree: ring degree inside SP. + rank: global rank of current process. + sp_groups: list of global-rank lists for each SP group. + use_ulysses_low: keep the same semantics as the original function. + """ + sp_degree = sp_ring_degree * sp_ulysses_degree + assert sp_degree > 0 + assert all( + len(g) == sp_degree for g in sp_groups + ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}" + + ulyssess_pg = None + ring_pg = None + + num_ulysses_pgs = sp_ring_degree + num_ring_pgs = sp_ulysses_degree + + def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]: + return [ranks[i] for i in indices] + + # Important: call torch.distributed.new_group in the same order on all ranks. + for sp_ranks in sp_groups: + if use_ulysses_low: + for i in range(num_ulysses_pgs): + idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree)) + ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) + group = torch.distributed.new_group(ulysses_ranks) + if rank in ulysses_ranks: + ulyssess_pg = group + + for i in range(num_ring_pgs): + idx = list(range(i, sp_degree, num_ring_pgs)) + ring_ranks = _map_indices_to_ranks(sp_ranks, idx) + group = torch.distributed.new_group(ring_ranks) + if rank in ring_ranks: + ring_pg = group + else: + for i in range(num_ring_pgs): + idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree)) + ring_ranks = _map_indices_to_ranks(sp_ranks, idx) + group = torch.distributed.new_group(ring_ranks) + if rank in ring_ranks: + ring_pg = group + + for i in range(num_ulysses_pgs): + idx = list(range(i, sp_degree, num_ulysses_pgs)) + ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx) + group = torch.distributed.new_group(ulysses_ranks) + if rank in ulysses_ranks: + ulyssess_pg = group + + PROCESS_GROUP.ULYSSES_PG = ulyssess_pg + PROCESS_GROUP.RING_PG = ring_pg diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index cbe9dd946..7f3516a4d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -21,7 +21,7 @@ except Exception as _e: pass from sglang.multimodal_gen.runtime.layers.activation import get_act_fn -from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP @@ -219,8 +219,12 @@ class ModulateProjection(nn.Module): super().__init__() self.factor = factor self.hidden_size = hidden_size - self.linear = ReplicatedLinear( - hidden_size, hidden_size * factor, bias=True, params_dtype=dtype + self.linear = ColumnParallelLinear( + hidden_size, + hidden_size * factor, + bias=True, + gather_output=True, + params_dtype=dtype, ) self.act = get_act_fn(act_layer) diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 52fc3ce37..4a7b34808 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -170,6 +170,7 @@ class Scheduler: if isinstance(req, Req): warmup_req = deepcopy(req) warmup_req.is_warmup = True + warmup_req.extra["cache_dit_num_inference_steps"] = req.num_inference_steps warmup_req.num_inference_steps = 1 recv_reqs.insert(0, (identity, warmup_req)) logger.info("Server warming up....") diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 60b98eac4..97ada9791 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -25,7 +25,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import ( ScaleResidual, ScaleResidualLayerNormScaleShift, ) -from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, @@ -132,10 +132,10 @@ class WanSelfAttention(nn.Module): self.parallel_attention = parallel_attention # layers - self.to_q = ReplicatedLinear(dim, dim) - self.to_k = ReplicatedLinear(dim, dim) - self.to_v = ReplicatedLinear(dim, dim) - self.to_out = ReplicatedLinear(dim, dim) + self.to_q = ColumnParallelLinear(dim, dim, gather_output=True) + self.to_k = ColumnParallelLinear(dim, dim, gather_output=True) + self.to_v = ColumnParallelLinear(dim, dim, gather_output=True) + self.to_out = ColumnParallelLinear(dim, dim, gather_output=True) self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() @@ -218,8 +218,8 @@ class WanI2VCrossAttention(WanSelfAttention): supported_attention_backends=supported_attention_backends, ) - self.add_k_proj = ReplicatedLinear(dim, dim) - self.add_v_proj = ReplicatedLinear(dim, dim) + self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=True) + self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=True) self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() @@ -272,11 +272,11 @@ class WanTransformerBlock(nn.Module): # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) + self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) + self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) - self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) if attention_type == "sla": self.attn1 = MinimalA2AAttnOp( SparseLinearAttention( @@ -454,12 +454,14 @@ class WanTransformerBlock_VSA(nn.Module): # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) - self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True) + self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) + self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) + self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) + self.to_gate_compress = ColumnParallelLinear( + dim, dim, bias=True, gather_output=True + ) - self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) self.attn1 = UlyssesAttention_VSA( num_heads=num_heads, head_size=dim // num_heads, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 6ca92bd8b..e6a9412a2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -487,6 +487,12 @@ class DenoisingStage(PipelineStage): """ assert self.transformer is not None pipeline = self.pipeline() if self.pipeline else None + # NOTE: In warmup requests we may override req.num_inference_steps (e.g. set to 1) + # for latency amortization, but cache-dit needs the *original* total steps to + # initialize/refresh its context correctly. + cache_dit_num_inference_steps = batch.extra.get( + "cache_dit_num_inference_steps", batch.num_inference_steps + ) if not server_args.model_loaded["transformer"]: # FIXME: reuse more code loader = TransformerLoader() @@ -494,13 +500,13 @@ class DenoisingStage(PipelineStage): server_args.model_paths["transformer"], server_args, "transformer" ) # enable cache-dit before torch.compile (delayed mounting) - self._maybe_enable_cache_dit(batch.num_inference_steps) + self._maybe_enable_cache_dit(cache_dit_num_inference_steps) self.compile_module_with_torch_compile(self.transformer) if pipeline: pipeline.add_module("transformer", self.transformer) server_args.model_loaded["transformer"] = True else: - self._maybe_enable_cache_dit(batch.num_inference_steps) + self._maybe_enable_cache_dit(cache_dit_num_inference_steps) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 95b45bff7..1683fd0a2 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -366,10 +366,14 @@ class ServerArgs: if self.vae_cpu_offload is None: self.vae_cpu_offload = False else: - self.dit_cpu_offload = True - self.text_encoder_cpu_offload = True - self.image_encoder_cpu_offload = True - self.vae_cpu_offload = True + if self.dit_cpu_offload is None: + self.dit_cpu_offload = True + if self.text_encoder_cpu_offload is None: + self.text_encoder_cpu_offload = True + if self.image_encoder_cpu_offload is None: + self.image_encoder_cpu_offload = True + if self.vae_cpu_offload is None: + self.vae_cpu_offload = True def __post_init__(self): # configure logger before use @@ -1005,9 +1009,9 @@ class ServerArgs: has_sp = self.sp_degree > 1 has_tp = self.tp_size > 1 if has_sp and has_tp: - raise ValueError( - "cache-dit does not support hybrid parallelism (SP + TP). " - "Please use either sequence parallelism or tensor parallelism, not both." + logger.warning( + "cache-dit is enabled with hybrid parallelism (SP + TP). " + "Proceeding anyway (SGLang integration may support this mode)." ) def _set_default_attention_backend(self) -> None: diff --git a/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py index 7ee33e4a9..7d7654ad4 100644 --- a/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py +++ b/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py @@ -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