From a0397e898e084dd02052741a17fc2dcd7362729c Mon Sep 17 00:00:00 2001 From: zhenggf Date: Thu, 25 Jun 2026 10:28:32 +0800 Subject: [PATCH 1/6] Optimize Hunyuan DiT non-attention compile (cherry picked from commit 8f06fb6c7e0859f432a329a84f8d5d8e3a386ad1) --- .../hunyuan_video/infer/transformer_infer.py | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index 87341b4f6..c5b14d922 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -1,8 +1,10 @@ +import os from typing import Tuple import torch import torch.nn.functional as F from einops import rearrange +from loguru import logger try: from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace @@ -38,6 +40,10 @@ def apply_gate(x, gate=None, tanh=False): return x * gate.unsqueeze(1) +def _env_flag(name, default="0"): + return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"} + + def apply_hunyuan_rope_with_flashinfer( xq: torch.Tensor, xk: torch.Tensor, @@ -119,6 +125,21 @@ def __init__(self, config): self.apply_rope_func = apply_hunyuan_rope_with_flashinfer else: self.apply_rope_func = apply_hunyuan_rope_with_torch + self.compile_non_attn = _env_flag("LIGHTX2V_COMPILE_DIT_NON_ATTN") + self.compile_before_attn = _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN") + self.compile_non_attn_mode = os.getenv("LIGHTX2V_COMPILE_DIT_MODE", "reduce-overhead") + self._compiled_non_attn = {} + self._compile_non_attn_failed = set() + if self.compile_non_attn or self.compile_before_attn: + try: + torch._dynamo.config.suppress_errors = True + except Exception as exc: + logger.warning(f"[Compile] Unable to enable Dynamo suppress_errors: {exc}") + logger.info( + "[Compile] Hunyuan DiT branch compile: " + f"after_attn={self.compile_non_attn}, before_attn={self.compile_before_attn}, " + f"mode={self.compile_non_attn_mode}" + ) def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -155,6 +176,15 @@ def infer_double_block(self, weights, infer_module_out): @torch.no_grad() def _infer_img_branch_before_attn(self, weights, infer_module_out): + return self._run_non_attn_branch( + "img_before_attn", + self._infer_img_branch_before_attn_eager, + weights, + infer_module_out, + compile_enabled=self.compile_before_attn, + ) + + def _infer_img_branch_before_attn_eager(self, weights, infer_module_out): ( img_mod1_shift, img_mod1_scale, @@ -188,6 +218,15 @@ def _infer_img_branch_before_attn(self, weights, infer_module_out): @torch.no_grad() def _infer_txt_branch_before_attn(self, weights, infer_module_out): + return self._run_non_attn_branch( + "txt_before_attn", + self._infer_txt_branch_before_attn_eager, + weights, + infer_module_out, + compile_enabled=self.compile_before_attn, + ) + + def _infer_txt_branch_before_attn_eager(self, weights, infer_module_out): ( txt_mod1_shift, txt_mod1_scale, @@ -246,8 +285,36 @@ def _infer_attn(self, weights, img_q, img_k, img_v, txt_q, txt_k, txt_v): img_attn, txt_attn = attn_out[:img_seqlen], attn_out[img_seqlen:] return img_attn, txt_attn + def _run_non_attn_branch(self, graph_name, eager_fn, *args, compile_enabled=None): + if compile_enabled is None: + compile_enabled = self.compile_non_attn + if not compile_enabled or graph_name in self._compile_non_attn_failed: + return eager_fn(*args) + + compiled_fn = self._compiled_non_attn.get(graph_name) + if compiled_fn is None: + try: + compiled_fn = torch.compile(eager_fn, fullgraph=False, dynamic=False, mode=self.compile_non_attn_mode) + self._compiled_non_attn[graph_name] = compiled_fn + logger.info(f"[Compile] Created compiled wrapper for {graph_name}") + except Exception as exc: + self._compile_non_attn_failed.add(graph_name) + logger.warning(f"[Compile] Failed to create compiled wrapper for {graph_name}, fallback to eager: {exc}") + return eager_fn(*args) + + try: + return compiled_fn(*args) + except Exception as exc: + self._compile_non_attn_failed.add(graph_name) + self._compiled_non_attn.pop(graph_name, None) + logger.warning(f"[Compile] Runtime failure in {graph_name}, disabling this graph and falling back to eager: {exc}") + return eager_fn(*args) + @torch.no_grad() def _infer_img_branch_after_attn(self, weights, img_attn, img, img_branch_out): + return self._run_non_attn_branch("img_after_attn", self._infer_img_branch_after_attn_eager, weights, img_attn, img, img_branch_out) + + def _infer_img_branch_after_attn_eager(self, weights, img_attn, img, img_branch_out): img = img + apply_gate(weights.img_branch.img_attn_proj.apply(img_attn).unsqueeze(0), gate=img_branch_out.img_mod1_gate) out = weights.img_branch.img_mlp_fc1.apply( self.modulate_func(weights.img_branch.img_norm2.apply(img.squeeze(0)), scale=img_branch_out.img_mod2_scale, shift=img_branch_out.img_mod2_shift).squeeze(0) @@ -258,6 +325,9 @@ def _infer_img_branch_after_attn(self, weights, img_attn, img, img_branch_out): @torch.no_grad() def _infer_txt_branch_after_attn(self, weights, txt_attn, txt, txt_branch_out): + return self._run_non_attn_branch("txt_after_attn", self._infer_txt_branch_after_attn_eager, weights, txt_attn, txt, txt_branch_out) + + def _infer_txt_branch_after_attn_eager(self, weights, txt_attn, txt, txt_branch_out): txt = txt + apply_gate(weights.txt_branch.txt_attn_proj.apply(txt_attn).unsqueeze(0), gate=txt_branch_out.txt_mod1_gate) out = weights.txt_branch.txt_mlp_fc1.apply( self.modulate_func(weights.txt_branch.txt_norm2.apply(txt.squeeze(0)), scale=txt_branch_out.txt_mod2_scale, shift=txt_branch_out.txt_mod2_shift).squeeze(0) From e7a3846e9bb14bf2c70e3385be96f051dceacbbd Mon Sep 17 00:00:00 2001 From: zhenggf Date: Mon, 29 Jun 2026 09:32:09 +0800 Subject: [PATCH 2/6] ulysses: add split qkv and async text gather Support split image/text QKV inputs, optional split attention outputs, async text all_gather, and profiler ranges for Ulysses sequence-parallel attention. (cherry picked from commit 8bb7c3e1784140a8f6d372fe429b468e3a502b8b) --- lightx2v/common/ops/attn/ulysses_attn.py | 184 ++++++++++++++++------ lightx2v/common/ops/attn/utils/all2all.py | 39 +++-- 2 files changed, 166 insertions(+), 57 deletions(-) diff --git a/lightx2v/common/ops/attn/ulysses_attn.py b/lightx2v/common/ops/attn/ulysses_attn.py index 93fff1eac..84c73ec42 100755 --- a/lightx2v/common/ops/attn/ulysses_attn.py +++ b/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 diff --git a/lightx2v/common/ops/attn/utils/all2all.py b/lightx2v/common/ops/attn/utils/all2all.py index 9bc717048..ef3e1abd9 100644 --- a/lightx2v/common/ops/attn/utils/all2all.py +++ b/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 # 返回转换后的输出张量 From 0430c86d8b1b9140158a59e0c0dc6461bbe07020 Mon Sep 17 00:00:00 2001 From: zhenggf Date: Mon, 29 Jun 2026 09:32:09 +0800 Subject: [PATCH 3/6] hunyuan: wire shared qkv quantization and split attention Reuse dynamic activation quantization across consecutive Q/K/V projections and route split image/text tensors through the Ulysses attention path when enabled. (cherry picked from commit 61c5df5c20106254d5294b910cdf3d1780970a97) --- .../hunyuan_video/infer/transformer_infer.py | 91 +++++++++++++++---- 1 file changed, 74 insertions(+), 17 deletions(-) diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index c5b14d922..32f48ad1e 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -111,11 +111,17 @@ def __init__(self, config): self.seq_p_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False) self.seq_p_fp4_comm = self.config["parallel"].get("seq_p_fp4_comm", False) self.enable_head_parallel = self.config["parallel"].get("seq_p_head_parallel", False) + self.seq_p_tensor_fusion = self.config["parallel"].get("seq_p_tensor_fusion", False) + self.seq_p_split_qkv_input = _env_flag("LIGHTX2V_SEQ_P_SPLIT_QKV_INPUT") + self.seq_p_split_attn_output = _env_flag("LIGHTX2V_SEQ_P_SPLIT_ATTN_OUTPUT") else: self.seq_p_group = None self.seq_p_fp8_comm = False self.seq_p_fp4_comm = False self.enable_head_parallel = False + self.seq_p_tensor_fusion = False + self.seq_p_split_qkv_input = False + self.seq_p_split_attn_output = False self.infer_func = self.infer_without_offload if self.config.get("modulate_type", "triton") == "triton": self.modulate_func = fuse_scale_shift_kernel @@ -128,8 +134,15 @@ def __init__(self, config): self.compile_non_attn = _env_flag("LIGHTX2V_COMPILE_DIT_NON_ATTN") self.compile_before_attn = _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN") self.compile_non_attn_mode = os.getenv("LIGHTX2V_COMPILE_DIT_MODE", "reduce-overhead") + self.share_qkv_act_quant = _env_flag("LIGHTX2V_INT8_SHARE_QKV_ACT_QUANT") self._compiled_non_attn = {} self._compile_non_attn_failed = set() + if self.share_qkv_act_quant: + logger.info("[Quant] Reusing dynamic activation quantization for consecutive Q/K/V projections") + if self.seq_p_split_qkv_input: + logger.info("[Ulysses] Passing split img/txt QKV tensors to avoid pre-attention concat/slice copies") + if self.seq_p_split_attn_output: + logger.info("[Ulysses] Returning split img/txt attention outputs to avoid post-attention concat/slice copies") if self.compile_non_attn or self.compile_before_attn: try: torch._dynamo.config.suppress_errors = True @@ -174,6 +187,22 @@ def infer_double_block(self, weights, infer_module_out): txt = self._infer_txt_branch_after_attn(weights, txt_attn, infer_module_out.txt, txt_branch_out) return img, txt + def _apply_qkv(self, q_weight, k_weight, v_weight, hidden_states): + use_shared_quant = ( + self.share_qkv_act_quant + and hasattr(q_weight, "prepare_quantized_input") + and hasattr(q_weight, "apply_quantized_input") + and not any(getattr(weight, "use_bf16_fallback", False) for weight in (q_weight, k_weight, v_weight)) + ) + if use_shared_quant: + quantized_input = q_weight.prepare_quantized_input(hidden_states) + return ( + q_weight.apply_quantized_input(hidden_states, quantized_input), + k_weight.apply_quantized_input(hidden_states, quantized_input), + v_weight.apply_quantized_input(hidden_states, quantized_input), + ) + return q_weight.apply(hidden_states), k_weight.apply(hidden_states), v_weight.apply(hidden_states) + @torch.no_grad() def _infer_img_branch_before_attn(self, weights, infer_module_out): return self._run_non_attn_branch( @@ -195,9 +224,12 @@ def _infer_img_branch_before_attn_eager(self, weights, infer_module_out): ) = weights.img_branch.img_mod.apply(infer_module_out.vec).chunk(6, dim=-1) img_modulated = weights.img_branch.img_norm1.apply(infer_module_out.img.squeeze(0)) img_modulated = self.modulate_func(img_modulated, scale=img_mod1_scale, shift=img_mod1_shift).squeeze(0) - img_q = weights.img_branch.img_attn_q.apply(img_modulated) - img_k = weights.img_branch.img_attn_k.apply(img_modulated) - img_v = weights.img_branch.img_attn_v.apply(img_modulated) + img_q, img_k, img_v = self._apply_qkv( + weights.img_branch.img_attn_q, + weights.img_branch.img_attn_k, + weights.img_branch.img_attn_v, + img_modulated, + ) img_q = rearrange(img_q, "L (H D) -> L H D", H=self.heads_num) img_k = rearrange(img_k, "L (H D) -> L H D", H=self.heads_num) img_v = rearrange(img_v, "L (H D) -> L H D", H=self.heads_num) @@ -237,9 +269,12 @@ def _infer_txt_branch_before_attn_eager(self, weights, infer_module_out): ) = weights.txt_branch.txt_mod.apply(infer_module_out.vec).chunk(6, dim=-1) txt_modulated = weights.txt_branch.txt_norm1.apply(infer_module_out.txt.squeeze(0)) txt_modulated = self.modulate_func(txt_modulated, scale=txt_mod1_scale, shift=txt_mod1_shift).squeeze(0) - txt_q = weights.txt_branch.txt_attn_q.apply(txt_modulated) - txt_k = weights.txt_branch.txt_attn_k.apply(txt_modulated) - txt_v = weights.txt_branch.txt_attn_v.apply(txt_modulated) + txt_q, txt_k, txt_v = self._apply_qkv( + weights.txt_branch.txt_attn_q, + weights.txt_branch.txt_attn_k, + weights.txt_branch.txt_attn_v, + txt_modulated, + ) txt_q = rearrange(txt_q, "L (H D) -> L H D", H=self.heads_num) txt_k = rearrange(txt_k, "L (H D) -> L H D", H=self.heads_num) txt_v = rearrange(txt_v, "L (H D) -> L H D", H=self.heads_num) @@ -260,29 +295,51 @@ def _infer_txt_branch_before_attn_eager(self, weights, infer_module_out): @torch.no_grad() def _infer_attn(self, weights, img_q, img_k, img_v, txt_q, txt_k, txt_v): img_seqlen = img_q.shape[1] - query = torch.cat([img_q, txt_q], dim=1) - key = torch.cat([img_k, txt_k], dim=1) - value = torch.cat([img_v, txt_v], dim=1) - seqlen = query.shape[1] + txt_seqlen = txt_q.shape[1] + seqlen = img_seqlen + txt_seqlen cu_seqlens_qkv = torch.tensor([0, seqlen], dtype=torch.int32, device="cpu") - if self.config["seq_parallel"]: + if self.config["seq_parallel"] and self.seq_p_split_qkv_input: attn_out = weights.self_attention_parallel.apply( - q=query, - k=key, - v=value, + q=(img_q, txt_q), + k=(img_k, txt_k), + v=(img_v, txt_v), slice_qkv_len=img_seqlen, cu_seqlens_qkv=cu_seqlens_qkv, attention_module=weights.self_attention, seq_p_group=self.seq_p_group, use_fp8_comm=self.seq_p_fp8_comm, use_fp4_comm=self.seq_p_fp4_comm, + use_tensor_fusion=self.seq_p_tensor_fusion, enable_head_parallel=self.enable_head_parallel, + return_split_output=self.seq_p_split_attn_output, ) else: - attn_out = weights.self_attention.apply(q=query, k=key, v=value, cu_seqlens_q=cu_seqlens_qkv, cu_seqlens_kv=cu_seqlens_qkv, max_seqlen_q=seqlen, max_seqlen_kv=seqlen) - - img_attn, txt_attn = attn_out[:img_seqlen], attn_out[img_seqlen:] + query = torch.cat([img_q, txt_q], dim=1) + key = torch.cat([img_k, txt_k], dim=1) + value = torch.cat([img_v, txt_v], dim=1) + if self.config["seq_parallel"]: + attn_out = weights.self_attention_parallel.apply( + q=query, + k=key, + v=value, + slice_qkv_len=img_seqlen, + cu_seqlens_qkv=cu_seqlens_qkv, + attention_module=weights.self_attention, + seq_p_group=self.seq_p_group, + use_fp8_comm=self.seq_p_fp8_comm, + use_fp4_comm=self.seq_p_fp4_comm, + use_tensor_fusion=self.seq_p_tensor_fusion, + enable_head_parallel=self.enable_head_parallel, + return_split_output=self.seq_p_split_attn_output, + ) + else: + attn_out = weights.self_attention.apply(q=query, k=key, v=value, cu_seqlens_q=cu_seqlens_qkv, cu_seqlens_kv=cu_seqlens_qkv, max_seqlen_q=seqlen, max_seqlen_kv=seqlen) + + if isinstance(attn_out, (tuple, list)): + img_attn, txt_attn = attn_out + else: + img_attn, txt_attn = attn_out[:img_seqlen], attn_out[img_seqlen:] return img_attn, txt_attn def _run_non_attn_branch(self, graph_name, eager_fn, *args, compile_enabled=None): From b7d6fa6b8272c8c07353f061a2c1dd9a6aaff4fc Mon Sep 17 00:00:00 2001 From: zhenggf Date: Tue, 30 Jun 2026 14:42:19 +0800 Subject: [PATCH 4/6] fix: harden Hunyuan DiT optimization switches --- lightx2v/common/ops/attn/ulysses_attn.py | 22 ++++++++++++++----- .../hunyuan_video/infer/transformer_infer.py | 11 ++++++++-- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/lightx2v/common/ops/attn/ulysses_attn.py b/lightx2v/common/ops/attn/ulysses_attn.py index 84c73ec42..caa16a51e 100755 --- a/lightx2v/common/ops/attn/ulysses_attn.py +++ b/lightx2v/common/ops/attn/ulysses_attn.py @@ -16,9 +16,17 @@ def _env_flag(name, default="0"): return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"} +def _env_int(name, default): + try: + return int(os.getenv(name, str(default))) + except (TypeError, ValueError): + return default + + _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") +_TEXT_GATHER_BUFFER_CACHE_MAX = max(1, _env_int("LIGHTX2V_ULYSSES_TEXT_GATHER_BUFFER_CACHE_MAX", 16)) def _profile_range(name): @@ -61,6 +69,8 @@ def _get_text_gather_buffers(self, tensor, world_size): key = (world_size, tuple(tensor.shape), tensor.dtype, tensor.device) buffers = self._text_gather_buffers.get(key) if buffers is None: + if len(self._text_gather_buffers) >= _TEXT_GATHER_BUFFER_CACHE_MAX: + self._text_gather_buffers.clear() buffers = [torch.empty_like(tensor) for _ in range(world_size)] self._text_gather_buffers[key] = buffers return buffers @@ -128,7 +138,7 @@ def apply( if split_qkv_input: img_qkv_len = img_q.shape[0] txt_qkv_len = txt_q.shape[0] - txt_mask_len = None + txt_mask_len = cu_seqlens_qkv[2] - img_qkv_len if img_first and len(cu_seqlens_qkv) == 3 else None elif img_first: img_qkv_len = slice_qkv_len if len(cu_seqlens_qkv) == 3: @@ -542,13 +552,15 @@ def apply( 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) - - # 收集所有进程的文本注意力结果 + # Finish the async gather before launching any later collective on the same process group. if _ASYNC_TEXT_GATHER_ENABLED: with _profile_range("text_all_gather_wait"): text_gather_work.wait() - else: + + img_attn = self._reshape_img_attn(img_attn, world_size, shard_seqlen, q_shard_heads, hidden_dims, seq_p_group, use_fp8_comm) + + # Gather text attention synchronously when async launch is disabled. + if not _ASYNC_TEXT_GATHER_ENABLED: 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) diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index 32f48ad1e..0ff6eb891 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -132,7 +132,8 @@ def __init__(self, config): else: self.apply_rope_func = apply_hunyuan_rope_with_torch self.compile_non_attn = _env_flag("LIGHTX2V_COMPILE_DIT_NON_ATTN") - self.compile_before_attn = _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN") + self.compile_before_attn_requested = _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN") + self.compile_before_attn = self.compile_before_attn_requested and _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN_UNSAFE") self.compile_non_attn_mode = os.getenv("LIGHTX2V_COMPILE_DIT_MODE", "reduce-overhead") self.share_qkv_act_quant = _env_flag("LIGHTX2V_INT8_SHARE_QKV_ACT_QUANT") self._compiled_non_attn = {} @@ -143,6 +144,12 @@ def __init__(self, config): logger.info("[Ulysses] Passing split img/txt QKV tensors to avoid pre-attention concat/slice copies") if self.seq_p_split_attn_output: logger.info("[Ulysses] Returning split img/txt attention outputs to avoid post-attention concat/slice copies") + if self.compile_before_attn_requested and not self.compile_before_attn: + logger.warning( + "[Compile] LIGHTX2V_COMPILE_DIT_BEFORE_ATTN is ignored unless " + "LIGHTX2V_COMPILE_DIT_BEFORE_ATTN_UNSAFE=1 is also set; before-attn graphs carry " + "block-specific weight objects and can grow Dynamo caches quickly." + ) if self.compile_non_attn or self.compile_before_attn: try: torch._dynamo.config.suppress_errors = True @@ -191,7 +198,7 @@ def _apply_qkv(self, q_weight, k_weight, v_weight, hidden_states): use_shared_quant = ( self.share_qkv_act_quant and hasattr(q_weight, "prepare_quantized_input") - and hasattr(q_weight, "apply_quantized_input") + and all(hasattr(weight, "apply_quantized_input") for weight in (q_weight, k_weight, v_weight)) and not any(getattr(weight, "use_bf16_fallback", False) for weight in (q_weight, k_weight, v_weight)) ) if use_shared_quant: From b215bb3044cf9d3515a2c4072ca128d0704e5a9f Mon Sep 17 00:00:00 2001 From: zhenggf Date: Wed, 1 Jul 2026 14:59:33 +0800 Subject: [PATCH 5/6] style: format Hunyuan DiT optimization changes --- .../networks/hunyuan_video/infer/transformer_infer.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index 0ff6eb891..cb505c03b 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -155,11 +155,7 @@ def __init__(self, config): torch._dynamo.config.suppress_errors = True except Exception as exc: logger.warning(f"[Compile] Unable to enable Dynamo suppress_errors: {exc}") - logger.info( - "[Compile] Hunyuan DiT branch compile: " - f"after_attn={self.compile_non_attn}, before_attn={self.compile_before_attn}, " - f"mode={self.compile_non_attn_mode}" - ) + logger.info(f"[Compile] Hunyuan DiT branch compile: after_attn={self.compile_non_attn}, before_attn={self.compile_before_attn}, mode={self.compile_non_attn_mode}") def set_scheduler(self, scheduler): self.scheduler = scheduler From 59ddfb8e437a7bba3705c67da4246f68e3a29648 Mon Sep 17 00:00:00 2001 From: zhenggf Date: Wed, 1 Jul 2026 17:59:48 +0800 Subject: [PATCH 6/6] refactor: remove env switches from Hunyuan DiT optimizations --- lightx2v/common/ops/attn/ulysses_attn.py | 152 +++++++----------- lightx2v/common/ops/attn/utils/all2all.py | 39 ++--- .../hunyuan_video/infer/transformer_infer.py | 44 ++--- 3 files changed, 96 insertions(+), 139 deletions(-) diff --git a/lightx2v/common/ops/attn/ulysses_attn.py b/lightx2v/common/ops/attn/ulysses_attn.py index caa16a51e..a0daf36e1 100755 --- a/lightx2v/common/ops/attn/ulysses_attn.py +++ b/lightx2v/common/ops/attn/ulysses_attn.py @@ -1,6 +1,3 @@ -import os -from contextlib import nullcontext - import torch import torch.distributed as dist from loguru import logger @@ -12,27 +9,6 @@ from .utils.all2all import all2all_head2seq -def _env_flag(name, default="0"): - return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"} - - -def _env_int(name, default): - try: - return int(os.getenv(name, str(default))) - except (TypeError, ValueError): - return default - - -_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") -_TEXT_GATHER_BUFFER_CACHE_MAX = max(1, _env_int("LIGHTX2V_ULYSSES_TEXT_GATHER_BUFFER_CACHE_MAX", 16)) - - -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)) @@ -62,14 +38,14 @@ 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: + def _get_text_gather_buffers(self, tensor, world_size, reuse_buffers=False, cache_max=16): + if not reuse_buffers: 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: - if len(self._text_gather_buffers) >= _TEXT_GATHER_BUFFER_CACHE_MAX: + if len(self._text_gather_buffers) >= max(1, cache_max): self._text_gather_buffers.clear() buffers = [torch.empty_like(tensor) for _ in range(world_size)] self._text_gather_buffers[key] = buffers @@ -91,6 +67,9 @@ def apply( img_first=True, q_only_img=False, return_split_output=False, + async_text_gather=False, + reuse_text_gather_buffers=False, + text_gather_buffer_cache_max=16, **kwargs, ): """ @@ -189,13 +168,12 @@ def apply( # 分割图像和文本的查询、键和值 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) + 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() @@ -226,17 +204,16 @@ def apply( img_k = k[txt_qkv_len:, :, :].contiguous() img_v = v[txt_qkv_len:, :, :].contiguous() - 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 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 模式" @@ -412,13 +389,12 @@ def apply( attn = torch.cat(head_attns, dim=1) else: - 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: + 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: @@ -491,13 +467,12 @@ 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: - 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) + 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) @@ -519,19 +494,17 @@ 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, :] - 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) + 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("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) + 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，直接还原通信格式 @@ -547,35 +520,30 @@ def apply( # 通信所有进程的图像注意力结果 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) + if async_text_gather: + gathered_txt_attn = self._get_text_gather_buffers(txt_attn, world_size, reuse_text_gather_buffers, text_gather_buffer_cache_max) + text_gather_work = dist.all_gather(gathered_txt_attn, txt_attn, group=seq_p_group, async_op=True) # Finish the async gather before launching any later collective on the same process group. - if _ASYNC_TEXT_GATHER_ENABLED: - with _profile_range("text_all_gather_wait"): - text_gather_work.wait() + if async_text_gather: + text_gather_work.wait() img_attn = self._reshape_img_attn(img_attn, world_size, shard_seqlen, q_shard_heads, hidden_dims, seq_p_group, use_fp8_comm) # Gather text attention synchronously when async launch is disabled. - if not _ASYNC_TEXT_GATHER_ENABLED: - 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 not async_text_gather: + gathered_txt_attn = self._get_text_gather_buffers(txt_attn, world_size, reuse_text_gather_buffers, text_gather_buffer_cache_max) + dist.all_gather(gathered_txt_attn, txt_attn, group=seq_p_group) + txt_attn = torch.cat(gathered_txt_attn, dim=1) # 合并所有进程的文本注意力结果 if return_split_output: return img_attn, txt_attn # 合并图像和文本的注意力结果 - 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) + if img_first: + attn = torch.cat([img_attn, txt_attn], dim=0) + else: + attn = torch.cat([txt_attn, img_attn], dim=0) return attn # 返回最终的注意力结果 @@ -587,13 +555,11 @@ 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])) - 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_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: - with _profile_range("output_all2all_head2seq"): - img_attn = all2all_head2seq(img_attn, group=seq_p_group) + img_attn = all2all_head2seq(img_attn, group=seq_p_group) img_attn = img_attn.reshape(shard_seqlen, -1) # 重塑为 [shard_seqlen, -1] 形状 return img_attn diff --git a/lightx2v/common/ops/attn/utils/all2all.py b/lightx2v/common/ops/attn/utils/all2all.py index ef3e1abd9..9bc717048 100644 --- a/lightx2v/common/ops/attn/utils/all2all.py +++ b/lightx2v/common/ops/attn/utils/all2all.py @@ -1,6 +1,3 @@ -import os -from contextlib import nullcontext - import torch import torch.distributed as dist @@ -12,17 +9,6 @@ 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. @@ -132,26 +118,23 @@ def all2all_head2seq(input, group=None): shard_seq_len = seq_len // world_size # 计算每个进程处理的序列长度 # 重塑输入张量以便进行 all-to-all 操作 - 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] - ) + 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 操作，将输入张量的内容分发到所有进程 - with _profile_range("head2seq_all_to_all"): - dist.all_to_all_single(output, input_t, group=group) + dist.all_to_all_single(output, input_t, group=group) - with _profile_range("head2seq_post_layout"): - # 重塑输出张量为 [heads, shard_seq_len, hidden_dims] 形状 - output = output.reshape(heads, shard_seq_len, hidden_dims) + # 重塑输出张量为 [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 # 返回转换后的输出张量 diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index cb505c03b..cd1c2c92a 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -1,4 +1,3 @@ -import os from typing import Tuple import torch @@ -40,10 +39,6 @@ def apply_gate(x, gate=None, tanh=False): return x * gate.unsqueeze(1) -def _env_flag(name, default="0"): - return os.getenv(name, default).strip().lower() in {"1", "true", "yes", "on"} - - def apply_hunyuan_rope_with_flashinfer( xq: torch.Tensor, xk: torch.Tensor, @@ -106,14 +101,18 @@ def __init__(self, config): self.config = config self.double_blocks_num = config["mm_double_blocks_depth"] self.heads_num = config["heads_num"] + parallel_config = self.config.get("parallel", {}) if self.config["seq_parallel"]: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") - self.seq_p_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False) - self.seq_p_fp4_comm = self.config["parallel"].get("seq_p_fp4_comm", False) - self.enable_head_parallel = self.config["parallel"].get("seq_p_head_parallel", False) - self.seq_p_tensor_fusion = self.config["parallel"].get("seq_p_tensor_fusion", False) - self.seq_p_split_qkv_input = _env_flag("LIGHTX2V_SEQ_P_SPLIT_QKV_INPUT") - self.seq_p_split_attn_output = _env_flag("LIGHTX2V_SEQ_P_SPLIT_ATTN_OUTPUT") + self.seq_p_fp8_comm = parallel_config.get("seq_p_fp8_comm", False) + self.seq_p_fp4_comm = parallel_config.get("seq_p_fp4_comm", False) + self.enable_head_parallel = parallel_config.get("seq_p_head_parallel", False) + self.seq_p_tensor_fusion = parallel_config.get("seq_p_tensor_fusion", False) + self.seq_p_split_qkv_input = parallel_config.get("seq_p_split_qkv_input", False) + self.seq_p_split_attn_output = parallel_config.get("seq_p_split_attn_output", False) + self.seq_p_async_text_gather = parallel_config.get("seq_p_async_text_gather", False) + self.seq_p_reuse_text_gather_buffers = parallel_config.get("seq_p_reuse_text_gather_buffers", False) + self.seq_p_text_gather_buffer_cache_max = parallel_config.get("seq_p_text_gather_buffer_cache_max", 16) else: self.seq_p_group = None self.seq_p_fp8_comm = False @@ -122,6 +121,9 @@ def __init__(self, config): self.seq_p_tensor_fusion = False self.seq_p_split_qkv_input = False self.seq_p_split_attn_output = False + self.seq_p_async_text_gather = False + self.seq_p_reuse_text_gather_buffers = False + self.seq_p_text_gather_buffer_cache_max = 16 self.infer_func = self.infer_without_offload if self.config.get("modulate_type", "triton") == "triton": self.modulate_func = fuse_scale_shift_kernel @@ -131,11 +133,11 @@ def __init__(self, config): self.apply_rope_func = apply_hunyuan_rope_with_flashinfer else: self.apply_rope_func = apply_hunyuan_rope_with_torch - self.compile_non_attn = _env_flag("LIGHTX2V_COMPILE_DIT_NON_ATTN") - self.compile_before_attn_requested = _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN") - self.compile_before_attn = self.compile_before_attn_requested and _env_flag("LIGHTX2V_COMPILE_DIT_BEFORE_ATTN_UNSAFE") - self.compile_non_attn_mode = os.getenv("LIGHTX2V_COMPILE_DIT_MODE", "reduce-overhead") - self.share_qkv_act_quant = _env_flag("LIGHTX2V_INT8_SHARE_QKV_ACT_QUANT") + self.compile_non_attn = self.config.get("compile_dit_non_attn", False) + self.compile_before_attn_requested = self.config.get("compile_dit_before_attn", False) + self.compile_before_attn = self.compile_before_attn_requested and self.config.get("compile_dit_before_attn_unsafe", False) + self.compile_non_attn_mode = self.config.get("compile_dit_mode", "reduce-overhead") + self.share_qkv_act_quant = self.config.get("share_qkv_act_quant", False) self._compiled_non_attn = {} self._compile_non_attn_failed = set() if self.share_qkv_act_quant: @@ -146,8 +148,8 @@ def __init__(self, config): logger.info("[Ulysses] Returning split img/txt attention outputs to avoid post-attention concat/slice copies") if self.compile_before_attn_requested and not self.compile_before_attn: logger.warning( - "[Compile] LIGHTX2V_COMPILE_DIT_BEFORE_ATTN is ignored unless " - "LIGHTX2V_COMPILE_DIT_BEFORE_ATTN_UNSAFE=1 is also set; before-attn graphs carry " + "[Compile] compile_dit_before_attn is ignored unless " + "compile_dit_before_attn_unsafe is also set; before-attn graphs carry " "block-specific weight objects and can grow Dynamo caches quickly." ) if self.compile_non_attn or self.compile_before_attn: @@ -316,6 +318,9 @@ def _infer_attn(self, weights, img_q, img_k, img_v, txt_q, txt_k, txt_v): use_tensor_fusion=self.seq_p_tensor_fusion, enable_head_parallel=self.enable_head_parallel, return_split_output=self.seq_p_split_attn_output, + async_text_gather=self.seq_p_async_text_gather, + reuse_text_gather_buffers=self.seq_p_reuse_text_gather_buffers, + text_gather_buffer_cache_max=self.seq_p_text_gather_buffer_cache_max, ) else: query = torch.cat([img_q, txt_q], dim=1) @@ -335,6 +340,9 @@ def _infer_attn(self, weights, img_q, img_k, img_v, txt_q, txt_k, txt_v): use_tensor_fusion=self.seq_p_tensor_fusion, enable_head_parallel=self.enable_head_parallel, return_split_output=self.seq_p_split_attn_output, + async_text_gather=self.seq_p_async_text_gather, + reuse_text_gather_buffers=self.seq_p_reuse_text_gather_buffers, + text_gather_buffer_cache_max=self.seq_p_text_gather_buffer_cache_max, ) else: attn_out = weights.self_attention.apply(q=query, k=key, v=value, cu_seqlens_q=cu_seqlens_qkv, cu_seqlens_kv=cu_seqlens_qkv, max_seqlen_q=seqlen, max_seqlen_kv=seqlen)