[diffusion] optimize: shard hunyuan text tokens under sp (#28319)

2026-06-20 18:21:38 +08:00
parent a38eba0f1a
commit 1109acc24b
2 changed files with 88 additions and 13 deletions
@@ -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


@@ -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)