From a816068a8b6b78d5d501cb104de9bc145ad22f71 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 3 Jul 2026 11:33:44 +0800 Subject: [PATCH 1/5] Optimize Hunyuan3D transformer sync and MoE --- configs/hunyuan3d/hunyuan3d_shape.json | 5 ++ .../hunyuan3d/infer/moe_fi_autotune.py | 50 +++++++++++++++++++ .../networks/hunyuan3d/infer/moe_infer.py | 28 +++++++++++ .../hunyuan3d/infer/transformer_infer.py | 30 ++++++++--- .../hunyuan3d/weights/transformer_weights.py | 46 +++++++++++++++++ .../hunyuan3d/hunyuan3d_shape_runner.py | 28 ++++++++--- 6 files changed, 174 insertions(+), 13 deletions(-) create mode 100644 lightx2v/models/networks/hunyuan3d/infer/moe_fi_autotune.py diff --git a/configs/hunyuan3d/hunyuan3d_shape.json b/configs/hunyuan3d/hunyuan3d_shape.json index 3dd41d8d2..130998b91 100644 --- a/configs/hunyuan3d/hunyuan3d_shape.json +++ b/configs/hunyuan3d/hunyuan3d_shape.json @@ -8,5 +8,10 @@ "num_chunks": 8000, "octree_resolution": 384, "attn_type": "torch_sdpa", + "moe_backend": "flashinfer", + "moe_flashinfer_setting": { + "autotune": true, + "tune_max_num_tokens": 8192 + }, "rms_norm_type": "torch" } diff --git a/lightx2v/models/networks/hunyuan3d/infer/moe_fi_autotune.py b/lightx2v/models/networks/hunyuan3d/infer/moe_fi_autotune.py new file mode 100644 index 000000000..1c57a95d7 --- /dev/null +++ b/lightx2v/models/networks/hunyuan3d/infer/moe_fi_autotune.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass + +from lightx2v.common.flashinfer_autotune import ( + FlashInferAutotune, + fi_autotune_cache_path, +) + +MOE_FI_CACHE_NAMESPACE = "hunyuan3d_moe" +MOE_FI_FORCE_RETUNE_ENV = "LIGHTX2V_HUNYUAN3D_MOE_FI_FORCE_RETUNE" + + +def _moe_intermediate(config) -> int: + hidden = int(config["hidden_size"]) + for key in ("moe_intermediate_size", "intermediate_size"): + if config.get(key): + return int(config[key]) + return int(hidden * float(config.get("mlp_ratio", 4))) + + +def build_moe_model_sig(config) -> str: + hidden = int(config["hidden_size"]) + intermediate = _moe_intermediate(config) + num_experts = int(config.get("num_experts", 8)) + top_k = int(config.get("moe_top_k", 2)) + return f"hunyuan3d_moe_e{num_experts}_k{top_k}_h{hidden}_i{intermediate}_gelu_bias" + + +def moe_fi_autotune_cache(config) -> str: + return fi_autotune_cache_path(MOE_FI_CACHE_NAMESPACE, build_moe_model_sig(config)) + + +@dataclass +class MoeFiAutotune(FlashInferAutotune): + tune_max_num_tokens: int = 8192 + + @classmethod + def from_hunyuan3d_config(cls, config) -> "MoeFiAutotune": + fi_cfg = config.get("moe_flashinfer_setting") or {} + tune_max = int(fi_cfg.get("tune_max_num_tokens", 8192)) + if str(config.get("moe_backend", "pytorch")).strip().lower() != "flashinfer": + return cls(tune_max_num_tokens=tune_max) + if not fi_cfg.get("autotune", False): + return cls(tune_max_num_tokens=tune_max) + return cls( + enabled=True, + cache_path=moe_fi_autotune_cache(config), + tune_max_num_tokens=tune_max, + force_retune_env=MOE_FI_FORCE_RETUNE_ENV, + log_prefix="Hunyuan3D Flashinfer MoE autotune", + ) diff --git a/lightx2v/models/networks/hunyuan3d/infer/moe_infer.py b/lightx2v/models/networks/hunyuan3d/infer/moe_infer.py index de6918d9e..b9068e8cf 100644 --- a/lightx2v/models/networks/hunyuan3d/infer/moe_infer.py +++ b/lightx2v/models/networks/hunyuan3d/infer/moe_infer.py @@ -1,6 +1,13 @@ import torch import torch.nn.functional as F +try: + from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe + from flashinfer.tllm_enums import ActivationType as FlashInferActivationType +except ImportError: + flashinfer_cutlass_fused_moe = None + FlashInferActivationType = None + @torch.no_grad() def infer_moe_ffn(ffn_weights, hidden_states): @@ -23,6 +30,27 @@ def infer_moe_block(moe_weights, hidden_states): flat_topk_idx = topk_idx.reshape(-1) flat_topk_weight = topk_weight.reshape(-1, 1) + if moe_weights.moe_backend == "flashinfer": + if flashinfer_cutlass_fused_moe is None: + raise RuntimeError("Hunyuan3D moe_backend=flashinfer but flashinfer.fused_moe is not available") + if not hasattr(moe_weights, "_fi_fc1_weight"): + moe_weights._build_flashinfer_weights() + routed = flashinfer_cutlass_fused_moe( + flat if flat.is_contiguous() else flat.contiguous(), + topk_idx.to(torch.int32), + topk_weight.to(torch.float32), + moe_weights._fi_fc1_weight, + moe_weights._fi_fc2_weight, + flat.dtype, + quant_scales=None, + fc1_expert_biases=moe_weights._fi_fc1_bias, + fc2_expert_biases=moe_weights._fi_fc2_bias, + tune_max_num_tokens=moe_weights.moe_flashinfer_tune_max_num_tokens, + activation_type=FlashInferActivationType.Gelu, + )[0].view(bsz, seq_len, hidden_dim) + shared = infer_moe_ffn(moe_weights.shared_experts, flat).view(bsz, seq_len, hidden_dim) + return routed + shared + expert_cache = torch.zeros_like(flat) idxs = flat_topk_idx.argsort() tokens_per_expert = flat_topk_idx.bincount(minlength=moe_weights.num_experts).cpu().numpy().cumsum(0) diff --git a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py index a51c2f99e..39dae59f5 100644 --- a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py @@ -1,8 +1,16 @@ +import os + import torch import torch.nn.functional as F +from loguru import logger +from lightx2v.common.flashinfer_autotune import flashinfer_autotune from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.hunyuan3d.infer.module_io import Hunyuan3DPreInferOutput +from lightx2v.models.networks.hunyuan3d.infer.moe_fi_autotune import ( + MOE_FI_FORCE_RETUNE_ENV, + MoeFiAutotune, +) from lightx2v.models.networks.hunyuan3d.infer.moe_infer import infer_moe_block @@ -15,6 +23,16 @@ def __init__(self, config): self.num_heads = config["num_heads"] self.head_dim = config["hidden_size"] // self.num_heads self.scheduler = None + self.fi_moe_autotune = MoeFiAutotune.from_hunyuan3d_config(config) + if self.fi_moe_autotune.enabled: + if flashinfer_autotune is None: + raise RuntimeError("Hunyuan3D FlashInfer MoE autotune enabled but flashinfer autotuner is not available") + logger.info( + f"Hunyuan3D FlashInfer MoE autotune enabled " + f"(cache={self.fi_moe_autotune.cache_path}, " + f"tune_max_num_tokens={self.fi_moe_autotune.tune_max_num_tokens}, " + f"{MOE_FI_FORCE_RETUNE_ENV}={os.environ.get(MOE_FI_FORCE_RETUNE_ENV, '0')})" + ) def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -70,8 +88,8 @@ def _run_attention(self, query, key, value, calculate, merge_batch=False): v = value[0] seqlen_q = q.shape[0] seqlen_k = k.shape[0] - cu_seqlens_q = torch.tensor([0, q.shape[0]], dtype=torch.int32, device=q.device) - cu_seqlens_k = torch.tensor([0, k.shape[0]], dtype=torch.int32, device=k.device) + cu_seqlens_q = torch.tensor([0, q.shape[0]], dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, k.shape[0]], dtype=torch.int32) attn_output = calculate.apply( q=q, k=k, @@ -90,8 +108,8 @@ def _run_attention(self, query, key, value, calculate, merge_batch=False): v = value.reshape(-1, self.num_heads, self.head_dim) seqlen_q = q.shape[0] // batch_size seqlen_k = k.shape[0] // batch_size - cu_seqlens_q = torch.tensor([0, q.shape[0]], dtype=torch.int32, device=q.device) - cu_seqlens_k = torch.tensor([0, k.shape[0]], dtype=torch.int32, device=k.device) + cu_seqlens_q = torch.tensor([0, q.shape[0]], dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, k.shape[0]], dtype=torch.int32) attn_output = calculate.apply( q=q, k=k, @@ -109,8 +127,8 @@ def _run_attention(self, query, key, value, calculate, merge_batch=False): q = query.reshape(-1, self.num_heads, self.head_dim) k = key.reshape(-1, self.num_heads, self.head_dim) v = value.reshape(-1, self.num_heads, self.head_dim) - cu_seqlens_q = torch.arange(0, batch_size + 1, dtype=torch.int32, device=q.device) * seqlen_q - cu_seqlens_k = torch.arange(0, batch_size + 1, dtype=torch.int32, device=k.device) * seqlen_k + cu_seqlens_q = torch.arange(0, batch_size + 1, dtype=torch.int32) * seqlen_q + cu_seqlens_k = torch.arange(0, batch_size + 1, dtype=torch.int32) * seqlen_k attn_output = calculate.apply( q=q, k=k, diff --git a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py index 189d84497..aa48017f1 100644 --- a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py +++ b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py @@ -1,3 +1,5 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER @@ -49,6 +51,13 @@ def __init__(self, config, block_idx, mm_type): num_experts = config.get("num_experts", 8) self.num_experts = num_experts self.moe_top_k = config.get("moe_top_k", 2) + self.moe_backend = str(config.get("moe_backend", "pytorch")).strip().lower() + if self.moe_backend not in ("pytorch", "flashinfer"): + raise ValueError(f"Invalid Hunyuan3D moe_backend={self.moe_backend!r}, expected 'pytorch' or 'flashinfer'") + fi_cfg = config.get("moe_flashinfer_setting") or {} + if fi_cfg.get("autotune") and self.moe_backend != "flashinfer": + raise ValueError("moe_flashinfer_setting.autotune=true requires moe_backend='flashinfer'") + self.moe_flashinfer_tune_max_num_tokens = int(fi_cfg.get("tune_max_num_tokens", 8192)) self.add_module( "gate", MM_WEIGHT_REGISTER[mm_type]( @@ -63,6 +72,43 @@ def __init__(self, config, block_idx, mm_type): experts = WeightModuleList(Hunyuan3DFeedForwardWeights(f"{prefix}.experts.{expert_idx}", mm_type) for expert_idx in range(num_experts)) self.add_module("experts", experts) + def load(self, weight_dict): + super().load(weight_dict) + if self.moe_backend == "flashinfer" and self.experts[0].fc1._get_actual_weight() is not None: + self._build_flashinfer_weights() + + def to_cuda(self, non_blocking=False): + super().to_cuda(non_blocking=non_blocking) + if self.moe_backend == "flashinfer": + self._build_flashinfer_weights() + + def to_cpu(self, non_blocking=False): + super().to_cpu(non_blocking=non_blocking) + for attr in ("_fi_fc1_weight", "_fi_fc2_weight", "_fi_fc1_bias", "_fi_fc2_bias"): + if hasattr(self, attr): + delattr(self, attr) + + @staticmethod + def _stack_optional_biases(biases): + if all(bias is None for bias in biases): + return None + if any(bias is None for bias in biases): + raise ValueError("FlashInfer Hunyuan3D MoE requires either all expert biases or no expert biases") + return torch.stack([bias.contiguous() for bias in biases], dim=0) + + def _build_flashinfer_weights(self): + fc1_list, fc2_list = [], [] + fc1_biases, fc2_biases = [], [] + for expert_w in self.experts: + fc1_list.append(expert_w.fc1._get_actual_weight().t().contiguous()) + fc2_list.append(expert_w.fc2._get_actual_weight().t().contiguous()) + fc1_biases.append(expert_w.fc1._get_actual_bias()) + fc2_biases.append(expert_w.fc2._get_actual_bias()) + self._fi_fc1_weight = torch.stack(fc1_list, dim=0) + self._fi_fc2_weight = torch.stack(fc2_list, dim=0) + self._fi_fc1_bias = self._stack_optional_biases(fc1_biases) + self._fi_fc2_bias = self._stack_optional_biases(fc2_biases) + class Hunyuan3DSelfAttentionWeights(WeightModule): def __init__(self, config, block_idx, mm_type, ln_type, rms_norm_type, attn_type, qkv_bias): diff --git a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py index 9d44c91f9..d9f9e790a 100644 --- a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py +++ b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py @@ -154,6 +154,14 @@ def _cast_cond_dtype(cond, dtype): return cond.to(dtype=dtype) return {k: Hunyuan3DShapeRunner._cast_cond_dtype(v, dtype) for k, v in cond.items()} + def _run_infer_step(self, step_index: int, dit_inputs: dict) -> None: + with ProfilingContext4DebugL1("step_pre"): + self.scheduler.step_pre(step_index) + with ProfilingContext4DebugL1("infer_main"): + self.model.infer(dit_inputs) + with ProfilingContext4DebugL1("step_post"): + self.scheduler.step_post() + @ProfilingContext4DebugL2("Run DiT") def run_main(self): latent_shape = (self.inputs["image_tensor"].shape[0], *self.vae_decoder.vae.latent_shape) @@ -178,13 +186,19 @@ def run_main(self): iterator = tqdm(iterator, desc="Diffusion Sampling:") - for step_index in iterator: - with ProfilingContext4DebugL1("step_pre"): - self.scheduler.step_pre(step_index) - with ProfilingContext4DebugL1("infer_main"): - self.model.infer(dit_inputs) - with ProfilingContext4DebugL1("step_post"): - self.scheduler.step_post() + at = self.model.transformer_infer.fi_moe_autotune + start_step = 0 + if at.cache_rebuild_needed(): + logger.info("Hunyuan3D FlashInfer MoE autotune: cache rebuild required; tuning on step 1 only, then cache-only for remaining steps") + with at.session(tune_mode=True): + self._run_infer_step(0, dit_inputs) + start_step = 1 + + with at.session(tune_mode=False): + for step_index in iterator: + if step_index < start_step: + continue + self._run_infer_step(step_index, dit_inputs) mesh = self._export_mesh(self.scheduler.latents) save_result_path = self.input_info.save_result_path From 49f61543c1f47876003060cf84ca700ed115b1f2 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 3 Jul 2026 15:49:12 +0800 Subject: [PATCH 2/5] Fuse Hunyuan3D attention small ops --- configs/hunyuan3d/hunyuan3d_shape.json | 1 + lightx2v/common/ops/norm/rms_norm_weight.py | 51 +++++++++++- lightx2v/common/ops/norm/triton_ops.py | 78 +++++++++++++++++++ .../hunyuan3d/infer/transformer_infer.py | 16 ++-- 4 files changed, 136 insertions(+), 10 deletions(-) diff --git a/configs/hunyuan3d/hunyuan3d_shape.json b/configs/hunyuan3d/hunyuan3d_shape.json index 130998b91..cba7d60ee 100644 --- a/configs/hunyuan3d/hunyuan3d_shape.json +++ b/configs/hunyuan3d/hunyuan3d_shape.json @@ -8,6 +8,7 @@ "num_chunks": 8000, "octree_resolution": 384, "attn_type": "torch_sdpa", + "use_fused_qk_rms_norm": true, "moe_backend": "flashinfer", "moe_flashinfer_setting": { "autotune": true, diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 340b747aa..48ff130e6 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -5,7 +5,12 @@ from loguru import logger from safetensors import safe_open -from lightx2v.common.ops.norm.triton_ops import fused_norm_3drope, fused_qk_norm_3drope, rms_norm_kernel +from lightx2v.common.ops.norm.triton_ops import ( + fused_norm_3drope, + fused_qk_norm_3drope, + fused_qk_rms_norm, + rms_norm_kernel, +) from lightx2v.common.ops.utils import * from lightx2v.utils.envs import * from lightx2v.utils.registry_factory import RMS_WEIGHT_REGISTER @@ -427,6 +432,50 @@ def apply(self, input_tensor): return rms_norm_kernel(input_tensor, w, self.eps) +def apply_qk_rms_norm( + query: torch.Tensor, + key: torch.Tensor, + norm_q, + norm_k, + *, + use_triton: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + if norm_q is None and norm_k is None: + return query, key + + if ( + use_triton + and norm_q is not None + and norm_k is not None + and norm_q.eps == norm_k.eps + and query.is_cuda + and key.is_cuda + and query.shape[-1] == key.shape[-1] + ): + q_shape = query.shape + k_shape = key.shape + head_dim = q_shape[-1] + q_flat = query.reshape(-1, head_dim) + k_flat = key.reshape(-1, head_dim) + q_flat, k_flat = fused_qk_rms_norm( + q_flat, + k_flat, + norm_q._get_actual_weight(), + norm_k._get_actual_weight(), + norm_q.eps, + match_torch_rms_cast=True, + ) + return q_flat.reshape(q_shape), k_flat.reshape(k_shape) + + if norm_q is not None: + q_shape = query.shape + query = norm_q.apply(query.reshape(-1, q_shape[-1])).reshape(q_shape) + if norm_k is not None: + k_shape = key.shape + key = norm_k.apply(key.reshape(-1, k_shape[-1])).reshape(k_shape) + return query, key + + class RMSWeightFusedQKNorm3DRope: """ Holds two pairs of dual-RMSNorm weights (Q and K) and applies diff --git a/lightx2v/common/ops/norm/triton_ops.py b/lightx2v/common/ops/norm/triton_ops.py index b73ea505d..cd5e93c24 100644 --- a/lightx2v/common/ops/norm/triton_ops.py +++ b/lightx2v/common/ops/norm/triton_ops.py @@ -980,6 +980,84 @@ def rms_norm_kernel(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): return y +@triton.jit +def _fused_qk_rms_norm_kernel( + q_ptr, + k_ptr, + w_q_ptr, + w_k_ptr, + n_q, + n_k, + D, + eps, + BLOCK_SIZE_DIM: tl.constexpr, + BLOCK_SIZE_SEQ: tl.constexpr, + MATCH_TORCH_RMS_CAST: tl.constexpr, +): + pid = tl.program_id(0) + n_q_tiles = tl.cdiv(n_q, BLOCK_SIZE_SEQ) + is_k = pid >= n_q_tiles + tile_id = pid - n_q_tiles if is_k else pid + n_rows = n_k if is_k else n_q + base_ptr = k_ptr if is_k else q_ptr + w_ptr = w_k_ptr if is_k else w_q_ptr + + row_start = tile_id * BLOCK_SIZE_SEQ + rows = row_start + tl.arange(0, BLOCK_SIZE_SEQ) + row_mask = rows < n_rows + + d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] + d_mask = d_offset < D + x_blk = base_ptr + rows[:, None] * D + d_offset + mask = row_mask[:, None] & d_mask + + x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32) + mean_square = tl.sum(x * x, axis=1, keep_dims=True) / D + rstd = tl.math.rsqrt(mean_square + eps) + w = tl.load(w_ptr + d_offset, mask=d_mask) + if MATCH_TORCH_RMS_CAST: + out = (x * rstd).to(w.dtype) * w + else: + out = (x * rstd * w.to(tl.float32)).to(w.dtype) + tl.store(x_blk, out, mask=mask) + + +def fused_qk_rms_norm( + q: torch.Tensor, + k: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + eps: float = 1e-6, + *, + match_torch_rms_cast: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """In-place RMSNorm on q [Nq, D] and k [Nk, D] in one Triton launch.""" + if not q.is_cuda or not k.is_cuda: + raise RuntimeError("fused_qk_rms_norm requires CUDA tensors") + q = q.contiguous() + k = k.contiguous() + n_q, d = q.shape + n_k = k.shape[0] + block_d = triton.next_power_of_2(d) + block_s = min(16, triton.next_power_of_2(max(1, max(n_q, n_k) // 512))) + grid = (triton.cdiv(n_q, block_s) + triton.cdiv(n_k, block_s),) + with torch.cuda.device(q.device): + torch.library.wrap_triton(_fused_qk_rms_norm_kernel)[grid]( + q, + k, + w_q, + w_k, + n_q, + n_k, + d, + eps, + BLOCK_SIZE_DIM=block_d, + BLOCK_SIZE_SEQ=block_s, + MATCH_TORCH_RMS_CAST=match_torch_rms_cast, + ) + return q, k + + # --------------------------------------------------------------------------- # Fused dual-RMSNorm + 3D Neox-RoPE (NeoPP Q/K, in-place, bfloat16) # --------------------------------------------------------------------------- diff --git a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py index 39dae59f5..13b223a0d 100644 --- a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py @@ -5,6 +5,7 @@ from loguru import logger from lightx2v.common.flashinfer_autotune import flashinfer_autotune +from lightx2v.common.ops.norm.rms_norm_weight import apply_qk_rms_norm from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.hunyuan3d.infer.module_io import Hunyuan3DPreInferOutput from lightx2v.models.networks.hunyuan3d.infer.moe_fi_autotune import ( @@ -23,6 +24,7 @@ def __init__(self, config): self.num_heads = config["num_heads"] self.head_dim = config["hidden_size"] // self.num_heads self.scheduler = None + self.use_fused_qk_rms_norm = bool(config.get("use_fused_qk_rms_norm", False)) self.fi_moe_autotune = MoeFiAutotune.from_hunyuan3d_config(config) if self.fi_moe_autotune.enabled: if flashinfer_autotune is None: @@ -67,6 +69,9 @@ def _reshape_cross_kv(k, v, num_heads, head_dim, batch_size, seq_len): v.reshape(batch_size, seq_len, num_heads, head_dim), ) + def _apply_qk_norm(self, query, key, norm_q, norm_k): + return apply_qk_rms_norm(query, key, norm_q, norm_k, use_triton=self.use_fused_qk_rms_norm) + def _project_qkv(self, hidden_states, to_q, to_k, to_v, norm_q, norm_k): batch_size, seq_len, hidden_dim = hidden_states.shape flat = hidden_states.reshape(-1, hidden_dim) @@ -74,10 +79,7 @@ def _project_qkv(self, hidden_states, to_q, to_k, to_v, norm_q, norm_k): k = to_k.apply(flat).reshape(batch_size, seq_len, hidden_dim) v = to_v.apply(flat).reshape(batch_size, seq_len, hidden_dim) query, key, value = self._reshape_self_qkv(q, k, v, self.num_heads, self.head_dim, batch_size, seq_len) - if norm_q is not None: - query = norm_q.apply(query.reshape(-1, self.head_dim)).reshape(batch_size, seq_len, self.num_heads, self.head_dim) - if norm_k is not None: - key = norm_k.apply(key.reshape(-1, self.head_dim)).reshape(batch_size, seq_len, self.num_heads, self.head_dim) + query, key = self._apply_qk_norm(query, key, norm_q, norm_k) return query, key, value def _run_attention(self, query, key, value, calculate, merge_batch=False): @@ -169,11 +171,7 @@ def _infer_cross_attention(self, block_weights, hidden_states, cond): k = k_flat.reshape(batch_size, cond_len, self.num_heads * self.head_dim) v = v_flat.reshape(batch_size, cond_len, self.num_heads * self.head_dim) key, value = self._reshape_cross_kv(k, v, self.num_heads, self.head_dim, batch_size, cond_len) - - if block_weights.attn2.norm_q is not None: - query = block_weights.attn2.norm_q.apply(query.reshape(-1, self.head_dim)).reshape(batch_size, seq_len, self.num_heads, self.head_dim) - if block_weights.attn2.norm_k is not None: - key = block_weights.attn2.norm_k.apply(key.reshape(-1, self.head_dim)).reshape(batch_size, cond_len, self.num_heads, self.head_dim) + query, key = self._apply_qk_norm(query, key, block_weights.attn2.norm_q, block_weights.attn2.norm_k) attn_output = self._run_attention(query, key, value, block_weights.attn2.calculate, merge_batch=False) flat = attn_output.reshape(-1, attn_output.shape[-1]) From 4ee7a348f769758005958fd52b1266b199b96bda Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Fri, 3 Jul 2026 16:54:16 +0800 Subject: [PATCH 3/5] Add Hunyuan3D FP8 DiT support without FP8 MoE Quantize non-MoE Hunyuan3D DiT linear weights through fp8-sgl while keeping MoE weights in BF16. Reject checkpoints that still contain FP8 MoE weights. --- configs/hunyuan3d/hunyuan3d_shape_fp8.json | 20 +++++++++++++ .../networks/hunyuan3d/infer/pre_infer.py | 2 ++ lightx2v/models/networks/hunyuan3d/model.py | 5 +++- .../hunyuan3d/weights/transformer_weights.py | 29 +++++++++++++++++-- .../hunyuan3d/hunyuan3d_shape_runner.py | 12 ++++---- 5 files changed, 58 insertions(+), 10 deletions(-) create mode 100644 configs/hunyuan3d/hunyuan3d_shape_fp8.json diff --git a/configs/hunyuan3d/hunyuan3d_shape_fp8.json b/configs/hunyuan3d/hunyuan3d_shape_fp8.json new file mode 100644 index 000000000..e0809ba78 --- /dev/null +++ b/configs/hunyuan3d/hunyuan3d_shape_fp8.json @@ -0,0 +1,20 @@ +{ + "infer_steps": 50, + "guidance_scale": 5.0, + "enable_rembg": true, + "enable_pbar": true, + "box_v": 1.01, + "mc_level": 0.0, + "num_chunks": 8000, + "octree_resolution": 384, + "attn_type": "torch_sdpa", + "use_fused_qk_rms_norm": true, + "moe_backend": "flashinfer", + "moe_flashinfer_setting": { + "autotune": true, + "tune_max_num_tokens": 8192 + }, + "dit_original_ckpt": "/data/nvme0/wangyingrui/Hunyuan3D-2.1-fp8/hunyuan3d-dit-v2-1/model.fp8.safetensors", + "dit_quant_scheme": "fp8-sgl", + "rms_norm_type": "torch" +} diff --git a/lightx2v/models/networks/hunyuan3d/infer/pre_infer.py b/lightx2v/models/networks/hunyuan3d/infer/pre_infer.py index 79e3d3348..7948f38cc 100644 --- a/lightx2v/models/networks/hunyuan3d/infer/pre_infer.py +++ b/lightx2v/models/networks/hunyuan3d/infer/pre_infer.py @@ -19,6 +19,8 @@ def set_scheduler(self, scheduler): def infer(self, weights, hidden_states, cond, timestep, guidance_cond=None): t_freq = apply_timesteps_embedding(timestep, self.config["hidden_size"]) weight_dtype = weights.t_embedder_mlp_0.weight.dtype + if weight_dtype == torch.float8_e4m3fn: + weight_dtype = hidden_states.dtype t_freq = t_freq.to(dtype=weight_dtype) if guidance_cond is not None and weights.t_embedder_cond_proj is not None: t_freq = t_freq + weights.t_embedder_cond_proj.apply(guidance_cond) diff --git a/lightx2v/models/networks/hunyuan3d/model.py b/lightx2v/models/networks/hunyuan3d/model.py index 6c39ee69d..b9457fdf6 100644 --- a/lightx2v/models/networks/hunyuan3d/model.py +++ b/lightx2v/models/networks/hunyuan3d/model.py @@ -95,7 +95,10 @@ def _build_weight_dict(state_dict, unified_dtype, sensitive_layer, config=None): for key, tensor in state_dict.items(): if config is None and not (unified_dtype or all(s not in key for s in sensitive_layer)): dtype = GET_SENSITIVE_DTYPE() - weight_dict[key] = tensor.to(dtype=dtype) + if tensor.dtype == torch.float8_e4m3fn or key.endswith(".weight_scale"): + weight_dict[key] = tensor + else: + weight_dict[key] = tensor.to(dtype=dtype) return weight_dict @property diff --git a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py index aa48017f1..4694b483b 100644 --- a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py +++ b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py @@ -49,6 +49,7 @@ def __init__(self, config, block_idx, mm_type): super().__init__() prefix = f"blocks.{block_idx}.moe" num_experts = config.get("num_experts", 8) + moe_mm_type = "Default" if str(mm_type).startswith("fp8") else mm_type self.num_experts = num_experts self.moe_top_k = config.get("moe_top_k", 2) self.moe_backend = str(config.get("moe_backend", "pytorch")).strip().lower() @@ -60,23 +61,45 @@ def __init__(self, config, block_idx, mm_type): self.moe_flashinfer_tune_max_num_tokens = int(fi_cfg.get("tune_max_num_tokens", 8192)) self.add_module( "gate", - MM_WEIGHT_REGISTER[mm_type]( + MM_WEIGHT_REGISTER[moe_mm_type]( f"{prefix}.gate.weight", None, ), ) self.add_module( "shared_experts", - Hunyuan3DFeedForwardWeights(f"{prefix}.shared_experts", mm_type), + Hunyuan3DFeedForwardWeights(f"{prefix}.shared_experts", moe_mm_type), ) - experts = WeightModuleList(Hunyuan3DFeedForwardWeights(f"{prefix}.experts.{expert_idx}", mm_type) for expert_idx in range(num_experts)) + experts = WeightModuleList(Hunyuan3DFeedForwardWeights(f"{prefix}.experts.{expert_idx}", moe_mm_type) for expert_idx in range(num_experts)) self.add_module("experts", experts) def load(self, weight_dict): super().load(weight_dict) + self._validate_no_fp8_moe_weights() if self.moe_backend == "flashinfer" and self.experts[0].fc1._get_actual_weight() is not None: self._build_flashinfer_weights() + @staticmethod + def _is_fp8_tensor(tensor): + return tensor is not None and tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + + def _iter_moe_mm_weights(self): + yield "gate", self.gate + yield "shared_experts.fc1", self.shared_experts.fc1 + yield "shared_experts.fc2", self.shared_experts.fc2 + for expert_idx, expert in enumerate(self.experts): + yield f"experts.{expert_idx}.fc1", expert.fc1 + yield f"experts.{expert_idx}.fc2", expert.fc2 + + def _validate_no_fp8_moe_weights(self): + for name, module in self._iter_moe_mm_weights(): + for attr in ("weight", "pin_weight", "weight_cuda_buffer"): + if self._is_fp8_tensor(getattr(module, attr, None)): + raise ValueError( + f"Hunyuan3D MoE FP8 is not supported yet, but {name}.{attr} is FP8. " + "Regenerate the checkpoint so .moe. weights stay BF16." + ) + def to_cuda(self, non_blocking=False): super().to_cuda(non_blocking=non_blocking) if self.moe_backend == "flashinfer": diff --git a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py index d9f9e790a..4268d1b0e 100644 --- a/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py +++ b/lightx2v/models/runners/hunyuan3d/hunyuan3d_shape_runner.py @@ -81,12 +81,12 @@ def load_transformer(self): return model def _build_dit_weight_dict(self): - state_dict = self._ckpt["model"] - dtype = GET_DTYPE() - weight_dict = {} - for key, tensor in state_dict.items(): - weight_dict[key] = tensor.to(dtype=dtype) - return weight_dict + ckpt = self._ckpt + ckpt_path = self.config.get("dit_original_ckpt") + if ckpt_path: + use_safetensors = bool(self.config.get("use_safetensors", False)) or str(ckpt_path).endswith(".safetensors") + ckpt = load_checkpoint_dict(ckpt_path, use_safetensors=use_safetensors) + return Hunyuan3DDiTModel._build_weight_dict(ckpt["model"], GET_DTYPE() == GET_SENSITIVE_DTYPE(), {}, config=self.config) def load_text_encoder(self): return [] From 955b96422b8ba770673fe859b704b55396706fbf Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Mon, 13 Jul 2026 19:49:38 +0800 Subject: [PATCH 4/5] Add Hunyuan3D fused QKV projection --- configs/hunyuan3d/hunyuan3d_shape_fp8.json | 1 + lightx2v/common/ops/norm/rms_norm_weight.py | 10 +-- .../hunyuan3d/infer/transformer_infer.py | 37 +++++++-- .../hunyuan3d/weights/transformer_weights.py | 78 ++++++++++++++++++- 4 files changed, 105 insertions(+), 21 deletions(-) diff --git a/configs/hunyuan3d/hunyuan3d_shape_fp8.json b/configs/hunyuan3d/hunyuan3d_shape_fp8.json index e0809ba78..16b2fd110 100644 --- a/configs/hunyuan3d/hunyuan3d_shape_fp8.json +++ b/configs/hunyuan3d/hunyuan3d_shape_fp8.json @@ -9,6 +9,7 @@ "octree_resolution": 384, "attn_type": "torch_sdpa", "use_fused_qk_rms_norm": true, + "use_fused_qkv_attn": true, "moe_backend": "flashinfer", "moe_flashinfer_setting": { "autotune": true, diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 48ff130e6..02cb5171a 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -443,15 +443,7 @@ def apply_qk_rms_norm( if norm_q is None and norm_k is None: return query, key - if ( - use_triton - and norm_q is not None - and norm_k is not None - and norm_q.eps == norm_k.eps - and query.is_cuda - and key.is_cuda - and query.shape[-1] == key.shape[-1] - ): + if use_triton and norm_q is not None and norm_k is not None and norm_q.eps == norm_k.eps and query.is_cuda and key.is_cuda and query.shape[-1] == key.shape[-1]: q_shape = query.shape k_shape = key.shape head_dim = q_shape[-1] diff --git a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py index 13b223a0d..5892841e6 100644 --- a/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan3d/infer/transformer_infer.py @@ -25,6 +25,7 @@ def __init__(self, config): self.head_dim = config["hidden_size"] // self.num_heads self.scheduler = None self.use_fused_qk_rms_norm = bool(config.get("use_fused_qk_rms_norm", False)) + self.use_fused_qkv_attn = bool(config.get("use_fused_qkv_attn", False)) self.fi_moe_autotune = MoeFiAutotune.from_hunyuan3d_config(config) if self.fi_moe_autotune.enabled: if flashinfer_autotune is None: @@ -49,6 +50,10 @@ def _flatten_norm(norm_weight, hidden_states): @staticmethod def _reshape_self_qkv(q, k, v, num_heads, head_dim, batch_size, seq_len): qkv = torch.cat((q, k, v), dim=-1) + return Hunyuan3DTransformerInfer._reshape_self_qkv_packed(qkv, num_heads, head_dim, batch_size, seq_len) + + @staticmethod + def _reshape_self_qkv_packed(qkv, num_heads, head_dim, batch_size, seq_len): split_size = qkv.shape[-1] // num_heads // 3 qkv = qkv.reshape(batch_size * seq_len, num_heads, split_size * 3) q, k, v = torch.split(qkv, split_size, dim=-1) @@ -82,6 +87,14 @@ def _project_qkv(self, hidden_states, to_q, to_k, to_v, norm_q, norm_k): query, key = self._apply_qk_norm(query, key, norm_q, norm_k) return query, key, value + def _project_fused_qkv(self, hidden_states, to_qkv, norm_q, norm_k): + batch_size, seq_len, hidden_dim = hidden_states.shape + flat = hidden_states.reshape(-1, hidden_dim) + qkv = to_qkv.apply(flat).reshape(batch_size, seq_len, hidden_dim * 3) + query, key, value = self._reshape_self_qkv_packed(qkv, self.num_heads, self.head_dim, batch_size, seq_len) + query, key = self._apply_qk_norm(query, key, norm_q, norm_k) + return query, key, value + def _run_attention(self, query, key, value, calculate, merge_batch=False): batch_size = query.shape[0] if batch_size == 1: @@ -145,14 +158,22 @@ def _run_attention(self, query, key, value, calculate, merge_batch=False): def _infer_self_attention(self, block_weights, hidden_states): norm_hidden = self._flatten_norm(block_weights.norm1, hidden_states) - query, key, value = self._project_qkv( - norm_hidden, - block_weights.attn1.to_q, - block_weights.attn1.to_k, - block_weights.attn1.to_v, - block_weights.attn1.norm_q, - block_weights.attn1.norm_k, - ) + if self.use_fused_qkv_attn and block_weights.attn1.has_fused_qkv: + query, key, value = self._project_fused_qkv( + norm_hidden, + block_weights.attn1.to_qkv, + block_weights.attn1.norm_q, + block_weights.attn1.norm_k, + ) + else: + query, key, value = self._project_qkv( + norm_hidden, + block_weights.attn1.to_q, + block_weights.attn1.to_k, + block_weights.attn1.to_v, + block_weights.attn1.norm_q, + block_weights.attn1.norm_k, + ) attn_output = self._run_attention(query, key, value, block_weights.attn1.calculate, merge_batch=False) batch_size, seq_len, _ = hidden_states.shape flat = attn_output.reshape(-1, attn_output.shape[-1]) diff --git a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py index 4694b483b..5349a867a 100644 --- a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py +++ b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py @@ -95,10 +95,7 @@ def _validate_no_fp8_moe_weights(self): for name, module in self._iter_moe_mm_weights(): for attr in ("weight", "pin_weight", "weight_cuda_buffer"): if self._is_fp8_tensor(getattr(module, attr, None)): - raise ValueError( - f"Hunyuan3D MoE FP8 is not supported yet, but {name}.{attr} is FP8. " - "Regenerate the checkpoint so .moe. weights stay BF16." - ) + raise ValueError(f"Hunyuan3D MoE FP8 is not supported yet, but {name}.{attr} is FP8. Regenerate the checkpoint so .moe. weights stay BF16.") def to_cuda(self, non_blocking=False): super().to_cuda(non_blocking=non_blocking) @@ -140,10 +137,18 @@ def __init__(self, config, block_idx, mm_type, ln_type, rms_norm_type, attn_type self.num_heads = config["num_heads"] self.head_dim = config["hidden_size"] // self.num_heads self.qk_norm = config.get("qk_norm", True) + self.use_fused_qkv_attn = bool(config.get("use_fused_qkv_attn", False)) self.add_module("to_q", MM_WEIGHT_REGISTER[mm_type](f"{prefix}.to_q.weight", f"{prefix}.to_q.bias" if qkv_bias else None)) self.add_module("to_k", MM_WEIGHT_REGISTER[mm_type](f"{prefix}.to_k.weight", f"{prefix}.to_k.bias" if qkv_bias else None)) self.add_module("to_v", MM_WEIGHT_REGISTER[mm_type](f"{prefix}.to_v.weight", f"{prefix}.to_v.bias" if qkv_bias else None)) + if self.use_fused_qkv_attn: + self.to_qkv = MM_WEIGHT_REGISTER[mm_type]( + f"{prefix}.to_qkv.weight", + f"{prefix}.to_qkv.bias" if qkv_bias else None, + ) + else: + self.to_qkv = None if self.qk_norm: self.add_module("norm_q", RMS_WEIGHT_REGISTER[rms_norm_type](f"{prefix}.q_norm.weight")) self.add_module("norm_k", RMS_WEIGHT_REGISTER[rms_norm_type](f"{prefix}.k_norm.weight")) @@ -153,6 +158,71 @@ def __init__(self, config, block_idx, mm_type, ln_type, rms_norm_type, attn_type self.add_module("out_proj", MM_WEIGHT_REGISTER[mm_type](f"{prefix}.out_proj.weight", f"{prefix}.out_proj.bias")) self.add_module("calculate", ATTN_WEIGHT_REGISTER[attn_type]()) + def load(self, weight_dict): + super().load(weight_dict) + self._build_fused_qkv() + + def to_cuda(self, non_blocking=False): + super().to_cuda(non_blocking=non_blocking) + self._build_fused_qkv() + + def to_cpu(self, non_blocking=False): + super().to_cpu(non_blocking=non_blocking) + self._build_fused_qkv() + + @staticmethod + def _output_dim(module): + if hasattr(module, "weight_need_transpose"): + return 1 if module.weight_need_transpose else 0 + return 1 + + @staticmethod + def _cat_optional_biases(biases): + if all(bias is None for bias in biases): + return None + if any(bias is None for bias in biases): + raise ValueError("Fused Hunyuan3D QKV requires either all QKV biases or no QKV biases") + return torch.cat([bias.contiguous() for bias in biases], dim=0) + + def _cat_output_weights(self, modules, tensors): + dim = self._output_dim(modules[0]) + if dim == 1: + return torch.cat([tensor.t().contiguous() for tensor in tensors], dim=0).t() + return torch.cat([tensor.contiguous() for tensor in tensors], dim=dim) + + def _cat_output_scales(self, scales): + if scales[0].dim() == 1: + dim = 0 + elif scales[0].shape[-1] == 1: + dim = 0 + else: + dim = scales[0].dim() - 1 + return torch.cat([scale.contiguous() for scale in scales], dim=dim) + + def _build_fused_qkv(self): + if self.to_qkv is None: + return + modules = (self.to_q, self.to_k, self.to_v) + weights = [module._get_actual_weight() for module in modules] + if any(weight is None for weight in weights): + return + + self.to_qkv.weight = self._cat_output_weights(modules, weights) + biases = [module._get_actual_bias() for module in modules] + self.to_qkv.bias = self._cat_optional_biases(biases) + self.to_qkv.has_lora_branch = False + + if all(hasattr(module, "weight_scale") for module in modules): + scales = [module.weight_scale for module in modules] + if all(scale is not None for scale in scales): + self.to_qkv.weight_scale = self._cat_output_scales(scales) + + @property + def has_fused_qkv(self): + if self.to_qkv is None or not hasattr(self.to_qkv, "weight"): + return False + return not any(getattr(module, "has_lora_branch", False) or getattr(module, "has_diff", False) for module in (self.to_q, self.to_k, self.to_v)) + class Hunyuan3DCrossAttentionWeights(WeightModule): def __init__(self, config, block_idx, mm_type, rms_norm_type, attn_type, qkv_bias): From 8815c8563f4b833526a480f26364d16d4a164a51 Mon Sep 17 00:00:00 2001 From: STwangyingrui Date: Mon, 13 Jul 2026 13:14:15 +0000 Subject: [PATCH 5/5] fix as gemini assist suggests --- configs/hunyuan3d/hunyuan3d_shape_fp8.json | 2 +- .../models/networks/hunyuan3d/weights/transformer_weights.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/configs/hunyuan3d/hunyuan3d_shape_fp8.json b/configs/hunyuan3d/hunyuan3d_shape_fp8.json index 16b2fd110..26bfca864 100644 --- a/configs/hunyuan3d/hunyuan3d_shape_fp8.json +++ b/configs/hunyuan3d/hunyuan3d_shape_fp8.json @@ -15,7 +15,7 @@ "autotune": true, "tune_max_num_tokens": 8192 }, - "dit_original_ckpt": "/data/nvme0/wangyingrui/Hunyuan3D-2.1-fp8/hunyuan3d-dit-v2-1/model.fp8.safetensors", + "dit_original_ckpt": "/path/to/hunyuan3d/model.fp8.safetensors", "dit_quant_scheme": "fp8-sgl", "rms_norm_type": "torch" } diff --git a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py index 5349a867a..e5c6499fd 100644 --- a/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py +++ b/lightx2v/models/networks/hunyuan3d/weights/transformer_weights.py @@ -117,6 +117,8 @@ def _stack_optional_biases(biases): return torch.stack([bias.contiguous() for bias in biases], dim=0) def _build_flashinfer_weights(self): + if self.experts[0].fc1._get_actual_weight() is None: + return fc1_list, fc2_list = [], [] fc1_biases, fc2_biases = [], [] for expert_w in self.experts: