From f7c543dfe44959ea7dba9d0e16414cf4f65423a7 Mon Sep 17 00:00:00 2001 From: Chernobyllight <62710385+Chernobyllight@users.noreply.github.com> Date: Fri, 21 Aug 2026 08:56:38 +0000 Subject: [PATCH] feat(hunyuan-image3): optimize four-GPU hybrid inference Add a production-ready HunyuanImage 3.0 path for the classic four-GPU topology. The AR phase runs TP4 with native-GQA paged FlashAttention 3 decode, full-decode CUDA Graph capture, and graph-aware vLLM custom all-reduce. The denoising phase remains TP2+SP2 and uses native-GQA FlashAttention 3 over NCCL collectives. Include portable T2I and TI2I configs and launch scripts, strict configuration validation, persistent paged KV-cache management, and deterministic distributed cleanup. The implementation intentionally excludes benchmark-only timing instrumentation. --- ...se_fa3_tp2_sp2_multi_micro_flashinfer.json | 79 ++++ ...se_fa3_tp2_sp2_multi_micro_flashinfer.json | 79 ++++ lightx2v/common/ops/attn/paged_flash_attn.py | 171 ++++++++ lightx2v/infer.py | 32 +- .../models/networks/hunyuan_image3/config.py | 38 ++ .../hunyuan_image3/custom_all_reduce.py | 319 +++++++++++++++ .../networks/hunyuan_image3/infer/kv_cache.py | 105 ++++- .../hunyuan_image3/infer/transformer_infer.py | 228 +++++++++-- .../models/networks/hunyuan_image3/model.py | 28 ++ .../networks/hunyuan_image3/parallel.py | 53 ++- .../hunyuan_image3/weights/hybrid_tp.py | 6 +- .../runners/hunyuan_image3/cuda_graph.py | 384 ++++++++++++++++++ .../hunyuan_image3/hunyuan_image3_runner.py | 57 ++- ...oise_fa3_tp2_sp2_multi_micro_flashinfer.sh | 25 ++ ...oise_fa3_tp2_sp2_multi_micro_flashinfer.sh | 26 ++ 15 files changed, 1560 insertions(+), 70 deletions(-) create mode 100644 configs/hunyuan_image3/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json create mode 100644 configs/hunyuan_image3/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json create mode 100644 lightx2v/common/ops/attn/paged_flash_attn.py create mode 100644 lightx2v/models/networks/hunyuan_image3/custom_all_reduce.py create mode 100644 lightx2v/models/runners/hunyuan_image3/cuda_graph.py create mode 100755 scripts/hunyuan_image3/run_hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh create mode 100755 scripts/hunyuan_image3/run_hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh diff --git a/configs/hunyuan_image3/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json b/configs/hunyuan_image3/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json new file mode 100644 index 000000000..105d4319e --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json @@ -0,0 +1,79 @@ +{ + "infer_steps": 50, + "sample_guide_scale": 5.0, + "flow_shift": 1.0, + "target_height": 1024, + "target_width": 1024, + "feature_caching": "NoCaching", + "moe_backend": "multi_micro", + "attn_impl": "torch_sdpa", + "denoise_attn_impl": "flash_attn3", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache_by_phase": { + "ar": "save_results/hunyuan_image3_flashinfer_autotune_t2i_ar_tp4.json" + }, + "flashinfer_tune_max_num_tokens": 16384, + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], + "flashinfer_autotune_round_up": true, + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "ar_decode_attn_impl": "flash_attn3_paged", + "ar_kv_page_size": 16, + "ar_flash_attn_max_num_splits": 32, + "enable_ar_custom_all_reduce": true, + "ar_custom_all_reduce_backend": "vllm", + "ar_custom_all_reduce_required": true, + "ar_custom_all_reduce_max_size_bytes": 8388608, + "ar_custom_all_reduce_skip_p2p_check": true, + "ar_custom_all_reduce_graph_mode": "workspace", + "enable_ar_cuda_graph": true, + "ar_cuda_graph_mode": "full_decode", + "ar_cuda_graph_kv_bucket_size": 128, + "ar_cuda_graph_capture_q_lens": [ + 1 + ], + "ar_cuda_graph_required": true, + "use_taylor_cache": false, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": false, + "taylor_cache_first_enhance_steps": 3, + "taylor_cache_enable_tailing_enhance": false, + "taylor_cache_tailing_enhance_steps": 1, + "taylor_cache_low_freqs_order": 2, + "taylor_cache_high_freqs_order": 2, + "parallel": { + "pipeline_parallel": false, + "phase_aware": true, + "storage_tensor_p_size": 2, + "ar": { + "tensor_p_size": 4, + "seq_p_size": 1 + }, + "denoise": { + "tensor_p_size": 2, + "seq_p_size": 2 + }, + "cfg_p_size": 1, + "seq_p_attn_type": "ulysses", + "cfg_mode": "serial" + }, + "enable_cfg": true, + "bot_task": "think_recaption", + "use_system_prompt": "en_unified", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6 +} diff --git a/configs/hunyuan_image3/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json b/configs/hunyuan_image3/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json new file mode 100644 index 000000000..4ad350f78 --- /dev/null +++ b/configs/hunyuan_image3/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json @@ -0,0 +1,79 @@ +{ + "infer_steps": 50, + "sample_guide_scale": 5.0, + "flow_shift": 1.0, + "image_size": "auto", + "feature_caching": "NoCaching", + "moe_backend": "multi_micro", + "attn_impl": "torch_sdpa", + "denoise_attn_impl": "flash_attn3", + "flashinfer_autotune_mode": "auto", + "flashinfer_autotune_cache_by_phase": { + "ar": "save_results/hunyuan_image3_flashinfer_autotune_ti2i_ar_tp4.json" + }, + "flashinfer_tune_max_num_tokens": 16384, + "flashinfer_tuning_buckets": [ + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 12288, + 16384 + ], + "flashinfer_autotune_round_up": true, + "enable_kv_cache": true, + "enable_text_kv_cache": true, + "ar_decode_attn_impl": "flash_attn3_paged", + "ar_kv_page_size": 16, + "ar_flash_attn_max_num_splits": 32, + "enable_ar_custom_all_reduce": true, + "ar_custom_all_reduce_backend": "vllm", + "ar_custom_all_reduce_required": true, + "ar_custom_all_reduce_max_size_bytes": 8388608, + "ar_custom_all_reduce_skip_p2p_check": true, + "ar_custom_all_reduce_graph_mode": "workspace", + "enable_ar_cuda_graph": true, + "ar_cuda_graph_mode": "full_decode", + "ar_cuda_graph_kv_bucket_size": 128, + "ar_cuda_graph_capture_q_lens": [ + 1 + ], + "ar_cuda_graph_required": true, + "use_taylor_cache": false, + "taylor_cache_interval": 5, + "taylor_cache_order": 2, + "taylor_cache_enable_first_enhance": false, + "taylor_cache_first_enhance_steps": 3, + "taylor_cache_enable_tailing_enhance": false, + "taylor_cache_tailing_enhance_steps": 1, + "taylor_cache_low_freqs_order": 2, + "taylor_cache_high_freqs_order": 2, + "parallel": { + "pipeline_parallel": false, + "phase_aware": true, + "storage_tensor_p_size": 2, + "ar": { + "tensor_p_size": 4, + "seq_p_size": 1 + }, + "denoise": { + "tensor_p_size": 2, + "seq_p_size": 2 + }, + "cfg_p_size": 1, + "seq_p_attn_type": "ulysses", + "cfg_mode": "serial" + }, + "enable_cfg": true, + "infer_align_image_size": true, + "bot_task": "think_recaption", + "use_system_prompt": "en_unified", + "max_new_tokens": 2048, + "text_do_sample": true, + "text_top_k": 1024, + "text_top_p": 0.95, + "text_temperature": 0.6 +} diff --git a/lightx2v/common/ops/attn/paged_flash_attn.py b/lightx2v/common/ops/attn/paged_flash_attn.py new file mode 100644 index 000000000..d8f6940dd --- /dev/null +++ b/lightx2v/common/ops/attn/paged_flash_attn.py @@ -0,0 +1,171 @@ +"""Paged FlashAttention 3 support for single-token autoregressive decode.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER + +from .template import AttnWeightTemplate + +try: + from flash_attn_interface import flash_attn_with_kvcache as flash_attn3_with_kvcache + from flash_attn_interface import get_scheduler_metadata as get_flash_attn3_scheduler_metadata +except ImportError: + flash_attn3_with_kvcache = None + get_flash_attn3_scheduler_metadata = None + + +def require_paged_flash_attn3() -> None: + if flash_attn3_with_kvcache is None or get_flash_attn3_scheduler_metadata is None: + raise ImportError("flash_attn3_paged requires the standalone FlashAttention 3 package with paged-KV and scheduler-metadata support.") + + +def build_flash_attn3_decode_scheduler_metadata( + *, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + num_query_heads: int, + num_key_value_heads: int, + head_dim: int, + page_size: int, + qkv_dtype: torch.dtype, + max_num_splits: int, +) -> torch.Tensor: + """Build FlashAttention 3 scheduling metadata for one-token decode.""" + + require_paged_flash_attn3() + if cache_seqlens.dtype != torch.int32 or not cache_seqlens.is_cuda: + raise ValueError("FlashAttention 3 cache_seqlens must be a CUDA int32 tensor.") + return get_flash_attn3_scheduler_metadata( + batch_size=int(cache_seqlens.numel()), + max_seqlen_q=1, + max_seqlen_k=int(max_seqlen_k), + num_heads_q=int(num_query_heads), + num_heads_kv=int(num_key_value_heads), + headdim=int(head_dim), + cache_seqlens=cache_seqlens, + qkv_dtype=qkv_dtype, + page_size=int(page_size), + max_seqlen_k_new=0, + causal=True, + num_splits=int(max_num_splits), + ) + + +@triton.jit +def _paged_kv_store_kernel( + key, + value, + key_cache, + value_cache, + page_table, + cache_seqlens, + key_stride_head: tl.constexpr, + key_stride_dim: tl.constexpr, + value_stride_head: tl.constexpr, + value_stride_dim: tl.constexpr, + num_kv_heads: tl.constexpr, + head_dim: tl.constexpr, + page_size: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = offsets < num_kv_heads * head_dim + head = offsets // head_dim + dim = offsets - head * head_dim + + sequence_length = tl.load(cache_seqlens) + token_position = sequence_length - 1 + logical_page = token_position // page_size + page_offset = token_position - logical_page * page_size + physical_page = tl.load(page_table + logical_page) + cache_offset = ((physical_page * page_size + page_offset) * num_kv_heads + head) * head_dim + dim + + key_value = tl.load(key + head * key_stride_head + dim * key_stride_dim, mask=valid) + value_value = tl.load(value + head * value_stride_head + dim * value_stride_dim, mask=valid) + tl.store(key_cache + cache_offset, key_value, mask=valid) + tl.store(value_cache + cache_offset, value_value, mask=valid) + + +def store_paged_kv(key, value, key_cache, value_cache, page_table, cache_seqlens) -> None: + """Store one strided K/V token in an NHD paged cache.""" + + num_kv_heads = int(key.shape[1]) + head_dim = int(key.shape[-1]) + total = num_kv_heads * head_dim + _paged_kv_store_kernel[(triton.cdiv(total, 256),)]( + key, + value, + key_cache, + value_cache, + page_table, + cache_seqlens, + key_stride_head=key.stride(1), + key_stride_dim=key.stride(3), + value_stride_head=value.stride(1), + value_stride_dim=value.stride(3), + num_kv_heads=num_kv_heads, + head_dim=head_dim, + page_size=int(key_cache.shape[1]), + BLOCK=256, + ) + + +@ATTN_WEIGHT_REGISTER("flash_attn3_paged") +class PagedFlashAttn3Weight(AttnWeightTemplate): + """Native-GQA FlashAttention 3 decode over an NHD paged KV cache.""" + + def __init__(self): + require_paged_flash_attn3() + self.config = {} + + def apply(self, q, k, v, **kwargs): + raise RuntimeError("flash_attn3_paged is decode-only; use apply_decode with a persistent paged KV cache.") + + def apply_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + scheduler_metadata: torch.Tensor, + max_num_splits: int, + softmax_scale: float | None = None, + ) -> torch.Tensor: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"Paged FA3 decode expects 4D Q/K/V, got {q.shape}, {k.shape}, {v.shape}.") + if q.shape[0] != 1 or q.shape[2] != 1 or k.shape[2] != 1 or v.shape[2] != 1: + raise ValueError(f"Paged FA3 decode requires batch=1 and q_len=1, got {q.shape}, {k.shape}, {v.shape}.") + if q.shape[1] % k.shape[1]: + raise ValueError(f"Paged FA3 GQA requires Q heads divisible by KV heads, got Hq={q.shape[1]}, Hkv={k.shape[1]}.") + if page_table.dtype != torch.int32 or cache_seqlens.dtype != torch.int32: + raise ValueError("Paged FA3 page_table and cache_seqlens must use torch.int32.") + + store_paged_kv(k, v, k_cache, v_cache, page_table, cache_seqlens) + output = flash_attn3_with_kvcache( + q.transpose(1, 2), + k_cache, + v_cache, + cache_seqlens=cache_seqlens, + page_table=page_table, + softmax_scale=softmax_scale, + causal=True, + scheduler_metadata=scheduler_metadata, + num_splits=int(max_num_splits), + ) + return output.transpose(1, 2) + + +__all__ = [ + "PagedFlashAttn3Weight", + "build_flash_attn3_decode_scheduler_metadata", + "require_paged_flash_attn3", + "store_paged_kv", +] diff --git a/lightx2v/infer.py b/lightx2v/infer.py index e2167a8c1..d824322d4 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -360,18 +360,26 @@ def main(): validate_config_paths(config) - with ProfilingContext4DebugL1("Total Cost"): - # init runner - runner = init_runner(config) - # start to infer - data = args.__dict__ - update_input_info_from_dict(input_info, data) - runner.run_pipeline(input_info) - - # Clean up distributed process group - if dist.is_initialized(): - dist.destroy_process_group() - logger.info("Distributed process group cleaned up") + runner = None + try: + with ProfilingContext4DebugL1("Total Cost"): + # init runner + runner = init_runner(config) + # start to infer + data = args.__dict__ + update_input_info_from_dict(input_info, data) + runner.run_pipeline(input_info) + finally: + try: + if runner is not None and config.get("model_cls") == "hunyuan_image3": + close = getattr(runner, "close", None) + if callable(close): + close() + finally: + # Graph-backed collectives must be released before process-group teardown. + if dist.is_initialized(): + dist.destroy_process_group() + logger.info("Distributed process group cleaned up") if __name__ == "__main__": diff --git a/lightx2v/models/networks/hunyuan_image3/config.py b/lightx2v/models/networks/hunyuan_image3/config.py index 439283c19..df29e1cde 100644 --- a/lightx2v/models/networks/hunyuan_image3/config.py +++ b/lightx2v/models/networks/hunyuan_image3/config.py @@ -1,5 +1,13 @@ SUPPORTED_TASKS = {"t2t", "t2i", "ti2t", "ti2i", "i2i"} SUPPORTED_BOT_TASKS = {"image", "auto", "think", "recaption", "think_recaption"} +_AR_CUSTOM_ALL_REDUCE_KEYS = { + "enable_ar_custom_all_reduce", + "ar_custom_all_reduce_backend", + "ar_custom_all_reduce_required", + "ar_custom_all_reduce_max_size_bytes", + "ar_custom_all_reduce_skip_p2p_check", + "ar_custom_all_reduce_graph_mode", +} def _config_ints(value): @@ -10,6 +18,34 @@ def _config_ints(value): return [int(value)] +def _normalize_ar_custom_all_reduce(config, *, phase_aware, parallel_config): + if not any(key in config for key in _AR_CUSTOM_ALL_REDUCE_KEYS): + return + + from lightx2v.models.networks.hunyuan_image3.custom_all_reduce import HunyuanImage3CustomAllReduceConfig + + custom_ar = HunyuanImage3CustomAllReduceConfig.from_mapping(config) + if custom_ar.enabled: + if not phase_aware: + raise ValueError("HunyuanImage3 AR custom all-reduce requires parallel.phase_aware=true.") + ar_tp_size = int((parallel_config.get("ar") or {}).get("tensor_p_size", 1)) + if ar_tp_size not in (2, 4): + raise ValueError(f"HunyuanImage3 AR custom all-reduce supports full-world TP2 or TP4, got AR TP size {ar_tp_size}.") + if custom_ar.graph_mode == "workspace" and config.get("enable_ar_cuda_graph") is not True: + raise ValueError("ar_custom_all_reduce_graph_mode='workspace' requires enable_ar_cuda_graph=true.") + + config.update( + { + "enable_ar_custom_all_reduce": custom_ar.enabled, + "ar_custom_all_reduce_backend": custom_ar.backend, + "ar_custom_all_reduce_required": custom_ar.required, + "ar_custom_all_reduce_max_size_bytes": custom_ar.max_size_bytes, + "ar_custom_all_reduce_skip_p2p_check": custom_ar.skip_p2p_check, + "ar_custom_all_reduce_graph_mode": custom_ar.graph_mode, + } + ) + + def normalize_hunyuan_image3_phase_parallel(config, parallel_config): moe_backend = config["moe_backend"] @@ -225,12 +261,14 @@ def normalize_hunyuan_image3_config(config): parallel_config = config.get("parallel") if not isinstance(parallel_config, dict): + _normalize_ar_custom_all_reduce(config, phase_aware=False, parallel_config={}) if moe_backend == "multi_micro": raise ValueError("HunyuanImage3 moe_backend='multi_micro' requires a phase-aware parallel configuration.") return config parallel_config = dict(parallel_config) phase_aware, pipeline_parallel, cfg_mode = _normalize_parallel_config(config, parallel_config, task) + _normalize_ar_custom_all_reduce(config, phase_aware=phase_aware, parallel_config=parallel_config) config["parallel"] = parallel_config config["pipeline_parallel"] = pipeline_parallel config["hunyuan_cfg_mode"] = cfg_mode diff --git a/lightx2v/models/networks/hunyuan_image3/custom_all_reduce.py b/lightx2v/models/networks/hunyuan_image3/custom_all_reduce.py new file mode 100644 index 000000000..c07770910 --- /dev/null +++ b/lightx2v/models/networks/hunyuan_image3/custom_all_reduce.py @@ -0,0 +1,319 @@ +"""CUDA-graph-aware custom all-reduce for HunyuanImage3 AR. + +The vLLM backend uses a Gloo process group to exchange CUDA IPC metadata and +CUDA kernels to perform the reduction. Denoising deliberately does not use +this wrapper; its TP2 reductions continue to use the active NCCL group. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} +_SUPPORTED_AR_TP_SIZES = frozenset((2, 4)) +_VLLM_SKIP_P2P_CHECK_ENV = "VLLM_SKIP_P2P_CHECK" + + +class HunyuanImage3CustomAllReduceError(RuntimeError): + """Raised when a required AR custom all-reduce cannot be honored.""" + + +def _as_bool(value: Any, *, name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, int) and value in (0, 1): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean, got {value!r}.") + + +@dataclass(frozen=True) +class HunyuanImage3CustomAllReduceConfig: + enabled: bool = False + backend: str = "vllm" + required: bool = False + max_size_bytes: int = 8 * 1024 * 1024 + skip_p2p_check: bool = True + graph_mode: str = "direct" + + @classmethod + def from_mapping(cls, config: Mapping[str, Any] | None) -> HunyuanImage3CustomAllReduceConfig: + values: Mapping[str, Any] = {} if config is None else config + if not isinstance(values, Mapping): + raise TypeError(f"custom all-reduce config must be a mapping, got {type(values).__name__}.") + + parsed = cls( + enabled=_as_bool(values.get("enable_ar_custom_all_reduce", False), name="enable_ar_custom_all_reduce"), + backend=str(values.get("ar_custom_all_reduce_backend", "vllm")).strip().lower(), + required=_as_bool(values.get("ar_custom_all_reduce_required", False), name="ar_custom_all_reduce_required"), + max_size_bytes=int(values.get("ar_custom_all_reduce_max_size_bytes", 8 * 1024 * 1024)), + skip_p2p_check=_as_bool( + values.get("ar_custom_all_reduce_skip_p2p_check", True), + name="ar_custom_all_reduce_skip_p2p_check", + ), + graph_mode=str(values.get("ar_custom_all_reduce_graph_mode", "direct")).strip().lower(), + ) + if parsed.required and not parsed.enabled: + raise ValueError("ar_custom_all_reduce_required=true requires enable_ar_custom_all_reduce=true.") + if parsed.enabled and parsed.backend != "vllm": + raise ValueError(f"Unsupported AR custom all-reduce backend {parsed.backend!r}; expected 'vllm'.") + if parsed.max_size_bytes <= 0: + raise ValueError("ar_custom_all_reduce_max_size_bytes must be positive.") + if parsed.graph_mode not in {"direct", "workspace"}: + raise ValueError("ar_custom_all_reduce_graph_mode must be 'direct' or 'workspace'.") + return parsed + + +def _create_vllm_backend(**kwargs: Any) -> Any: + # vLLM remains optional for configurations that do not enable custom AR. + from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + return CustomAllreduce(**kwargs) + + +class HunyuanImage3CustomAllReduce: + """Dispatch eligible AR SUM reductions to vLLM custom all-reduce.""" + + def __init__( + self, + *, + metadata_group: dist.ProcessGroup, + fallback_group: dist.ProcessGroup | Callable[[], dist.ProcessGroup], + config: Mapping[str, Any] | HunyuanImage3CustomAllReduceConfig | None, + device: torch.device | str | int, + phase_getter: Callable[[], str], + ) -> None: + self.config = config if isinstance(config, HunyuanImage3CustomAllReduceConfig) else HunyuanImage3CustomAllReduceConfig.from_mapping(config) + self.metadata_group = metadata_group + self._fallback_group = fallback_group + self.device = torch.device(device) + self._phase_getter = phase_getter + self._backend: Any | None = None + self._initialization_attempted = False + self._available = False + self._closed = False + self._capture_depth = 0 + self._unavailable_reason: str | None = None + + @property + def enabled(self) -> bool: + return self.config.enabled + + @property + def required(self) -> bool: + return self.config.required + + @property + def available(self) -> bool: + return self._available and not self._closed + + @staticmethod + def _backend_name(group: dist.ProcessGroup) -> str: + return str(dist.get_backend(group)).strip().lower() + + def _initialization_failure(self, reason: str, error: BaseException | None = None) -> bool: + self._available = False + self._unavailable_reason = reason + if self.required: + message = f"Required HunyuanImage3 AR custom all-reduce initialization failed: {reason}." + if error is None: + raise HunyuanImage3CustomAllReduceError(message) + raise HunyuanImage3CustomAllReduceError(message) from error + return False + + def initialize(self) -> bool: + """Collectively initialize the vLLM backend on full-world AR ranks.""" + + if self._closed: + raise HunyuanImage3CustomAllReduceError("Cannot initialize a closed custom all-reduce.") + if not self.enabled: + return False + if self._initialization_attempted: + if self.required and not self._available: + raise HunyuanImage3CustomAllReduceError(f"Required HunyuanImage3 AR custom all-reduce is unavailable: {self._unavailable_reason}.") + return self._available + + self._initialization_attempted = True + if self.metadata_group is None: + return self._initialization_failure("metadata_group_missing") + if not dist.is_available() or not dist.is_initialized(): + return self._initialization_failure("torch_distributed_not_initialized") + + try: + metadata_backend = self._backend_name(self.metadata_group) + if "gloo" not in metadata_backend: + return self._initialization_failure(f"metadata_group_backend_{metadata_backend or 'unknown'}") + world_size = int(dist.get_world_size(group=self.metadata_group)) + if world_size not in _SUPPORTED_AR_TP_SIZES: + return self._initialization_failure(f"unsupported_ar_tp_size_{world_size}") + + previous_skip_check = os.environ.get(_VLLM_SKIP_P2P_CHECK_ENV) + if self.config.skip_p2p_check: + os.environ[_VLLM_SKIP_P2P_CHECK_ENV] = "1" + try: + backend = _create_vllm_backend( + group=self.metadata_group, + device=self.device, + max_size=self.config.max_size_bytes, + symm_mem_enabled=False, + ) + finally: + if self.config.skip_p2p_check: + if previous_skip_check is None: + os.environ.pop(_VLLM_SKIP_P2P_CHECK_ENV, None) + else: + os.environ[_VLLM_SKIP_P2P_CHECK_ENV] = previous_skip_check + + self._backend = backend + if backend is None: + return self._initialization_failure("backend_factory_returned_none") + if bool(getattr(backend, "disabled", False)): + return self._initialization_failure("vllm_backend_disabled") + except HunyuanImage3CustomAllReduceError: + raise + except Exception as error: + return self._initialization_failure(type(error).__name__, error) + + self._available = True + self._unavailable_reason = None + return True + + @staticmethod + def _weakly_contiguous(tensor: torch.Tensor) -> bool: + if tensor.is_contiguous(): + return True + try: + storage_bytes = tensor.untyped_storage().nbytes() + offset_bytes = tensor.storage_offset() * tensor.element_size() + return storage_bytes - offset_bytes == tensor.numel() * tensor.element_size() + except (AttributeError, RuntimeError): + return False + + def _tensor_eligibility(self, tensor: torch.Tensor) -> tuple[bool, str]: + if tensor.device.type != "cuda": + return False, "not_cuda" + if tensor.dtype not in _SUPPORTED_DTYPES: + return False, f"unsupported_dtype_{tensor.dtype}" + if tensor.numel() == 0: + return False, "empty_tensor" + num_bytes = tensor.numel() * tensor.element_size() + if num_bytes % 16: + return False, "size_not_16_byte_aligned" + if num_bytes >= self.config.max_size_bytes: + return False, "size_exceeds_custom_ar_limit" + if not self._weakly_contiguous(tensor): + return False, "not_weakly_contiguous" + return True, "eligible" + + def _custom_decision(self, tensor: torch.Tensor) -> tuple[bool, str, bool]: + locally_eligible, reason = self._tensor_eligibility(tensor) + if not locally_eligible: + return False, reason, False + if not self.enabled: + return False, "custom_ar_disabled", True + if not self._initialization_attempted: + return False, "custom_ar_not_initialized", True + if not self.available or self._backend is None: + return False, self._unavailable_reason or "custom_ar_unavailable", True + try: + if not bool(self._backend.should_custom_ar(tensor)): + return False, "vllm_should_custom_ar_rejected", True + except Exception as error: + return False, f"vllm_should_custom_ar_{type(error).__name__}", True + return True, "eligible", True + + def _nccl_fallback(self, tensor: torch.Tensor, *, reason: str) -> torch.Tensor: + group = self._fallback_group() if callable(self._fallback_group) else self._fallback_group + if group is None: + raise HunyuanImage3CustomAllReduceError(f"NCCL fallback group is missing for AR all-reduce ({reason}).") + backend = self._backend_name(group) + if "nccl" not in backend: + raise HunyuanImage3CustomAllReduceError(f"HunyuanImage3 AR fallback must use NCCL, got backend {backend!r}.") + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group) + return tensor + + def all_reduce(self, tensor: torch.Tensor, *, is_decode: bool = False) -> torch.Tensor: + """SUM-reduce an AR tensor and fail closed for required decode work.""" + + if self._closed: + raise HunyuanImage3CustomAllReduceError("Cannot use a closed custom all-reduce.") + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"all_reduce expects torch.Tensor, got {type(tensor).__name__}.") + use_custom, reason, locally_eligible = self._custom_decision(tensor) + if use_custom: + try: + if self.config.graph_mode == "workspace" and self._capture_depth and torch.cuda.is_current_stream_capturing(): + output = self._backend.all_reduce(tensor, registered=False) + else: + output = self._backend.custom_all_reduce(tensor) + except Exception as error: + reason = f"vllm_custom_all_reduce_{type(error).__name__}" + if self.required and is_decode and locally_eligible: + raise HunyuanImage3CustomAllReduceError(f"Required AR decode custom all-reduce failed: {reason}.") from error + else: + if output is not None: + return output + reason = "vllm_custom_all_reduce_returned_none" + + if self.required and is_decode and locally_eligible: + raise HunyuanImage3CustomAllReduceError(f"Required AR decode tensor cannot fall back to NCCL: {reason}.") + return self._nccl_fallback(tensor, reason=reason) + + @contextmanager + def capture(self) -> Iterator[None]: + """Register stable custom-all-reduce buffers during AR graph capture.""" + + if self._closed: + raise HunyuanImage3CustomAllReduceError("Cannot capture with a closed custom all-reduce.") + if not self.enabled: + yield + return + if self._phase_getter().strip().lower() != "ar": + raise HunyuanImage3CustomAllReduceError("Custom all-reduce graph capture is only valid during the AR phase.") + if not self.available or self._backend is None: + if self.required: + raise HunyuanImage3CustomAllReduceError(f"Required AR custom all-reduce is unavailable during graph capture: {self._unavailable_reason or 'not_initialized'}.") + yield + return + if self._capture_depth: + raise HunyuanImage3CustomAllReduceError("Nested custom all-reduce capture contexts are not supported.") + + self._capture_depth = 1 + try: + with self._backend.capture(): + yield + finally: + self._capture_depth = 0 + + def close(self) -> None: + """Release vLLM CUDA IPC resources.""" + + if self._closed: + return + if self._capture_depth: + raise HunyuanImage3CustomAllReduceError("Cannot close custom all-reduce during graph capture.") + backend = self._backend + self._backend = None + self._available = False + self._closed = True + if backend is not None: + backend.close() + + +__all__ = [ + "HunyuanImage3CustomAllReduce", + "HunyuanImage3CustomAllReduceConfig", + "HunyuanImage3CustomAllReduceError", +] diff --git a/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py b/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py index e0402d2ee..58c445d8a 100644 --- a/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py +++ b/lightx2v/models/networks/hunyuan_image3/infer/kv_cache.py @@ -13,23 +13,95 @@ class HunyuanImage3KVCacheLayer: class HunyuanImage3StaticKVCache: """Per-layer KV cache matching HunyuanImage3 gen_text/gen_image inference.""" - def __init__(self, num_layers, max_cache_len, dynamic=False): + def __init__(self, num_layers, max_cache_len, dynamic=False, paged=False, page_size=16): self.num_layers = int(num_layers) - self.max_cache_len = int(max_cache_len) + requested_cache_len = int(max_cache_len) self.dynamic = bool(dynamic) + self.paged = bool(paged) + self.page_size = int(page_size) + if self.page_size < 1: + raise ValueError(f"HunyuanImage3 KV page_size must be positive, got {self.page_size}.") + self.num_pages = math.ceil(requested_cache_len / self.page_size) if self.paged else 0 + self.max_cache_len = self.num_pages * self.page_size if self.paged else requested_cache_len self.layers = [HunyuanImage3KVCacheLayer() for _ in range(self.num_layers)] + self.page_table = None + self.cache_seqlens = None + self.scheduler_metadata = None + self.num_key_value_heads = None + self.head_dim = None def _ensure_layer(self, layer_idx, key_states, value_states): layer = self.layers[layer_idx] if layer.key is None: - key_shape = (*key_states.shape[:2], self.max_cache_len, key_states.shape[-1]) - value_shape = (*value_states.shape[:2], self.max_cache_len, value_states.shape[-1]) + if self.paged: + if key_states.shape[0] != 1 or value_states.shape[0] != 1: + raise ValueError("HunyuanImage3 paged KV cache currently requires batch size 1.") + self._ensure_paged_metadata(key_states.device) + self.num_key_value_heads = int(key_states.shape[1]) + self.head_dim = int(key_states.shape[-1]) + key_shape = (self.num_pages, self.page_size, key_states.shape[1], key_states.shape[-1]) + value_shape = (self.num_pages, self.page_size, value_states.shape[1], value_states.shape[-1]) + else: + key_shape = (*key_states.shape[:2], self.max_cache_len, key_states.shape[-1]) + value_shape = (*value_states.shape[:2], self.max_cache_len, value_states.shape[-1]) layer.key = torch.zeros(key_shape, device=key_states.device, dtype=key_states.dtype) layer.value = torch.zeros(value_shape, device=value_states.device, dtype=value_states.dtype) return layer + def _ensure_paged_metadata(self, device): + if not self.paged: + raise RuntimeError("Paged metadata was requested from a dense HunyuanImage3 KV cache.") + if self.page_table is None: + self.page_table = torch.arange(self.num_pages, device=device, dtype=torch.int32).reshape(1, self.num_pages) + self.cache_seqlens = torch.zeros(1, device=device, dtype=torch.int32) + + def set_paged_decode_length(self, length): + if not self.paged or self.cache_seqlens is None: + raise RuntimeError("HunyuanImage3 paged decode metadata is not initialized; run prefill first.") + length = int(length) + if length < 1 or length > self.max_cache_len: + raise ValueError(f"HunyuanImage3 paged decode length must be in [1, {self.max_cache_len}], got {length}.") + self.cache_seqlens.fill_(length) + + def prepare_paged_decode_scheduler(self, *, valid_length, num_query_heads, max_num_splits): + """Refresh FA3 scheduler data without changing its persistent address.""" + + if not self.paged or self.cache_seqlens is None or self.num_key_value_heads is None or self.head_dim is None: + raise RuntimeError("HunyuanImage3 paged KV cache must be populated by prefill before decode scheduling.") + self.set_paged_decode_length(valid_length) + from lightx2v.common.ops.attn.paged_flash_attn import build_flash_attn3_decode_scheduler_metadata + + runtime_metadata = build_flash_attn3_decode_scheduler_metadata( + cache_seqlens=self.cache_seqlens, + max_seqlen_k=self.max_cache_len, + num_query_heads=int(num_query_heads), + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + page_size=self.page_size, + qkv_dtype=self.layers[0].key.dtype, + max_num_splits=int(max_num_splits), + ) + if self.scheduler_metadata is None: + fixed_size = max(17, int(runtime_metadata.numel())) + self.scheduler_metadata = torch.zeros(fixed_size, device=runtime_metadata.device, dtype=torch.int32) + if runtime_metadata.numel() > self.scheduler_metadata.numel(): + raise RuntimeError(f"FA3 scheduler metadata grew from {self.scheduler_metadata.numel()} to {runtime_metadata.numel()} entries after allocation.") + self.scheduler_metadata.zero_() + self.scheduler_metadata[: runtime_metadata.numel()].copy_(runtime_metadata) + return self.scheduler_metadata[: runtime_metadata.numel()] + + def get_paged_layer(self, layer_idx): + if not self.paged: + raise RuntimeError("Paged layer storage was requested from a dense HunyuanImage3 KV cache.") + layer = self.layers[int(layer_idx)] + if layer.key is None or layer.value is None: + raise RuntimeError(f"HunyuanImage3 paged KV layer {layer_idx} has not been allocated by prefill.") + return layer.key, layer.value + def update(self, key_states, value_states, layer_idx, cache_position=None): layer = self._ensure_layer(layer_idx, key_states, value_states) + if self.paged: + return self._update_paged(layer, key_states, value_states, cache_position) if cache_position is None: layer.key[:, :, : key_states.shape[2]].copy_(key_states) layer.value[:, :, : value_states.shape[2]].copy_(value_states) @@ -51,6 +123,31 @@ def update(self, key_states, value_states, layer_idx, cache_position=None): layer.value[batch_idx].index_copy_(1, cache_position[batch_idx], value_states[batch_idx]) return self._slice_dynamic(layer, int(cache_position.max().item()) + 1) + def _update_paged(self, layer, key_states, value_states, cache_position): + if cache_position is None: + positions = torch.arange(key_states.shape[2], device=key_states.device, dtype=torch.long) + else: + positions = cache_position.to(device=key_states.device, dtype=torch.long) + if positions.dim() == 2: + if positions.shape[0] != 1: + raise ValueError("HunyuanImage3 paged KV cache currently requires batch size 1.") + positions = positions[0] + elif positions.dim() != 1: + raise ValueError(f"HunyuanImage3 cache_position must be 1D or 2D, got {positions.shape}.") + + flat_key = layer.key.reshape(self.max_cache_len, layer.key.shape[2], layer.key.shape[3]) + flat_value = layer.value.reshape(self.max_cache_len, layer.value.shape[2], layer.value.shape[3]) + new_key = key_states.transpose(1, 2).reshape(-1, key_states.shape[1], key_states.shape[-1]) + new_value = value_states.transpose(1, 2).reshape(-1, value_states.shape[1], value_states.shape[-1]) + flat_key.index_copy_(0, positions, new_key) + flat_value.index_copy_(0, positions, new_value) + + end = key_states.shape[2] if cache_position is None else int(positions[-1].item()) + 1 + end = min(int(end), self.max_cache_len) + dense_key = flat_key[:end].unsqueeze(0).transpose(1, 2) + dense_value = flat_value[:end].unsqueeze(0).transpose(1, 2) + return dense_key, dense_value + def _slice_dynamic(self, layer, end): if not self.dynamic: return layer.key, layer.value diff --git a/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py index 31a2f193a..199fa6bbf 100644 --- a/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py @@ -1,3 +1,5 @@ +import weakref + import torch import torch.distributed as dist import torch.nn.functional as F @@ -44,8 +46,18 @@ def __init__(self, config): self.num_key_value_heads = self.global_num_key_value_heads // self.tp_size self.hidden_act = config.get("hidden_act", "silu") self.attn_impl = self._normalize_attention_impl(config.get("attn_impl", "torch_sdpa")) - self.attn_kernel = None if self.attn_impl == "torch_sdpa" else self._build_attention_kernel(self.attn_impl) + self.ar_attn_impl = self._normalize_attention_impl(config.get("ar_attn_impl", self.attn_impl)) + self.denoise_attn_impl = self._normalize_attention_impl(config.get("denoise_attn_impl", self.attn_impl)) + self._attn_kernels = {impl: None if impl == "torch_sdpa" else self._build_attention_kernel(impl) for impl in {self.attn_impl, self.ar_attn_impl, self.denoise_attn_impl}} + self.attn_kernel = self._attn_kernels[self.attn_impl] + self.ar_decode_attn_impl = str(config.get("ar_decode_attn_impl", "disabled") or "disabled").strip().lower() + if self.ar_decode_attn_impl in ("none", "off", "false"): + self.ar_decode_attn_impl = "disabled" + if self.ar_decode_attn_impl not in ("disabled", "flash_attn3_paged"): + raise ValueError(f"Unsupported HunyuanImage3 ar_decode_attn_impl={self.ar_decode_attn_impl!r}; expected 'disabled' or 'flash_attn3_paged'.") + self.ar_decode_attn_kernel = None if self.ar_decode_attn_impl == "disabled" else self._build_attention_kernel(self.ar_decode_attn_impl) self._attn_cu_seqlens_cache = {} + self._attn_segment_specs_cache = {} self._attn_fallback_warnings = set() self._sp_gather_buffers = {} self._pre_infer_device_cache = {} @@ -77,6 +89,18 @@ def _parallel_context_value(context, *names, default=None): def _active_phase(self): return str(self._parallel_context_value(self.parallel_context, "phase", default="legacy")) + def _active_attention_impl(self): + phase = self._active_phase().strip().lower() + if phase == "ar": + return self.ar_attn_impl + if phase == "denoise": + return self.denoise_attn_impl + return self.attn_impl + + def _active_attention_kernel(self, attn_impl=None): + attn_impl = self._active_attention_impl() if attn_impl is None else attn_impl + return self._attn_kernels.get(attn_impl) + def _active_tp_state(self): group = self._parallel_context_value(self.parallel_context, "active_tp_group", "tp_group", default=self.tp_group) size = self._parallel_context_value(self.parallel_context, "active_tp_size", "tp_size") @@ -128,6 +152,10 @@ def _build_attention_kernel(self, attn_impl): if flash_attn_func_v3 is None or flash_attn_varlen_func_v3 is None: raise ImportError("HunyuanImage3 attn_impl='flash_attn3' requires flash-attn v3 / flash_attn_interface.") + elif attn_impl == "flash_attn3_paged": + from lightx2v.common.ops.attn.paged_flash_attn import require_paged_flash_attn3 + + require_paged_flash_attn3() elif attn_impl == "sage_attn2": from lightx2v.common.ops.attn.sage_attn import sageattn @@ -149,12 +177,12 @@ def _normalize_attention_dtype(self, tensor): return tensor.to(torch.bfloat16) return tensor.to(torch.float32) - def _get_cu_seqlens(self, name, batch, seq_len, device): - key = (name, batch, seq_len, device.type, device.index) + def _get_cu_seqlens(self, name, batch, seq_len, device, attn_impl): + key = (attn_impl, name, batch, seq_len, device.type, device.index) cu_seqlens = self._attn_cu_seqlens_cache.get(key) if cu_seqlens is None: cu_seqlens = torch.arange(0, batch * seq_len + 1, seq_len, dtype=torch.int32) - if self.attn_impl in ("flash_attn2", "flash_attn3"): + if attn_impl in ("flash_attn2", "flash_attn3"): cu_seqlens = cu_seqlens.to(device, non_blocking=True) self._attn_cu_seqlens_cache[key] = cu_seqlens return cu_seqlens @@ -178,34 +206,39 @@ def _attention_mask_mode(self, attention_mask, q_len, kv_len): return "causal" return "custom" - def _warn_attention_fallback_once(self, mask_mode): - key = (self.attn_impl, mask_mode) + def _warn_attention_fallback_once(self, attn_impl, mask_mode): + key = (attn_impl, mask_mode) if key in self._attn_fallback_warnings: return self._attn_fallback_warnings.add(key) logger.warning( "HunyuanImage3 attn_impl='{}' does not support {} attention masks in the low-intrusion path; falling back to PyTorch SDPA for this attention call.", - self.attn_impl, + attn_impl, mask_mode, ) def _sdpa_attention(self, query_states, key_states, value_states, attention_mask): + if query_states.shape[1] != key_states.shape[1]: + if query_states.shape[1] % key_states.shape[1]: + raise ValueError(f"HunyuanImage3 SDPA fallback requires Q heads to be divisible by KV heads; got Q={query_states.shape[1]}, KV={key_states.shape[1]}.") + repeat_groups = query_states.shape[1] // key_states.shape[1] + key_states = repeat_kv(key_states, repeat_groups) + value_states = repeat_kv(value_states, repeat_groups) if query_states.device.type == "cuda" and attention_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() return F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=0.0) - def _apply_registered_attention_kernel(self, query_states, key_states, value_states, causal): - batch, query_heads, q_len, _ = query_states.shape - kv_len = key_states.shape[2] - original_dtype = query_states.dtype - q = self._normalize_attention_dtype(query_states.transpose(1, 2)).contiguous() - k = self._normalize_attention_dtype(key_states.transpose(1, 2)).contiguous() - v = self._normalize_attention_dtype(value_states.transpose(1, 2)).contiguous() - cu_seqlens_q = self._get_cu_seqlens("q", batch, q_len, q.device) - cu_seqlens_kv = self._get_cu_seqlens("kv", batch, kv_len, k.device) - attn_output = self.attn_kernel.apply( + def _apply_registered_attention_kernel_bshd(self, q, k, v, causal, attn_impl): + batch, q_len, query_heads, _ = q.shape + kv_len = k.shape[1] + attn_kernel = self._active_attention_kernel(attn_impl) + if attn_kernel is None: + raise RuntimeError(f"HunyuanImage3 attn_impl={attn_impl!r} has no registered attention kernel.") + cu_seqlens_q = self._get_cu_seqlens("q", batch, q_len, q.device, attn_impl) + cu_seqlens_kv = self._get_cu_seqlens("kv", batch, kv_len, k.device, attn_impl) + attn_output = attn_kernel.apply( q=q, k=k, v=v, @@ -220,7 +253,15 @@ def _apply_registered_attention_kernel(self, query_states, key_states, value_sta elif attn_output.dim() == 3: attn_output = attn_output.reshape(batch, q_len, query_heads, self.head_dim) else: - raise RuntimeError(f"HunyuanImage3 attn_impl={self.attn_impl!r} returned unexpected shape {tuple(attn_output.shape)}.") + raise RuntimeError(f"HunyuanImage3 attn_impl={attn_impl!r} returned unexpected shape {tuple(attn_output.shape)}.") + return attn_output + + def _apply_registered_attention_kernel(self, query_states, key_states, value_states, causal, attn_impl): + original_dtype = query_states.dtype + q = self._normalize_attention_dtype(query_states.transpose(1, 2)).contiguous() + k = self._normalize_attention_dtype(key_states.transpose(1, 2)).contiguous() + v = self._normalize_attention_dtype(value_states.transpose(1, 2)).contiguous() + attn_output = self._apply_registered_attention_kernel_bshd(q, k, v, causal, attn_impl) return attn_output.to(original_dtype).transpose(1, 2) def _normalize_full_attn_slices(self, full_attn_slices, batch): @@ -297,7 +338,53 @@ def _build_segment_specs(self, position_ids, full_slices, kv_len): local_start = local_end return segments - def _segmented_flash_attention(self, query_states, key_states, value_states, position_ids, full_attn_slices, segment_specs=None): + def _segment_specs_cache_key(self, position_ids, batch_full_slices, kv_len): + if not torch.is_tensor(position_ids): + return None + return ( + id(position_ids), + tuple(position_ids.shape), + tuple(position_ids.stride()), + int(position_ids.storage_offset()), + int(getattr(position_ids, "_version", 0)), + int(kv_len), + tuple(tuple(sample_slices) for sample_slices in batch_full_slices), + ) + + def _lookup_segment_specs_cache(self, key, position_ids): + if key is None: + return None + entry = self._attn_segment_specs_cache.get(key) + if entry is None: + return None + tensor_ref, segment_specs = entry + if tensor_ref() is position_ids: + return segment_specs + self._attn_segment_specs_cache.pop(key, None) + return None + + def _store_segment_specs_cache(self, key, position_ids, segment_specs): + if key is None: + return + if len(self._attn_segment_specs_cache) >= 32: + stale_keys = [cache_key for cache_key, (tensor_ref, _) in self._attn_segment_specs_cache.items() if tensor_ref() is None] + for cache_key in stale_keys: + self._attn_segment_specs_cache.pop(cache_key, None) + if len(self._attn_segment_specs_cache) >= 32: + self._attn_segment_specs_cache.pop(next(iter(self._attn_segment_specs_cache))) + self._attn_segment_specs_cache[key] = (weakref.ref(position_ids), segment_specs) + + def _segmented_flash_attention( + self, + query_states, + key_states, + value_states, + position_ids, + full_attn_slices, + segment_specs=None, + attn_impl=None, + ): + attn_impl = self._active_attention_impl() if attn_impl is None else attn_impl if position_ids is None: return None batch, _, q_len, _ = query_states.shape @@ -315,17 +402,22 @@ def _segmented_flash_attention(self, query_states, key_states, value_states, pos if len(segment_specs) != batch or any(specs is None for specs in segment_specs): return None - output = torch.empty_like(query_states) + original_dtype = query_states.dtype + q = self._normalize_attention_dtype(query_states.transpose(1, 2)).contiguous() + k = self._normalize_attention_dtype(key_states.transpose(1, 2)).contiguous() + v = self._normalize_attention_dtype(value_states.transpose(1, 2)).contiguous() + output = torch.empty_like(q) for batch_idx in range(batch): for q_start, q_stop, kv_stop, causal in segment_specs[batch_idx]: - segment_output = self._apply_registered_attention_kernel( - query_states[batch_idx : batch_idx + 1, :, q_start:q_stop], - key_states[batch_idx : batch_idx + 1, :, :kv_stop], - value_states[batch_idx : batch_idx + 1, :, :kv_stop], + segment_output = self._apply_registered_attention_kernel_bshd( + q[batch_idx : batch_idx + 1, q_start:q_stop], + k[batch_idx : batch_idx + 1, :kv_stop], + v[batch_idx : batch_idx + 1, :kv_stop], causal=causal, + attn_impl=attn_impl, ) - output[batch_idx : batch_idx + 1, :, q_start:q_stop] = segment_output - return output + output[batch_idx : batch_idx + 1, q_start:q_stop] = segment_output + return output.to(original_dtype).transpose(1, 2) def _registered_attention( self, @@ -337,7 +429,8 @@ def _registered_attention( full_attn_slices=None, segment_specs=None, ): - if self.attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None: + attn_impl = self._active_attention_impl() + if attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None: segmented_output = self._segmented_flash_attention( query_states, key_states, @@ -345,6 +438,7 @@ def _registered_attention( position_ids, full_attn_slices, segment_specs=segment_specs, + attn_impl=attn_impl, ) if segmented_output is not None: return segmented_output @@ -354,9 +448,9 @@ def _registered_attention( kv_len = key_states.shape[2] mask_mode = self._attention_mask_mode(attention_mask, q_len, kv_len) - if self.attn_impl == "torch_sdpa": + if attn_impl == "torch_sdpa": return self._sdpa_attention(query_states, key_states, value_states, attention_mask) - if self.attn_impl in ("flash_attn2", "flash_attn3"): + if attn_impl in ("flash_attn2", "flash_attn3"): if mask_mode not in ("none", "full", "causal"): segmented_output = self._segmented_flash_attention( query_states, @@ -365,21 +459,28 @@ def _registered_attention( position_ids, full_attn_slices, segment_specs=segment_specs, + attn_impl=attn_impl, ) if segmented_output is not None: return segmented_output - self._warn_attention_fallback_once(mask_mode) + self._warn_attention_fallback_once(attn_impl, mask_mode) return self._sdpa_attention(query_states, key_states, value_states, attention_mask) causal = mask_mode == "causal" - elif self.attn_impl in ("sage_attn2", "sage_attn3"): + elif attn_impl in ("sage_attn2", "sage_attn3"): if mask_mode not in ("none", "full"): - self._warn_attention_fallback_once(mask_mode) + self._warn_attention_fallback_once(attn_impl, mask_mode) return self._sdpa_attention(query_states, key_states, value_states, attention_mask) causal = False else: - raise ValueError(f"Unsupported HunyuanImage3 normalized attn_impl={self.attn_impl!r}.") + raise ValueError(f"Unsupported HunyuanImage3 normalized attn_impl={attn_impl!r}.") - return self._apply_registered_attention_kernel(query_states, key_states, value_states, causal=causal) + return self._apply_registered_attention_kernel( + query_states, + key_states, + value_states, + causal=causal, + attn_impl=attn_impl, + ) @torch.no_grad() def infer(self, weights, pre_infer_out): @@ -397,7 +498,7 @@ def infer_block(self, block_idx, block, hidden_states, pre_infer_out): if device is not None and device.type == "cuda" and device.index is not None: torch.cuda.set_device(device.index) hidden_states = to_device(hidden_states, device) - use_segment_specs = self.attn_impl in ("flash_attn2", "flash_attn3") and pre_infer_out.attention_segment_specs is not None + use_segment_specs = self._active_attention_impl() in ("flash_attn2", "flash_attn3") and pre_infer_out.attention_segment_specs is not None attention_mask = None if use_segment_specs else self._cached_pre_infer_to_device("attention_mask", pre_infer_out.attention_mask, device) position_ids = self._cached_pre_infer_to_device("position_ids", pre_infer_out.position_ids, device) custom_pos_emb = self._cached_pre_infer_to_device("custom_pos_emb", pre_infer_out.custom_pos_emb, device) @@ -437,6 +538,7 @@ def infer_attention( segment_specs=None, ): batch, q_len, _ = hidden_states.shape + attn_impl = self._active_attention_impl() _, _, active_tp_size, _ = self._active_tp_state() if self.global_num_heads % active_tp_size or self.global_num_key_value_heads % active_tp_size: raise ValueError(f"HunyuanImage3 active TP size must divide Q and KV heads: Q={self.global_num_heads}, KV={self.global_num_key_value_heads}, active_tp_size={active_tp_size}.") @@ -492,7 +594,7 @@ def infer_attention( key_states.device, ) position_ids = cache_position_ids - if not (self.attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None): + if not (attn_impl in ("flash_attn2", "flash_attn3") and segment_specs is not None): attention_mask = self._cached_pre_infer_to_device( "sp_global_attention_mask", sequence_parallel_state.global_attention_mask, @@ -501,15 +603,40 @@ def infer_attention( else: raise ValueError(f"Unsupported HunyuanImage3 sequence parallel attention type: {sequence_parallel_state.attn_type!r}.") - if past_key_values is not None: + paged_decode = ( + past_key_values is not None and bool(getattr(past_key_values, "paged", False)) and self._active_phase().strip().lower() == "ar" and q_len == 1 and self.ar_decode_attn_kernel is not None + ) + + if paged_decode: + if sequence_parallel_state is not None: + raise RuntimeError("HunyuanImage3 paged AR decode does not support sequence parallel attention.") + if past_key_values.page_table is None or past_key_values.cache_seqlens is None or past_key_values.scheduler_metadata is None: + raise RuntimeError("HunyuanImage3 paged AR decode metadata must be prepared before model execution.") + key_cache, value_cache = past_key_values.get_paged_layer(block_idx) + query_states = query_states.to(key_cache.dtype) + key_states = key_states.to(key_cache.dtype) + value_states = value_states.to(value_cache.dtype) + attn_output = self.ar_decode_attn_kernel.apply_decode( + query_states, + key_states, + value_states, + k_cache=key_cache, + v_cache=value_cache, + page_table=past_key_values.page_table, + cache_seqlens=past_key_values.cache_seqlens, + scheduler_metadata=past_key_values.scheduler_metadata, + max_num_splits=int(self.config.get("ar_flash_attn_max_num_splits", 32)), + ) + elif past_key_values is not None: if cache_position_ids is None: raise ValueError("HunyuanImage3 KV cache requires position_ids.") key_states, value_states = past_key_values.update(key_states, value_states, block_idx, cache_position_ids) query_states = query_states.to(key_states.dtype) - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - if sequence_parallel_state is not None and sequence_parallel_state.attn_type == "kv_all_gather": + if not paged_decode and attn_impl not in ("flash_attn2", "flash_attn3"): + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + if not paged_decode and sequence_parallel_state is not None and sequence_parallel_state.attn_type == "kv_all_gather": valid_q_len = sequence_parallel_state.valid_local_seq_len attn_output = torch.zeros_like(query_states) if valid_q_len: @@ -522,7 +649,7 @@ def infer_attention( full_attn_slices=full_attn_slices, segment_specs=segment_specs, ) - else: + elif not paged_decode: attn_output = self._registered_attention( query_states, key_states, @@ -541,7 +668,7 @@ def infer_attention( return attn_output.reshape(batch, q_len, -1) def _prepare_attention_segment_specs(self, pre_infer_out): - if self.attn_impl not in ("flash_attn2", "flash_attn3"): + if self._active_attention_impl() not in ("flash_attn2", "flash_attn3"): return None state = pre_infer_out.sequence_parallel_state @@ -564,13 +691,20 @@ def _prepare_attention_segment_specs(self, pre_infer_out): return None batch, q_len = position_ids.shape kv_len = attention_mask.shape[-1] - if self._attention_mask_mode(attention_mask, q_len, kv_len) != "custom": + if attention_mask.dtype != torch.bool or attention_mask.dim() != 4 or attention_mask.shape[0] != batch or attention_mask.shape[1] != 1 or attention_mask.shape[-2] != q_len: return None batch_full_slices = self._normalize_full_attn_slices(pre_infer_out.full_attn_slices, batch) if batch_full_slices is None or not any(batch_full_slices): return None + cache_key = self._segment_specs_cache_key(position_ids, batch_full_slices, kv_len) + cached_specs = self._lookup_segment_specs_cache(cache_key, position_ids) + if cached_specs is not None: + return cached_specs segment_specs = [self._build_segment_specs(position_ids[batch_idx], batch_full_slices[batch_idx], kv_len) for batch_idx in range(batch)] - return None if any(specs is None for specs in segment_specs) else segment_specs + if any(specs is None for specs in segment_specs): + return None + self._store_segment_specs_cache(cache_key, position_ids, segment_specs) + return segment_specs def _cached_pre_infer_to_device(self, name, value, device): if value is None: @@ -669,7 +803,11 @@ def infer_mlp(self, phase, hidden_states): if tp_size > 1: if tp_group is None: raise RuntimeError("HunyuanImage3 active tensor parallelism requires an active TP process group.") - dist.all_reduce(output, op=dist.ReduceOp.SUM, group=tp_group) + phase_all_reduce = getattr(self.parallel_context, "tensor_parallel_all_reduce", None) + if callable(phase_all_reduce): + output = phase_all_reduce(output) + else: + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=tp_group) return output def _moe_topk(self, moe, hidden_states): diff --git a/lightx2v/models/networks/hunyuan_image3/model.py b/lightx2v/models/networks/hunyuan_image3/model.py index 32936d363..f34e8bb8b 100644 --- a/lightx2v/models/networks/hunyuan_image3/model.py +++ b/lightx2v/models/networks/hunyuan_image3/model.py @@ -450,6 +450,34 @@ def _infer_transformer_with_taylor_cache(self, pre_infer_out, cache_dic): self.taylor_cache.clear_derivatives() return hidden_states + @torch.no_grad() + def prepare_ar_pre_infer(self, inputs): + """Prepare embeddings and multimodal inputs for an AR forward.""" + + if self._active_seq_parallel(): + raise RuntimeError("HunyuanImage3 prepared AR inference requires sequence parallelism to be inactive.") + if inputs.get("cache_dic") is not None: + raise RuntimeError("HunyuanImage3 prepared AR inference does not support Taylor cache state.") + if inputs.get("_cfg_parallel_branch", False): + raise RuntimeError("HunyuanImage3 prepared AR inference does not support a CFG-parallel branch.") + if hasattr(self, "scheduler"): + self.scheduler.infer_condition = True + pre_infer_out = self.pre_infer.infer(self.pre_weight, inputs) + if pre_infer_out.sequence_parallel_state is not None: + raise RuntimeError("HunyuanImage3 prepared AR inference received unexpected sequence-parallel metadata.") + return pre_infer_out + + @torch.no_grad() + def infer_ar_prepared(self, pre_infer_out): + """Run transformer and output projection from prepared AR inputs.""" + + if self._active_seq_parallel(): + raise RuntimeError("HunyuanImage3 prepared AR inference requires sequence parallelism to be inactive.") + if pre_infer_out.sequence_parallel_state is not None: + raise RuntimeError("HunyuanImage3 prepared AR inference received sequence-parallel metadata.") + hidden_states = self._infer_transformer(pre_infer_out) + return self.post_infer.infer(self.post_weight, hidden_states, pre_infer_out) + @torch.no_grad() def _infer_cond_uncond(self, inputs, infer_condition=True): if hasattr(self, "scheduler"): diff --git a/lightx2v/models/networks/hunyuan_image3/parallel.py b/lightx2v/models/networks/hunyuan_image3/parallel.py index 315619de7..e40e54d51 100644 --- a/lightx2v/models/networks/hunyuan_image3/parallel.py +++ b/lightx2v/models/networks/hunyuan_image3/parallel.py @@ -1,6 +1,6 @@ from __future__ import annotations -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Iterator, Mapping @@ -72,6 +72,7 @@ def __init__( self.denoise_seq_size = int(denoise_seq_size) self._phase_states = dict(phase_states) self._phase = self._normalize_phase(initial_phase) + self.ar_custom_all_reduce = None @staticmethod def _normalize_phase(name: str) -> str: @@ -175,6 +176,34 @@ def stage(self, name: str) -> Iterator[HunyuanImage3ParallelContext]: finally: self.activate_phase(previous) + def tensor_parallel_all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + """SUM across active TP, using custom all-reduce only during AR.""" + + if self.active_tp_size <= 1: + return tensor + if self.active_tp_group is None: + raise RuntimeError(f"HunyuanImage3 phase={self.phase!r} has TP size {self.active_tp_size} without a TP group.") + + reducer = self.ar_custom_all_reduce + if self.phase == "ar" and reducer is not None: + token_count = tensor.numel() // tensor.shape[-1] if tensor.ndim else 0 + return reducer.all_reduce(tensor, is_decode=token_count == 1) + + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=self.active_tp_group) + return tensor + + def custom_all_reduce_capture(self): + reducer = self.ar_custom_all_reduce + if reducer is None: + return nullcontext() + return reducer.capture() + + def close_custom_all_reduce(self) -> None: + reducer = self.ar_custom_all_reduce + self.ar_custom_all_reduce = None + if reducer is not None: + reducer.close() + @property def _active(self) -> _PhaseParallelState: return self._phase_states[self._phase] @@ -324,6 +353,13 @@ def build_hunyuan_image3_parallel_context(config) -> HunyuanImage3ParallelContex raise RuntimeError(f"Invalid HunyuanImage3 AR logical TP mapping: {physical_to_logical}.") logical_gather_order = tuple(sorted(range(world_size), key=physical_to_logical.__getitem__)) + ar_metadata_group = None + if config.get("enable_ar_custom_all_reduce", False): + # vLLM uses this CPU group only to exchange CUDA IPC metadata. AR is + # full-world in the official phase-aware topology, so every rank is a + # member and enters initialization collectives in the same order. + ar_metadata_group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + phase_states = { "ar": _PhaseParallelState( tp_group=dist.group.WORLD, @@ -347,7 +383,7 @@ def build_hunyuan_image3_parallel_context(config) -> HunyuanImage3ParallelContex ), } - return HunyuanImage3ParallelContext( + context = HunyuanImage3ParallelContext( device_mesh=device_mesh, storage_tp_group=storage_tp_group, storage_tp_rank=storage_tp_rank, @@ -359,6 +395,19 @@ def build_hunyuan_image3_parallel_context(config) -> HunyuanImage3ParallelContex denoise_seq_size=denoise_seq_size, phase_states=phase_states, ) + if ar_metadata_group is not None: + from lightx2v.models.networks.hunyuan_image3.custom_all_reduce import HunyuanImage3CustomAllReduce + + reducer = HunyuanImage3CustomAllReduce( + metadata_group=ar_metadata_group, + fallback_group=lambda: context.active_tp_group, + config=config, + device=torch.device(AI_DEVICE, torch.cuda.current_device()), + phase_getter=lambda: context.phase, + ) + context.ar_custom_all_reduce = reducer + reducer.initialize() + return context def initialize_hunyuan_image3_parallel_runtime(config) -> HunyuanImage3ParallelContext: diff --git a/lightx2v/models/networks/hunyuan_image3/weights/hybrid_tp.py b/lightx2v/models/networks/hunyuan_image3/weights/hybrid_tp.py index 49e08aca1..76bbf2d0a 100644 --- a/lightx2v/models/networks/hunyuan_image3/weights/hybrid_tp.py +++ b/lightx2v/models/networks/hunyuan_image3/weights/hybrid_tp.py @@ -324,7 +324,11 @@ def apply(self, input_tensor): output = torch.mm(input_tensor, weight) if self.reduce_output: if self.active_tp_size > 1 and self.active_tp_group is not None: - dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.active_tp_group) + phase_all_reduce = getattr(self.parallel_context, "tensor_parallel_all_reduce", None) + if callable(phase_all_reduce): + output = phase_all_reduce(output) + else: + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.active_tp_group) if self._row_split_bias is not None: output = output + self._row_split_bias return output diff --git a/lightx2v/models/runners/hunyuan_image3/cuda_graph.py b/lightx2v/models/runners/hunyuan_image3/cuda_graph.py new file mode 100644 index 000000000..75dd2165c --- /dev/null +++ b/lightx2v/models/runners/hunyuan_image3/cuda_graph.py @@ -0,0 +1,384 @@ +"""CUDA Graph execution for HunyuanImage3 autoregressive decoding.""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from loguru import logger + +from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache + + +def _config_bool(value, *, name): + if isinstance(value, bool): + return value + if isinstance(value, int) and value in (0, 1): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean, got {value!r}.") + + +def _tensor_signature(tensor: torch.Tensor | None): + if tensor is None: + return None + return ( + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device.type, + tensor.device.index, + ) + + +@dataclass(frozen=True) +class HunyuanImage3ARCudaGraphKey: + q_len: int + hidden_states: tuple + position_ids: tuple + rope_cos: tuple + rope_sin: tuple + cache_capacity: int + + +@dataclass +class _HunyuanImage3ARCudaGraphEntry: + key: HunyuanImage3ARCudaGraphKey + graph: torch.cuda.CUDAGraph + pre_infer_out: Any + hidden_states: torch.Tensor + position_ids: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + logits: torch.Tensor | None + + +class HunyuanImage3ARCudaGraphController: + """Own persistent buffers and capture the full q_len=1 AR forward.""" + + def __init__(self, config, model, device): + self.config = config + self.model = model + self.device = torch.device(device) + if self.device.type == "cuda" and self.device.index is None and torch.cuda.is_available(): + self.device = torch.device("cuda", torch.cuda.current_device()) + + self.enabled = _config_bool(config.get("enable_ar_cuda_graph", False), name="enable_ar_cuda_graph") + self.required = _config_bool(config.get("ar_cuda_graph_required", False), name="ar_cuda_graph_required") + self.mode = str(config.get("ar_cuda_graph_mode", "full_decode")).strip().lower() + self.kv_bucket_size = int(config.get("ar_cuda_graph_kv_bucket_size", 128)) + self.capture_q_lens = tuple(int(value) for value in config.get("ar_cuda_graph_capture_q_lens", [1])) + self.capture_warmups = int(config.get("ar_cuda_graph_capture_warmups", 2)) + self.decode_attn_impl = str(config.get("ar_decode_attn_impl", "disabled") or "disabled").strip().lower() + self.page_size = int(config.get("ar_kv_page_size", 16)) + self.max_num_splits = int(config.get("ar_flash_attn_max_num_splits", 32)) + + if self.required and not self.enabled: + raise ValueError("ar_cuda_graph_required=true requires enable_ar_cuda_graph=true.") + if self.enabled: + if self.mode != "full_decode": + raise ValueError(f"Unsupported HunyuanImage3 ar_cuda_graph_mode={self.mode!r}; expected 'full_decode'.") + if self.capture_q_lens != (1,): + raise ValueError("HunyuanImage3 full-decode CUDA Graph supports only ar_cuda_graph_capture_q_lens=[1].") + if self.decode_attn_impl != "flash_attn3_paged": + raise ValueError("HunyuanImage3 AR CUDA Graph requires ar_decode_attn_impl='flash_attn3_paged'.") + if self.device.type != "cuda": + raise ValueError(f"HunyuanImage3 AR CUDA Graph requires a CUDA device, got {self.device}.") + if self.kv_bucket_size < 1: + raise ValueError(f"ar_cuda_graph_kv_bucket_size must be positive, got {self.kv_bucket_size}.") + if self.capture_warmups < 1: + raise ValueError(f"ar_cuda_graph_capture_warmups must be positive, got {self.capture_warmups}.") + if self.page_size < 1: + raise ValueError(f"ar_kv_page_size must be positive, got {self.page_size}.") + if self.max_num_splits < 1: + raise ValueError(f"ar_flash_attn_max_num_splits must be positive, got {self.max_num_splits}.") + + self._entries: dict[HunyuanImage3ARCudaGraphKey, _HunyuanImage3ARCudaGraphEntry] = {} + self._disabled_keys: set[HunyuanImage3ARCudaGraphKey] = set() + self._kv_cache: HunyuanImage3StaticKVCache | None = None + self._pool = None + self._capture_stream = None + self._closed = False + + @staticmethod + def _round_up(value, multiple): + value = int(value) + multiple = int(multiple) + return ((value + multiple - 1) // multiple) * multiple + + def clear(self): + if self._closed: + return + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(self.device) + self._entries.clear() + self._disabled_keys.clear() + self._pool = None + self._capture_stream = None + + def close(self): + if self._closed: + return + self.clear() + self._kv_cache = None + self.model = None + context = self.config.get("parallel_context") + close_custom = getattr(context, "close_custom_all_reduce", None) + if callable(close_custom): + close_custom() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(self.device) + self._closed = True + + def acquire_kv_cache(self, *, num_layers, max_cache_len): + if not self.enabled: + raise RuntimeError("Cannot acquire a graph KV cache while AR CUDA Graph is disabled.") + configured_capacity = int(self.config.get("ar_cuda_graph_max_cache_len", 0) or 0) + requested_capacity = max(int(max_cache_len), configured_capacity) + capacity = self._round_up(requested_capacity, self.kv_bucket_size) + cache = self._kv_cache + if cache is None or cache.num_layers != int(num_layers) or cache.max_cache_len < capacity: + if cache is not None: + logger.warning( + "Growing HunyuanImage3 AR graph KV cache from {} to {}; cached graphs will be rebuilt.", + cache.max_cache_len, + capacity, + ) + self.clear() + cache = HunyuanImage3StaticKVCache( + num_layers=num_layers, + max_cache_len=capacity, + dynamic=True, + paged=True, + page_size=self.page_size, + ) + self._kv_cache = cache + return cache + + def is_target_decode(self, pre_infer_out): + hidden_states = getattr(pre_infer_out, "hidden_states", None) + return self.enabled and torch.is_tensor(hidden_states) and hidden_states.ndim == 3 and int(hidden_states.shape[1]) in self.capture_q_lens + + def _validate_runtime(self, pre_infer_out, valid_kv_len): + if self._closed: + raise RuntimeError("HunyuanImage3 AR CUDA Graph controller is closed.") + if not self.enabled: + raise RuntimeError("HunyuanImage3 AR CUDA Graph controller was called while disabled.") + if not torch.cuda.is_available(): + raise RuntimeError("HunyuanImage3 AR CUDA Graph requires CUDA.") + + hidden_states = getattr(pre_infer_out, "hidden_states", None) + position_ids = getattr(pre_infer_out, "position_ids", None) + custom_pos_emb = getattr(pre_infer_out, "custom_pos_emb", None) + cache = getattr(pre_infer_out, "past_key_values", None) + if not torch.is_tensor(hidden_states) or hidden_states.ndim != 3 or int(hidden_states.shape[1]) != 1: + raise RuntimeError(f"HunyuanImage3 AR CUDA Graph requires hidden_states [B, 1, H], got {getattr(hidden_states, 'shape', None)}.") + if int(hidden_states.shape[0]) != 1: + raise RuntimeError(f"HunyuanImage3 AR CUDA Graph requires batch size 1, got {hidden_states.shape[0]}.") + if not torch.is_tensor(position_ids) or tuple(position_ids.shape) != tuple(hidden_states.shape[:2]): + raise RuntimeError("HunyuanImage3 AR CUDA Graph requires position_ids matching the token axes.") + if not isinstance(custom_pos_emb, tuple) or len(custom_pos_emb) != 2 or not all(torch.is_tensor(value) for value in custom_pos_emb): + raise RuntimeError("HunyuanImage3 AR CUDA Graph requires a static (cos, sin) rotary embedding tuple.") + if not isinstance(cache, HunyuanImage3StaticKVCache) or cache is not self._kv_cache or not cache.paged: + raise RuntimeError("HunyuanImage3 AR CUDA Graph requires its persistent paged KV cache.") + if not bool(getattr(pre_infer_out, "use_cache", False)): + raise RuntimeError("HunyuanImage3 AR CUDA Graph requires use_cache=true.") + unsupported = { + name + for name in ( + "attention_mask", + "image_mask", + "timesteps", + "token_hw", + "full_attn_slices", + "sequence_parallel_state", + "attention_segment_specs", + ) + if getattr(pre_infer_out, name, None) is not None + } + if unsupported: + raise RuntimeError(f"HunyuanImage3 q_len=1 graph inputs contain unsupported fields: {sorted(unsupported)}.") + if int(valid_kv_len) > cache.max_cache_len: + raise RuntimeError(f"Valid KV length {valid_kv_len} exceeds graph cache capacity {cache.max_cache_len}.") + if any(layer.key is None or layer.value is None for layer in cache.layers): + raise RuntimeError("The eager AR prefill must allocate every paged KV layer before graph capture.") + + def _make_key(self, pre_infer_out): + cos, sin = pre_infer_out.custom_pos_emb + cache = pre_infer_out.past_key_values + return HunyuanImage3ARCudaGraphKey( + q_len=int(pre_infer_out.hidden_states.shape[1]), + hidden_states=_tensor_signature(pre_infer_out.hidden_states), + position_ids=_tensor_signature(pre_infer_out.position_ids), + rope_cos=_tensor_signature(cos), + rope_sin=_tensor_signature(sin), + cache_capacity=int(cache.max_cache_len), + ) + + def _active_tp_group(self): + context = self.config.get("parallel_context") + if context is not None: + return getattr(context, "active_tp_group", getattr(context, "tp_group", None)) + return getattr(self.model, "tp_group", None) + + def _active_tp_size(self): + context = self.config.get("parallel_context") + if context is not None: + return int(getattr(context, "active_tp_size", getattr(context, "tp_size", 1))) + group = self._active_tp_group() + if group is not None and dist.is_available() and dist.is_initialized(): + return int(dist.get_world_size(group)) + return 1 + + def _local_query_heads(self): + transformer = getattr(self.model, "transformer_infer", None) + global_heads = getattr(transformer, "global_num_heads", None) + if global_heads is None: + global_heads = self.config.get("num_attention_heads", self.config.get("num_heads")) + if global_heads is None: + raise RuntimeError("Could not resolve the HunyuanImage3 query-head count.") + tp_size = self._active_tp_size() + if int(global_heads) % tp_size: + raise RuntimeError(f"Attention heads {global_heads} are not divisible by active TP size {tp_size}.") + return int(global_heads) // tp_size + + def _copy_runtime_inputs(self, entry, pre_infer_out, valid_kv_len): + entry.hidden_states.copy_(pre_infer_out.hidden_states) + entry.position_ids.copy_(pre_infer_out.position_ids) + cos, sin = pre_infer_out.custom_pos_emb + entry.rope_cos.copy_(cos) + entry.rope_sin.copy_(sin) + pre_infer_out.past_key_values.prepare_paged_decode_scheduler( + valid_length=int(valid_kv_len), + num_query_heads=self._local_query_heads(), + max_num_splits=self.max_num_splits, + ) + + def _tp_barrier(self): + group = self._active_tp_group() + if group is None or not dist.is_available() or not dist.is_initialized() or dist.get_world_size(group) <= 1: + return + if "nccl" in str(dist.get_backend(group)).lower(): + dist.barrier(group=group, device_ids=[self.device.index]) + else: + dist.barrier(group=group) + + def _build_entry(self, key, pre_infer_out): + hidden_states = torch.empty_like(pre_infer_out.hidden_states) + position_ids = torch.empty_like(pre_infer_out.position_ids) + runtime_cos, runtime_sin = pre_infer_out.custom_pos_emb + rope_cos = torch.empty_like(runtime_cos) + rope_sin = torch.empty_like(runtime_sin) + static_pre_infer_out = type(pre_infer_out)( + hidden_states=hidden_states, + attention_mask=None, + position_ids=position_ids, + custom_pos_emb=(rope_cos, rope_sin), + past_key_values=pre_infer_out.past_key_values, + use_cache=True, + image_mask=None, + timesteps=None, + token_hw=None, + first_step=pre_infer_out.first_step, + full_attn_slices=None, + sequence_parallel_state=None, + attention_segment_specs=None, + ) + return _HunyuanImage3ARCudaGraphEntry( + key=key, + graph=torch.cuda.CUDAGraph(), + pre_infer_out=static_pre_infer_out, + hidden_states=hidden_states, + position_ids=position_ids, + rope_cos=rope_cos, + rope_sin=rope_sin, + logits=None, + ) + + def _capture(self, key, pre_infer_out, valid_kv_len): + entry = self._build_entry(key, pre_infer_out) + self._copy_runtime_inputs(entry, pre_infer_out, valid_kv_len) + if self._pool is None: + self._pool = torch.cuda.graph_pool_handle() + if self._capture_stream is None: + self._capture_stream = torch.cuda.Stream(device=self.device) + + current_stream = torch.cuda.current_stream(self.device) + self._capture_stream.wait_stream(current_stream) + self._tp_barrier() + context = self.config.get("parallel_context") + custom_capture = getattr(context, "custom_all_reduce_capture", None) + + with torch.cuda.stream(self._capture_stream): + for _ in range(self.capture_warmups): + self.model.infer_ar_prepared(entry.pre_infer_out)["logits"][:, -1, :] + self._capture_stream.synchronize() + self._tp_barrier() + + capture_context = custom_capture() if callable(custom_capture) else nullcontext() + with capture_context: + with torch.cuda.graph(entry.graph, pool=self._pool, stream=self._capture_stream): + entry.logits = self.model.infer_ar_prepared(entry.pre_infer_out)["logits"][:, -1, :] + current_stream.wait_stream(self._capture_stream) + self._tp_barrier() + self._entries[key] = entry + logger.info( + "Captured HunyuanImage3 AR CUDA Graph: q_len={} cache_capacity={} cached_graphs={} rank={}.", + key.q_len, + key.cache_capacity, + len(self._entries), + dist.get_rank() if dist.is_available() and dist.is_initialized() else 0, + ) + return entry + + def _prepare_paged_eager(self, pre_infer_out, valid_kv_len): + pre_infer_out.past_key_values.prepare_paged_decode_scheduler( + valid_length=int(valid_kv_len), + num_query_heads=self._local_query_heads(), + max_num_splits=self.max_num_splits, + ) + + def prepare_replay(self, pre_infer_out, *, valid_kv_len): + self._validate_runtime(pre_infer_out, valid_kv_len) + key = self._make_key(pre_infer_out) + if key in self._disabled_keys: + self._prepare_paged_eager(pre_infer_out, valid_kv_len) + return None + + if key not in self._entries: + try: + self._capture(key, pre_infer_out, valid_kv_len) + except Exception as error: + self._disabled_keys.add(key) + if self.required: + raise RuntimeError(f"Required HunyuanImage3 AR CUDA Graph capture failed for cache capacity {key.cache_capacity}.") from error + logger.exception( + "HunyuanImage3 AR graph capture failed for cache capacity {}; using eager paged FA3.", + key.cache_capacity, + ) + self._prepare_paged_eager(pre_infer_out, valid_kv_len) + return None + + entry = self._entries[key] + self._copy_runtime_inputs(entry, pre_infer_out, valid_kv_len) + return entry + + def run(self, pre_infer_out, *, valid_kv_len): + entry = self.prepare_replay(pre_infer_out, valid_kv_len=valid_kv_len) + if entry is None: + return self.model.infer_ar_prepared(pre_infer_out)["logits"][:, -1, :] + entry.graph.replay() + if entry.logits is None: + raise RuntimeError("HunyuanImage3 AR CUDA Graph replay completed without a captured logits tensor.") + return entry.logits + + +__all__ = ["HunyuanImage3ARCudaGraphController", "HunyuanImage3ARCudaGraphKey"] diff --git a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py index 069642542..d84305eae 100644 --- a/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py +++ b/lightx2v/models/runners/hunyuan_image3/hunyuan_image3_runner.py @@ -17,6 +17,7 @@ from lightx2v.models.networks.hunyuan_image3.infer.kv_cache import HunyuanImage3StaticKVCache from lightx2v.models.networks.hunyuan_image3.model import HunyuanImage3Model from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.hunyuan_image3.cuda_graph import HunyuanImage3ARCudaGraphController from lightx2v.models.runners.hunyuan_image3.flashinfer_autotune import ( DistributedAutotuneContext, FlashInferAutotuneController, @@ -363,6 +364,17 @@ def _hunyuan_text_kv_cache_enabled(self): return bool(self.config["enable_kv_cache"]) return hasattr(self, "hunyuan_config") + def _get_ar_cuda_graph_controller(self): + controller = getattr(self, "_hunyuan_ar_cuda_graph_controller", None) + if controller is None: + controller = HunyuanImage3ARCudaGraphController( + config=self.config, + model=self.model, + device=self._pipeline_latent_device(), + ) + self._hunyuan_ar_cuda_graph_controller = controller + return controller + def _hunyuan_num_layers(self): hunyuan_config = getattr(self, "hunyuan_config", None) fallback = self.config.get("num_hidden_layers", getattr(hunyuan_config, "num_hidden_layers", 1)) @@ -964,14 +976,24 @@ def _generate_text_tokens( pending_tokens = [] generated = [] use_kv_cache = self._hunyuan_text_kv_cache_enabled() + graph_controller = self._get_ar_cuda_graph_controller() + if graph_controller.enabled and not use_kv_cache: + raise ValueError("HunyuanImage3 enable_ar_cuda_graph=true requires enable_text_kv_cache=true.") kv_cache = None cache_filled_length = 0 if use_kv_cache: - kv_cache = HunyuanImage3StaticKVCache( - num_layers=self._hunyuan_num_layers(), - max_cache_len=input_ids.shape[1] + max_new_tokens + sum(len(tokens) for tokens in transition_map.values()), - dynamic=True, - ) + max_cache_len = input_ids.shape[1] + max_new_tokens + sum(len(tokens) for tokens in transition_map.values()) + if graph_controller.enabled: + kv_cache = graph_controller.acquire_kv_cache( + num_layers=self._hunyuan_num_layers(), + max_cache_len=max_cache_len, + ) + else: + kv_cache = HunyuanImage3StaticKVCache( + num_layers=self._hunyuan_num_layers(), + max_cache_len=max_cache_len, + dynamic=True, + ) for _ in range(max_new_tokens): if pending_tokens: @@ -999,7 +1021,17 @@ def _generate_text_tokens( else: model_inputs = self._build_text_model_inputs(input_ids, tokenizer_output, cond_inputs=cond_inputs, rope_image_info=rope_image_info) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): - logits = self.model.infer(model_inputs)["logits"][:, -1, :] + if graph_controller.enabled: + pre_infer_out = self.model.prepare_ar_pre_infer(model_inputs) + if not first_cache_step and graph_controller.is_target_decode(pre_infer_out): + logits = graph_controller.run( + pre_infer_out, + valid_kv_len=cache_filled_length, + ) + else: + logits = self.model.infer_ar_prepared(pre_infer_out)["logits"][:, -1, :] + else: + logits = self.model.infer(model_inputs)["logits"][:, -1, :] next_token = self._sample_text_token(logits, generator, generation_options=generation_options) next_token = next_token.to(device=device, dtype=input_ids.dtype) next_token = self._broadcast_parallel_tensor(next_token) @@ -1651,3 +1683,16 @@ def run_pipeline(self, input_info): images[0].save(save_result_path) logger.info(f"HunyuanImage3 image saved successfully to: {save_result_path}") return {"image": None} + + def close(self): + """Release graph and IPC resources before distributed teardown.""" + + controller = getattr(self, "_hunyuan_ar_cuda_graph_controller", None) + self._hunyuan_ar_cuda_graph_controller = None + if controller is not None: + controller.close() + return + context = self.config.get("parallel_context") + close_custom = getattr(context, "close_custom_all_reduce", None) + if callable(close_custom): + close_custom() diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh new file mode 100755 index 000000000..428d86f68 --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +# By default, keep the three repositories side by side. Override either external +# path when using a different workspace layout. +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="$(cd -- "${script_dir}/../.." && pwd)" +workspace_path="$(dirname -- "${lightx2v_path}")" +model_path="${HUNYUAN_IMAGE3_MODEL_PATH:-${workspace_path}/HunyuanImage-3-Instruct}" +hunyuan_image3_path="${HUNYUAN_IMAGE3_SOURCE_PATH:-${workspace_path}/HunyuanImage-3.0}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export PYTHONPATH="${hunyuan_image3_path}:${PYTHONPATH:-}" +export TORCH_NCCL_ASYNC_ERROR_HANDLING=0 + +source "${lightx2v_path}/scripts/base/base.sh" +export PROFILING_DEBUG_LEVEL=0 + +torchrun --standalone --nproc_per_node=4 -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task t2i \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json" \ + --prompt "生成图片：一辆汽车行驶在高速公路上，驾驶员在打电话，副驾驶坐着一只狗" \ + --save_result_path "${lightx2v_path}/save_results/hunyuan_image3_t2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.png" \ + --seed 42 diff --git a/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh new file mode 100755 index 000000000..9851911bc --- /dev/null +++ b/scripts/hunyuan_image3/run_hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# By default, keep the three repositories side by side. Override either external +# path when using a different workspace layout. +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="$(cd -- "${script_dir}/../.." && pwd)" +workspace_path="$(dirname -- "${lightx2v_path}")" +model_path="${HUNYUAN_IMAGE3_MODEL_PATH:-${workspace_path}/HunyuanImage-3-Instruct}" +hunyuan_image3_path="${HUNYUAN_IMAGE3_SOURCE_PATH:-${workspace_path}/HunyuanImage-3.0}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export PYTHONPATH="${hunyuan_image3_path}:${PYTHONPATH:-}" +export TORCH_NCCL_ASYNC_ERROR_HANDLING=0 + +source "${lightx2v_path}/scripts/base/base.sh" +export PROFILING_DEBUG_LEVEL=0 + +torchrun --standalone --nproc_per_node=4 -m lightx2v.infer \ + --model_cls hunyuan_image3 \ + --task ti2i \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/hunyuan_image3/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.json" \ + --prompt "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。鱼眼镜头，背景是房间门口，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清绒毛细节，面部特写，宝丽莱相纸，写实胶片摄影，复古颗粒感。" \ + --image_path "${model_path}/assets/demo_instruct_imgs/input_0_0.png" \ + --save_result_path "${lightx2v_path}/save_results/hunyuan_image3_ti2i_ar_tp4_cudagraph_denoise_fa3_tp2_sp2_multi_micro_flashinfer.png" \ + --seed 42