feat(hunyuan-image3): optimize four-GPU hybrid inference

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
+}

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
+}

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",
+]

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__":