From 1109acc24b06f7f72c89ee474b4b498d3a6af0c1 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 20 Jun 2026 18:21:38 +0800 Subject: [PATCH] [diffusion] optimize: shard hunyuan text tokens under sp (#28319) --- .../runtime/layers/attention/layer.py | 24 +++++- .../runtime/models/dits/hunyuanvideo.py | 77 ++++++++++++++++--- 2 files changed, 88 insertions(+), 13 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 71628fd30..3c8c8be95 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -40,7 +40,9 @@ from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import ( ) from sglang.multimodal_gen.runtime.layers.usp import ( _usp_input_all_to_all, + _usp_input_all_to_all_varlen, _usp_output_all_to_all, + _usp_output_all_to_all_varlen, ring_attn, ) from sglang.multimodal_gen.runtime.managers.forward_context import ( @@ -143,6 +145,7 @@ class UlyssesAttention(nn.Module): replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, + seq_lens: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Forward pass for distributed attention. @@ -168,11 +171,21 @@ class UlyssesAttention(nn.Module): forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata + if seq_lens is not None: + assert ( + replicated_q is None and replicated_k is None and replicated_v is None + ), "Varlen Ulysses attention does not support replicated QKV" + # Stack QKV qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] # Redistribute heads across sequence dimension - qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) + if seq_lens is None: + qkv = sequence_model_parallel_all_to_all_4D( + qkv, scatter_dim=2, gather_dim=1 + ) + else: + qkv = _usp_input_all_to_all_varlen(qkv, seq_lens, head_dim=2) # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) @@ -204,9 +217,12 @@ class UlyssesAttention(nn.Module): # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) - output = sequence_model_parallel_all_to_all_4D( - output, scatter_dim=1, gather_dim=2 - ) + if seq_lens is None: + output = sequence_model_parallel_all_to_all_4D( + output, scatter_dim=1, gather_dim=2 + ) + else: + output = _usp_output_all_to_all_varlen(output, seq_lens, head_dim=2) return output, replicated_output diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index c173e54d5..30d764b8d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -11,7 +11,11 @@ import torch.nn as nn from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size -from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_parallel_world_size, + get_sp_parallel_rank, + get_sp_world_size, +) from sglang.multimodal_gen.runtime.layers.attention import ( LocalAttention, UlyssesAttention, @@ -223,6 +227,8 @@ class MMDoubleStreamBlock(nn.Module): txt: torch.Tensor, vec: torch.Tensor, freqs_cis: tuple, + txt_is_sharded: bool = False, + seq_lens: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # Process modulation vectors img_mod_outputs = self.img_mod(vec) @@ -285,7 +291,16 @@ class MMDoubleStreamBlock(nn.Module): txt_k = self.txt_attn_k_norm(txt_k.contiguous()).to(txt_k.dtype) # Run distributed attention - img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v) + if txt_is_sharded: + attn, _ = self.attn( + torch.cat((img_q, txt_q), dim=1), + torch.cat((img_k, txt_k), dim=1), + torch.cat((img_v, txt_v), dim=1), + seq_lens=seq_lens, + ) + img_attn, txt_attn = attn.split([image_seq_len, text_seq_len], dim=1) + else: + img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v) img_attn_out, _ = self.img_attn_proj( img_attn.reshape(batch_size, image_seq_len, -1) ) @@ -406,6 +421,8 @@ class MMSingleStreamBlock(nn.Module): vec: torch.Tensor, txt_len: int, freqs_cis: tuple[torch.Tensor, torch.Tensor], + txt_is_sharded: bool = False, + seq_lens: list[int] | None = None, ) -> torch.Tensor: # Process modulation mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) @@ -445,12 +462,19 @@ class MMSingleStreamBlock(nn.Module): ) # Run distributed attention - img_attn_output, txt_attn_output = self.attn( - img_q, img_k, img_v, txt_q, txt_k, txt_v - ) - attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1).view( - batch_size, seq_len, -1 - ) + if txt_is_sharded: + attn_output, _ = self.attn( + torch.cat((img_q, txt_q), dim=1), + torch.cat((img_k, txt_k), dim=1), + torch.cat((img_v, txt_v), dim=1), + seq_lens=seq_lens, + ) + else: + img_attn_output, txt_attn_output = self.attn( + img_q, img_k, img_v, txt_q, txt_k, txt_v + ) + attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1) + attn_output = attn_output.view(batch_size, seq_len, -1) # Process MLP activation mlp_output = self.mlp_act(mlp) @@ -690,7 +714,33 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi img = self.img_in(img) txt = self.txt_in(txt, t) txt_seq_len = txt.shape[1] + sp_size = get_sp_world_size() + txt_is_sharded = ( + sp_size > 1 + and get_ring_parallel_world_size() == 1 + and txt_seq_len >= sp_size + and not torch.is_grad_enabled() + ) + seq_lens = None + if txt_is_sharded: + sp_rank = get_sp_parallel_rank() + base_text_shard_len = txt_seq_len // sp_size + extra_text_tokens = txt_seq_len % sp_size + text_seq_lens = [ + base_text_shard_len + (1 if rank < extra_text_tokens else 0) + for rank in range(sp_size) + ] + text_shard_start = base_text_shard_len * sp_rank + min( + sp_rank, extra_text_tokens + ) + text_shard_len = text_seq_lens[sp_rank] + txt = txt[ + :, text_shard_start : text_shard_start + text_shard_len + ].contiguous() + txt_seq_len = text_shard_len img_seq_len = img.shape[1] + if txt_is_sharded: + seq_lens = [img_seq_len + text_len for text_len in text_seq_lens] freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None @@ -706,7 +756,14 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi # Process through double stream blocks for index, block in enumerate(self.double_blocks): - double_block_args = [img, txt, vec, freqs_cis] + double_block_args = [ + img, + txt, + vec, + freqs_cis, + txt_is_sharded, + seq_lens, + ] img, txt = block(*double_block_args) # Merge txt and img to pass through single stream blocks x = torch.cat((img, txt), 1) @@ -719,6 +776,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi vec, txt_seq_len, freqs_cis, + txt_is_sharded, + seq_lens, ] x = block(*single_block_args)