This Hunyuan Image 3.0 version applies T2I, T2T, TI2I, and TI2T tasks.

configs/hunyuan_image3/hunyuan_image3_t2i.json → configs/hunyuan_image3/hunyuan_image3_t2i_tp_sp_cfg_flashinfer.json

@@ -6,12 +6,21 @@
     "target_width": 1024,
     "feature_caching": "NoCaching",
     "moe_impl": "flashinfer",
-    "attn_impl": "sdpa",
-    "hunyuan_cfg_mode": "batch",
-    "flashinfer_autotune_mode": "tune",
-    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2i.json",
+    "attn_impl": "torch_sdpa",
+    "flashinfer_autotune_mode": "auto",
+    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2i_tp2_sp2_cfg2.json",
     "flashinfer_tune_max_num_tokens": 16384,
-    "flashinfer_tuning_buckets": [128, 256, 512, 1024, 2048, 4096, 8192, 12288, 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,
@@ -24,7 +33,14 @@
     "taylor_cache_tailing_enhance_steps": 1,
     "taylor_cache_low_freqs_order": 2,
     "taylor_cache_high_freqs_order": 2,
-    "pipeline_parallel": true,
+    "parallel": {
+        "pipeline_parallel": false,
+        "tensor_p_size": 2,
+        "seq_p_size": 2,
+        "cfg_p_size": 2,
+        "seq_p_attn_type": "ulysses",
+        "cfg_mode": "parallel"
+    },
     "enable_cfg": true,
     "bot_task": "think_recaption",
     "use_system_prompt": "en_unified",

configs/hunyuan_image3/hunyuan_image3_t2t_tp_sp_cfg_flashinfer.json

@@ -0,0 +1,39 @@
+{
+    "feature_caching": "NoCaching",
+    "moe_impl": "flashinfer",
+    "attn_impl": "torch_sdpa",
+    "flashinfer_autotune_mode": "auto",
+    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_t2t_tp2_sp2_cfg1.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,
+    "parallel": {
+        "pipeline_parallel": false,
+        "tensor_p_size": 2,
+        "seq_p_size": 2,
+        "cfg_p_size": 1,
+        "seq_p_attn_type": "ulysses",
+        "cfg_mode": "batch"
+    },
+    "enable_cfg": false,
+    "bot_task": "auto",
+    "use_system_prompt": "None",
+    "max_new_tokens": 2048,
+    "text_do_sample": true,
+    "text_top_k": 1024,
+    "text_top_p": 0.95,
+    "text_temperature": 0.6,
+    "text_stream_output": true
+}

configs/hunyuan_image3/hunyuan_image3_i2i.json → configs/hunyuan_image3/hunyuan_image3_ti2i_tp_sp_cfg_flashinfer.json

@@ -6,11 +6,20 @@
     "feature_caching": "NoCaching",
     "moe_impl": "flashinfer",
     "attn_impl": "torch_sdpa",
-    "hunyuan_cfg_mode": "batch",
-    "flashinfer_autotune_mode": "tune",
-    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_i2i.json",
+    "flashinfer_autotune_mode": "auto",
+    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_ti2i_tp2_sp2_cfg2.json",
     "flashinfer_tune_max_num_tokens": 16384,
-    "flashinfer_tuning_buckets": [128, 256, 512, 1024, 2048, 4096, 8192, 12288, 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,
@@ -23,7 +32,14 @@
     "taylor_cache_tailing_enhance_steps": 1,
     "taylor_cache_low_freqs_order": 2,
     "taylor_cache_high_freqs_order": 2,
-    "pipeline_parallel": true,
+    "parallel": {
+        "pipeline_parallel": false,
+        "tensor_p_size": 2,
+        "seq_p_size": 2,
+        "cfg_p_size": 2,
+        "seq_p_attn_type": "ulysses",
+        "cfg_mode": "parallel"
+    },
     "enable_cfg": true,
     "infer_align_image_size": true,
     "bot_task": "think_recaption",

configs/hunyuan_image3/hunyuan_image3_ti2t_tp_sp_cfg_flashinfer.json

@@ -0,0 +1,40 @@
+{
+    "feature_caching": "NoCaching",
+    "moe_impl": "flashinfer",
+    "attn_impl": "torch_sdpa",
+    "flashinfer_autotune_mode": "auto",
+    "flashinfer_autotune_cache": "save_results/hunyuan_image3_flashinfer_autotune_ti2t_tp2_sp2_cfg1.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,
+    "parallel": {
+        "pipeline_parallel": false,
+        "tensor_p_size": 2,
+        "seq_p_size": 2,
+        "cfg_p_size": 1,
+        "seq_p_attn_type": "ulysses",
+        "cfg_mode": "batch"
+    },
+    "enable_cfg": false,
+    "infer_align_image_size": true,
+    "bot_task": "auto",
+    "use_system_prompt": "None",
+    "max_new_tokens": 2048,
+    "text_do_sample": true,
+    "text_top_k": 1024,
+    "text_top_p": 0.95,
+    "text_temperature": 0.6,
+    "text_stream_output": true
+}

lightx2v/common/ops/mm/mm_weight.py

@@ -2444,6 +2444,7 @@ def __init__(
         is_post_adapter=False,
         lora_prefix="diffusion_model.blocks",
         lora_path="",
+        reduce_output=True,
     ):
         super().__init__(
             weight_name,
@@ -2460,6 +2461,7 @@ def __init__(
         self.tp_rank = tp_rank
         self.tp_size = tp_size
         self.split_dim = split_dim  # "col" for column split, "row" for row split
+        self.reduce_output = reduce_output
         assert split_dim in ["col", "row"], f"split_dim must be 'col' or 'row', got {split_dim}"
 
         self._mm = MM_WEIGHT_REGISTER.get(mm_type, MMWeight)(
@@ -2497,7 +2499,7 @@ def apply(self, input_tensor):
         output = self._mm.apply(input_tensor)
 
         # For row split, need all-reduce to combine results from all ranks
-        if self.split_dim == "row" and self.tp_size > 1 and self.tp_group is not None:
+        if self.split_dim == "row" and self.reduce_output and self.tp_size > 1 and self.tp_group is not None:
             dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group)
             # Add bias after all-reduce (bias is not split for row split)
             if self._row_split_bias is not None:

lightx2v/disagg/examples/run_controller.py

@@ -32,12 +32,9 @@ def _parallel_world_size_from_config(config: dict[str, Any] | None) -> int:
         return 1
 
     tensor_p_size = int(parallel.get("tensor_p_size", 1) or 1)
-    if tensor_p_size > 1:
-        return tensor_p_size
-
     cfg_p_size = int(parallel.get("cfg_p_size", 1) or 1)
     seq_p_size = int(parallel.get("seq_p_size", 1) or 1)
-    return max(1, cfg_p_size * seq_p_size)
+    return max(1, tensor_p_size * cfg_p_size * seq_p_size)
 
 
 def _load_base_config_json(path: str) -> dict[str, Any]:

lightx2v/infer.py

@@ -138,7 +138,10 @@ def main():
     )
 
     parser.add_argument(
-        "--task", type=str, choices=["t2v", "i2v", "t2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "i2va", "v2av", "ltx2_s2v", "sr", "recon", "i23d"], default="t2v"
+        "--task",
+        type=str,
+        choices=["t2v", "i2v", "t2t", "t2i", "ti2t", "ti2i", "i2i", "flf2v", "vace", "animate", "s2v", "rs2v", "t2av", "i2av", "i2va", "v2av", "ltx2_s2v", "sr", "recon", "i23d"],
+        default="t2v",
     )
     parser.add_argument("--support_tasks", type=str, nargs="+", default=[], help="Set supported tasks for the model")
     parser.add_argument("--model_path", type=str, required=True)
@@ -152,7 +155,7 @@ def main():
         "--image_path",
         type=str,
         default="",
-        help="The path to input image file(s) for image-to-video (i2v) or image-to-audio-video (i2av) task. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'",
+        help="The path to input image file(s), including HunyuanImage3 ti2t/ti2i. Multiple paths should be comma-separated. Example: 'path1.jpg,path2.jpg'",
     )
     parser.add_argument("--state_path", type=str, default="", help="The path to input robot state file for robot i2v/i2va inference.")
     parser.add_argument("--last_frame_path", type=str, default="", help="The path to last frame file for first-last-frame-to-video (flf2v) task")

lightx2v/models/networks/hunyuan_image3/infer/post_infer.py

@@ -1,4 +1,5 @@
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 
 from lightx2v.models.networks.hunyuan_image3.infer.utils import apply_linear, apply_timestep_embedder, first_weight_device, to_device
@@ -7,6 +8,12 @@
 class HunyuanImage3PostInfer:
     def __init__(self, config):
         self.config = config
+        if config.get("tensor_parallel", False):
+            self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p")
+            self.tp_size = dist.get_world_size(self.tp_group)
+        else:
+            self.tp_group = None
+            self.tp_size = 1
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
@@ -52,6 +59,15 @@ def infer(self, weights, hidden_states, pre_infer_out):
                 image_output = hidden_states.masked_select(image_mask.unsqueeze(-1).bool()).reshape(-1, token_h * token_w, hidden)
             return {"diffusion_prediction": self._final_layer(weights, image_output, timesteps, token_h, token_w)}
 
-        normed = weights.final_norm.apply(hidden_states)
+        # Autoregressive generation only consumes the last position. Limiting
+        # the TP vocabulary projection to that token avoids materializing and
+        # gathering prompt-length logits on every rank.
+        logits_hidden_states = hidden_states[:, -1:, :] if self.tp_size > 1 else hidden_states
+        normed = weights.final_norm.apply(logits_hidden_states)
         logits = apply_linear(weights.lm_head, normed.reshape(-1, normed.shape[-1])).reshape(*normed.shape[:-1], -1)
+        if self.tp_size > 1:
+            local_logits = logits.contiguous()
+            gathered_logits = [torch.empty_like(local_logits) for _ in range(self.tp_size)]
+            dist.all_gather(gathered_logits, local_logits, group=self.tp_group)
+            logits = torch.cat(gathered_logits, dim=-1)
         return {"logits": logits}

lightx2v/models/networks/hunyuan_image3/infer/transformer_infer.py

@@ -37,10 +37,20 @@ def __init__(self, config):
         self.config = config
         self.num_layers = int(config.get("num_layers") or config["num_hidden_layers"])
         self.hidden_size = config["hidden_size"]
-        self.num_heads = int(config.get("num_attention_heads") or config["num_heads"])
-        self.num_key_value_heads = config.get("num_key_value_heads") or self.num_heads
-        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-        self.head_dim = config.get("attention_head_dim", self.hidden_size // self.num_heads)
+        self.global_num_heads = int(config.get("num_attention_heads") or config["num_heads"])
+        self.global_num_key_value_heads = int(config.get("num_key_value_heads") or self.global_num_heads)
+        self.num_key_value_groups = self.global_num_heads // self.global_num_key_value_heads
+        self.head_dim = config.get("attention_head_dim", self.hidden_size // self.global_num_heads)
+        if config.get("tensor_parallel", False):
+            self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p")
+            self.tp_rank = dist.get_rank(self.tp_group)
+            self.tp_size = dist.get_world_size(self.tp_group)
+        else:
+            self.tp_group = None
+            self.tp_rank = 0
+            self.tp_size = 1
+        self.num_heads = self.global_num_heads // self.tp_size
+        self.num_key_value_heads = self.global_num_key_value_heads // self.tp_size
         self.hidden_act = config.get("hidden_act", "silu")
         self.flashinfer_tune_max_num_tokens = int(config.get("flashinfer_tune_max_num_tokens", 8192))
         self.attn_impl = self._normalize_attention_impl(config.get("attn_impl", "torch_sdpa"))
@@ -551,7 +561,7 @@ def _sequence_parallel_gather_kv(self, key_states, value_states, state):
 
     def _sequence_parallel_ulysses_seq_to_head(self, query_states, key_states, value_states, state):
         if query_states.shape[0] != 1:
-            raise ValueError("HunyuanImage3 Ulysses expects batch size 1; use hunyuan_cfg_mode='serial' for pure SP or 'parallel' for CFG+SP.")
+            raise ValueError("HunyuanImage3 Ulysses expects batch size 1; use parallel.cfg_mode='serial' when cfg_p_size=1 (including TP+SP), or 'parallel' for CFG+SP.")
         query = all2all_seq2head(query_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group)
         key = all2all_seq2head(key_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group)
         value = all2all_seq2head(value_states[0].transpose(0, 1).contiguous(), group=self.seq_p_group)
@@ -577,11 +587,15 @@ def infer_mlp(self, phase, hidden_states):
         moe = phase.moe
         moe_impl = getattr(moe, "moe_impl", self.config.get("moe_impl", "eager"))
         if moe_impl == "flashinfer":
-            return self._infer_mlp_flashinfer(moe, hidden_states)
-        if moe_impl != "eager":
+            output = self._infer_mlp_flashinfer(moe, hidden_states)
+        elif moe_impl == "eager":
+            output = self._infer_mlp_eager(moe, hidden_states)
+        else:
             raise ValueError(f"Unsupported HunyuanImage3 moe_impl={moe_impl!r}. Expected 'eager' or 'flashinfer'.")
 
-        return self._infer_mlp_eager(moe, hidden_states)
+        if self.tp_size > 1:
+            dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group)
+        return output
 
     def _moe_easy_topk(self, moe, hidden_states):
         flat = hidden_states.reshape(-1, hidden_states.shape[-1])
@@ -634,6 +648,8 @@ def _infer_mlp_flashinfer(self, moe, hidden_states):
             compute_dtype,
             output=combined_output,
             quant_scales=None,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
             tune_max_num_tokens=self.flashinfer_tune_max_num_tokens,
         )
         output = combined_output.reshape_as(hidden_states).to(original_dtype)

lightx2v/models/networks/hunyuan_image3/infer/utils.py

@@ -18,7 +18,13 @@ def _linear_tensor_attr(linear, attr):
     tensor = getattr(linear, attr, None)
     if tensor is not None:
         return tensor
-    return getattr(linear, f"pin_{attr}", None)
+    tensor = getattr(linear, f"pin_{attr}", None)
+    if tensor is not None:
+        return tensor
+    wrapped_linear = getattr(linear, "_mm", None)
+    if wrapped_linear is not None and wrapped_linear is not linear:
+        return _linear_tensor_attr(wrapped_linear, attr)
+    return None
 
 
 def match_linear_dtype(input_tensor, linear):
@@ -67,8 +73,8 @@ def repeat_kv(hidden_states, n_rep):
 
 
 def weight_device(module):
-    for attr in ("weight", "pin_weight", "bias", "pin_bias"):
-        tensor = getattr(module, attr, None)
+    for attr in ("weight", "bias"):
+        tensor = _linear_tensor_attr(module, attr)
         if tensor is not None:
             return tensor.device
     return None