Optimize hunyuan3d

configs/hunyuan3d/hunyuan3d_shape.json

@@ -8,5 +8,11 @@
     "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
+    },
     "rms_norm_type": "torch"
 }

configs/hunyuan3d/hunyuan3d_shape_fp8.json

@@ -0,0 +1,21 @@
+{
+    "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,
+    "use_fused_qkv_attn": 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"
+}

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,42 @@ 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

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)
 # ---------------------------------------------------------------------------

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

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)

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)