perf: optimize Hunyuan DiT Ulysses and non-attention paths

lightx2v/common/ops/attn/ulysses_attn.py

@@ -1,3 +1,6 @@
+import os
+from contextlib import nullcontext
+
 import torch
 import torch.distributed as dist
 from loguru import logger
@@ -8,6 +11,34 @@
 from .template import AttnWeightTemplate
 from .utils.all2all import all2all_head2seq
 
+
+def _env_flag(name, default="0"):
+    return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"}
+
+
+_PROFILE_RANGES_ENABLED = _env_flag("LIGHTX2V_ULYSSES_PROFILE_RANGES")
+_ASYNC_TEXT_GATHER_ENABLED = _env_flag("LIGHTX2V_ULYSSES_ASYNC_TEXT_GATHER")
+_REUSE_TEXT_GATHER_BUFFERS_ENABLED = _env_flag("LIGHTX2V_ULYSSES_REUSE_TEXT_GATHER_BUFFERS")
+
+
+def _profile_range(name):
+    return torch.profiler.record_function(f"ulysses::{name}") if _PROFILE_RANGES_ENABLED else nullcontext()
+
+
+def _is_split_qkv_input(tensor_or_pair):
+    return isinstance(tensor_or_pair, (tuple, list))
+
+
+def _to_3d_qkv(tensor):
+    if len(tensor.shape) == 4:
+        return tensor.reshape(-1, tensor.shape[-2], tensor.shape[-1])
+    return tensor
+
+
+def _contiguous_if_needed(tensor):
+    return tensor if tensor.is_contiguous() else tensor.contiguous()
+
+
 try:
     from sageattn3_sparse import dequant_fp4 as dequant_fp4_sage3
     from sageattn3_sparse import quant_fp4 as quant_fp4_sage3
@@ -21,6 +52,18 @@
 class UlyssesAttnWeight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
+        self._text_gather_buffers = {}
+
+    def _get_text_gather_buffers(self, tensor, world_size):
+        if not _REUSE_TEXT_GATHER_BUFFERS_ENABLED:
+            return [torch.empty_like(tensor) for _ in range(world_size)]
+
+        key = (world_size, tuple(tensor.shape), tensor.dtype, tensor.device)
+        buffers = self._text_gather_buffers.get(key)
+        if buffers is None:
+            buffers = [torch.empty_like(tensor) for _ in range(world_size)]
+            self._text_gather_buffers[key] = buffers
+        return buffers
 
     def apply(
         self,
@@ -37,6 +80,7 @@ def apply(
         enable_head_parallel=False,
         img_first=True,
         q_only_img=False,
+        return_split_output=False,
         **kwargs,
     ):
         """
@@ -62,8 +106,16 @@ def apply(
         assert not (use_fp8_comm and use_fp4_comm), "use_fp8_comm and use_fp4_comm can't be enabled at the same time."
 
         use_qkv_fusion = use_tensor_fusion
+        split_qkv_input = _is_split_qkv_input(q)
 
-        if len(q.shape) == 4:
+        if split_qkv_input:
+            if q_only_img:
+                raise NotImplementedError("split QKV input does not support q_only_img yet.")
+            assert _is_split_qkv_input(k) and _is_split_qkv_input(v), "q/k/v must all use split img/txt input."
+            img_q, txt_q = (_to_3d_qkv(tensor) for tensor in q)
+            img_k, txt_k = (_to_3d_qkv(tensor) for tensor in k)
+            img_v, txt_v = (_to_3d_qkv(tensor) for tensor in v)
+        elif len(q.shape) == 4:
             q = q.reshape(-1, q.shape[-2], q.shape[-1])
             k = k.reshape(-1, k.shape[-2], k.shape[-1])
             v = v.reshape(-1, v.shape[-2], v.shape[-1])
@@ -73,7 +125,11 @@ def apply(
         cur_rank = dist.get_rank(seq_p_group)
 
         # 获取序列长度和文本相关的长度
-        if img_first:
+        if split_qkv_input:
+            img_qkv_len = img_q.shape[0]
+            txt_qkv_len = txt_q.shape[0]
+            txt_mask_len = None
+        elif img_first:
             img_qkv_len = slice_qkv_len
             if len(cu_seqlens_qkv) == 3:
                 txt_qkv_len = cu_seqlens_qkv[1] - slice_qkv_len  # 文本查询、键和值的长度
@@ -87,8 +143,12 @@ def apply(
             txt_mask_len = None
 
         # 分别获取 q 和 kv 的头数，支持 GQA（k/v 头数可能少于 q）
-        _, q_heads, hidden_dims = q.shape
-        _, kv_heads, _ = k.shape
+        if split_qkv_input:
+            _, q_heads, hidden_dims = img_q.shape
+            _, kv_heads, _ = img_k.shape
+        else:
+            _, q_heads, hidden_dims = q.shape
+            _, kv_heads, _ = k.shape
         is_gqa = q_heads != kv_heads
         q_shard_heads = q_heads // world_size  # q 每个进程处理的头数
         kv_shard_heads = kv_heads // world_size  # k/v 每个进程处理的头数
@@ -118,7 +178,15 @@ def apply(
             max_seqlen_q = max_seqlen_kv
 
         # 分割图像和文本的查询、键和值
-        if q_only_img:
+        if split_qkv_input:
+            with _profile_range("split_qkv_contiguous"):
+                img_q = _contiguous_if_needed(img_q)
+                img_k = _contiguous_if_needed(img_k)
+                img_v = _contiguous_if_needed(img_v)
+                txt_q = _contiguous_if_needed(txt_q)
+                txt_k = _contiguous_if_needed(txt_k)
+                txt_v = _contiguous_if_needed(txt_v)
+        elif q_only_img:
             # q 只含图像 token，无需分割；仅 k/v 需要拆出图像和文本部分
             img_q = q.contiguous()
             txt_q = None
@@ -148,16 +216,17 @@ def apply(
                 img_k = k[txt_qkv_len:, :, :].contiguous()
                 img_v = v[txt_qkv_len:, :, :].contiguous()
 
-        if use_qkv_fusion:
-            # fusion 路径：q_shard_heads == kv_shard_heads（非 GQA、非 q_only_img 时才走此分支）
-            img_qkv = torch.stack([img_q, img_k, img_v], dim=0).reshape(3, img_qkv_len, world_size, shard_heads, hidden_dims)
-            original_dtype = img_qkv.dtype
-        else:
-            # 非 fusion：q 和 kv 分别 reshape，支持 GQA 下头数不同
-            img_q = img_q.reshape(img_qkv_len, world_size, q_shard_heads, hidden_dims)
-            img_k = img_k.reshape(img_qkv_len, world_size, kv_shard_heads, hidden_dims)
-            img_v = img_v.reshape(img_qkv_len, world_size, kv_shard_heads, hidden_dims)
-            original_dtype = img_q.dtype
+        with _profile_range("seq2head_initial_reshape"):
+            if use_qkv_fusion:
+                # fusion 路径：q_shard_heads == kv_shard_heads（非 GQA、非 q_only_img 时才走此分支）
+                img_qkv = torch.stack([img_q, img_k, img_v], dim=0).reshape(3, img_qkv_len, world_size, shard_heads, hidden_dims)
+                original_dtype = img_qkv.dtype
+            else:
+                # 非 fusion：q 和 kv 分别 reshape，支持 GQA 下头数不同
+                img_q = img_q.reshape(img_qkv_len, world_size, q_shard_heads, hidden_dims)
+                img_k = img_k.reshape(img_qkv_len, world_size, kv_shard_heads, hidden_dims)
+                img_v = img_v.reshape(img_qkv_len, world_size, kv_shard_heads, hidden_dims)
+                original_dtype = img_q.dtype
 
         if enable_head_parallel:
             assert not is_gqa, "GQA（q_heads != kv_heads）暂不支持 enable_head_parallel 模式"
@@ -333,12 +402,13 @@ def apply(
             attn = torch.cat(head_attns, dim=1)
 
         else:
-            if use_qkv_fusion:
-                img_qkv = img_qkv.permute(2, 1, 0, 3, 4).contiguous()  # (world_size, img_qkv_len, 3, shard_heads, hidden_dims)
-            else:
-                img_q = img_q.permute(1, 0, 2, 3).contiguous()  # (world_size, img_qkv_len, q_shard_heads, hidden_dims)
-                img_k = img_k.permute(1, 0, 2, 3).contiguous()  # (world_size, img_qkv_len, kv_shard_heads, hidden_dims)
-                img_v = img_v.permute(1, 0, 2, 3).contiguous()
+            with _profile_range("pre_all_to_all_layout"):
+                if use_qkv_fusion:
+                    img_qkv = img_qkv.permute(2, 1, 0, 3, 4).contiguous()  # (world_size, img_qkv_len, 3, shard_heads, hidden_dims)
+                else:
+                    img_q = img_q.permute(1, 0, 2, 3).contiguous()  # (world_size, img_qkv_len, q_shard_heads, hidden_dims)
+                    img_k = img_k.permute(1, 0, 2, 3).contiguous()  # (world_size, img_qkv_len, kv_shard_heads, hidden_dims)
+                    img_v = img_v.permute(1, 0, 2, 3).contiguous()
 
             # 通信图像的查询、键和值
             if use_qkv_fusion:
@@ -411,12 +481,13 @@ def apply(
                         output_k = dequant_fp4_sage3(output_k_quant.reshape(1, 1, -1, hidden_dims // 2), output_k_scale.reshape(1, 1, -1, hidden_dims // 16))
                         output_v = dequant_fp4_sage3(output_v_quant.reshape(1, 1, -1, hidden_dims // 2), output_v_scale.reshape(1, 1, -1, hidden_dims // 16))
                 else:
-                    output_q = torch.empty_like(img_q)
-                    output_k = torch.empty_like(img_k)
-                    output_v = torch.empty_like(img_v)
-                    dist.all_to_all_single(output_q, img_q, group=seq_p_group)
-                    dist.all_to_all_single(output_k, img_k, group=seq_p_group)
-                    dist.all_to_all_single(output_v, img_v, group=seq_p_group)
+                    with _profile_range("pre_attn_all_to_all_qkv"):
+                        output_q = torch.empty_like(img_q)
+                        output_k = torch.empty_like(img_k)
+                        output_v = torch.empty_like(img_v)
+                        dist.all_to_all_single(output_q, img_q, group=seq_p_group)
+                        dist.all_to_all_single(output_k, img_k, group=seq_p_group)
+                        dist.all_to_all_single(output_v, img_v, group=seq_p_group)
                 # q 与 kv 使用各自对应的 shard_heads 进行 reshape
                 shard_img_q = output_q.reshape(global_img_seqlen, q_shard_heads, hidden_dims)
                 shard_img_k = output_k.reshape(global_img_seqlen, kv_shard_heads, hidden_dims)
@@ -438,17 +509,19 @@ def apply(
                 shard_txt_q = txt_q[:, cur_rank * q_shard_heads : (cur_rank + 1) * q_shard_heads, :]
                 shard_txt_k = txt_k[:, cur_rank * kv_shard_heads : (cur_rank + 1) * kv_shard_heads, :]
                 shard_txt_v = txt_v[:, cur_rank * kv_shard_heads : (cur_rank + 1) * kv_shard_heads, :]
-                if img_first:
-                    q = torch.cat((shard_img_q, shard_txt_q), dim=0)
-                    k = torch.cat((shard_img_k, shard_txt_k), dim=0)
-                    v = torch.cat((shard_img_v, shard_txt_v), dim=0)
-                else:
-                    q = torch.cat((shard_txt_q, shard_img_q), dim=0)
-                    k = torch.cat((shard_txt_k, shard_img_k), dim=0)
-                    v = torch.cat((shard_txt_v, shard_img_v), dim=0)
+                with _profile_range("attn_input_cat"):
+                    if img_first:
+                        q = torch.cat((shard_img_q, shard_txt_q), dim=0)
+                        k = torch.cat((shard_img_k, shard_txt_k), dim=0)
+                        v = torch.cat((shard_img_v, shard_txt_v), dim=0)
+                    else:
+                        q = torch.cat((shard_txt_q, shard_img_q), dim=0)
+                        k = torch.cat((shard_txt_k, shard_img_k), dim=0)
+                        v = torch.cat((shard_txt_v, shard_img_v), dim=0)
 
             # 调用注意力函数计算注意力结果
-            attn = attention_module.apply(q=q, k=k, v=v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, **kwargs)
+            with _profile_range("attention_apply"):
+                attn = attention_module.apply(q=q, k=k, v=v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, **kwargs)
 
         if q_only_img:
             # q 只含图像 token：attn 全部是图像侧结果，无 txt_attn，直接还原通信格式
@@ -462,18 +535,35 @@ def apply(
             txt_attn, img_attn = attn[:txt_qkv_len, :], attn[txt_qkv_len:]
 
         # 通信所有进程的图像注意力结果
+        gathered_txt_attn = None
+        text_gather_work = None
+        if _ASYNC_TEXT_GATHER_ENABLED:
+            with _profile_range("text_all_gather_launch"):
+                gathered_txt_attn = self._get_text_gather_buffers(txt_attn, world_size)
+                text_gather_work = dist.all_gather(gathered_txt_attn, txt_attn, group=seq_p_group, async_op=True)
+
         img_attn = self._reshape_img_attn(img_attn, world_size, shard_seqlen, q_shard_heads, hidden_dims, seq_p_group, use_fp8_comm)
 
         # 收集所有进程的文本注意力结果
-        gathered_txt_attn = [torch.empty_like(txt_attn) for _ in range(world_size)]
-        dist.all_gather(gathered_txt_attn, txt_attn, group=seq_p_group)
-        txt_attn = torch.cat(gathered_txt_attn, dim=1)  # 合并所有进程的文本注意力结果
+        if _ASYNC_TEXT_GATHER_ENABLED:
+            with _profile_range("text_all_gather_wait"):
+                text_gather_work.wait()
+        else:
+            with _profile_range("text_all_gather"):
+                gathered_txt_attn = self._get_text_gather_buffers(txt_attn, world_size)
+                dist.all_gather(gathered_txt_attn, txt_attn, group=seq_p_group)
+        with _profile_range("text_all_gather_cat"):
+            txt_attn = torch.cat(gathered_txt_attn, dim=1)  # 合并所有进程的文本注意力结果
+
+        if return_split_output:
+            return img_attn, txt_attn
 
         # 合并图像和文本的注意力结果
-        if img_first:
-            attn = torch.cat([img_attn, txt_attn], dim=0)
-        else:
-            attn = torch.cat([txt_attn, img_attn], dim=0)
+        with _profile_range("output_img_txt_cat"):
+            if img_first:
+                attn = torch.cat([img_attn, txt_attn], dim=0)
+            else:
+                attn = torch.cat([txt_attn, img_attn], dim=0)
 
         return attn  # 返回最终的注意力结果
 
@@ -485,11 +575,13 @@ def _reshape_img_attn(self, img_attn, world_size, shard_seqlen, shard_heads, hid
             original_dtype = img_attn.dtype
             original_shape = img_attn.shape
             img_attn_quant, attn_scale = quant_fp8_vllm(img_attn.reshape(-1, original_shape[-1]))
-            img_attn_quant = all2all_head2seq(img_attn_quant.reshape(original_shape), group=seq_p_group)
-            attn_scale = all2all_head2seq(attn_scale.reshape(original_shape[0], original_shape[1], 1), group=seq_p_group)
+            with _profile_range("output_all2all_head2seq"):
+                img_attn_quant = all2all_head2seq(img_attn_quant.reshape(original_shape), group=seq_p_group)
+                attn_scale = all2all_head2seq(attn_scale.reshape(original_shape[0], original_shape[1], 1), group=seq_p_group)
             img_attn = dequant_fp8_vllm(img_attn_quant, attn_scale, original_dtype)
         else:
-            img_attn = all2all_head2seq(img_attn, group=seq_p_group)
+            with _profile_range("output_all2all_head2seq"):
+                img_attn = all2all_head2seq(img_attn, group=seq_p_group)
 
         img_attn = img_attn.reshape(shard_seqlen, -1)  # 重塑为 [shard_seqlen, -1] 形状
         return img_attn

lightx2v/common/ops/attn/utils/all2all.py

@@ -1,3 +1,6 @@
+import os
+from contextlib import nullcontext
+
 import torch
 import torch.distributed as dist
 
@@ -9,6 +12,17 @@
     dequant_fp4_sage3 = None
 
 
+def _env_flag(name, default="0"):
+    return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"}
+
+
+_PROFILE_RANGES_ENABLED = _env_flag("LIGHTX2V_ULYSSES_PROFILE_RANGES")
+
+
+def _profile_range(name):
+    return torch.profiler.record_function(f"ulysses::{name}") if _PROFILE_RANGES_ENABLED else nullcontext()
+
+
 def _fp8_all_to_all(input_t, group=None):
     """All-to-all with per-token fp8 compression along the last dim.
 
@@ -118,23 +132,26 @@ def all2all_head2seq(input, group=None):
     shard_seq_len = seq_len // world_size  # 计算每个进程处理的序列长度
 
     # 重塑输入张量以便进行 all-to-all 操作
-    input_t = (
-        input.reshape(world_size, shard_seq_len, shard_heads, hidden_dims)  # 重塑为 [world_size, shard_seq_len, shard_heads, hidden_dims]
-        .transpose(1, 2)  # 转置以便进行 all-to-all 操作
-        .contiguous()  # 确保内存连续
-        .reshape(world_size, shard_heads, shard_seq_len, hidden_dims)  # 再次重塑为 [world_size, shard_heads, shard_seq_len, hidden_dims]
-    )
+    with _profile_range("head2seq_pre_layout"):
+        input_t = (
+            input.reshape(world_size, shard_seq_len, shard_heads, hidden_dims)  # 重塑为 [world_size, shard_seq_len, shard_heads, hidden_dims]
+            .transpose(1, 2)  # 转置以便进行 all-to-all 操作
+            .contiguous()  # 确保内存连续
+            .reshape(world_size, shard_heads, shard_seq_len, hidden_dims)  # 再次重塑为 [world_size, shard_heads, shard_seq_len, hidden_dims]
+        )
 
     # 创建一个与输入张量相同形状的输出张量
     output = torch.empty_like(input_t)
 
     # 执行 all-to-all 操作，将输入张量的内容分发到所有进程
-    dist.all_to_all_single(output, input_t, group=group)
+    with _profile_range("head2seq_all_to_all"):
+        dist.all_to_all_single(output, input_t, group=group)
 
-    # 重塑输出张量为 [heads, shard_seq_len, hidden_dims] 形状
-    output = output.reshape(heads, shard_seq_len, hidden_dims)
+    with _profile_range("head2seq_post_layout"):
+        # 重塑输出张量为 [heads, shard_seq_len, hidden_dims] 形状
+        output = output.reshape(heads, shard_seq_len, hidden_dims)
 
-    # 转置输出张量并重塑为 [shard_seq_len, heads, hidden_dims] 形状
-    output = output.transpose(0, 1).contiguous().reshape(shard_seq_len, heads, hidden_dims)
+        # 转置输出张量并重塑为 [shard_seq_len, heads, hidden_dims] 形状
+        output = output.transpose(0, 1).contiguous().reshape(shard_seq_len, heads, hidden_dims)
 
     return output  # 返回转换后的输出张量