diff --git a/DEVELOPER.md b/DEVELOPER.md index 92474406..22c7f599 100644 --- a/DEVELOPER.md +++ b/DEVELOPER.md @@ -100,7 +100,8 @@ difflet/ └── models/ flux/ wan/ hunyuan_video/ qwen_image/ ltx_2/ — each: modeling, pipeline, application, entry, checkpoint (where needed), and spec.py for models on the GPU flow. HunyuanVideo 1.5 is served from - hunyuan_video/ (model_version="1.5"). + hunyuan_video/ (model_version="1.5"); on CUDA it has its own path, + hunyuan_video_15/spec.py and backends/cuda/hunyuan_video_15/. ``` `difflet/backends/trainium/{core,modules}` follows upstream Neuron coding style so periodic @@ -251,7 +252,8 @@ mesh `{dp, cfg, cp, tp}` (tp innermost, dp outermost): true-CFG models (Wan, LTX-2, HunyuanVideo 1.5) — guidance-distilled models (Flux, HunyuanVideo) and Qwen-Image run a single forward pass and reject it. HunyuanVideo 1.5 is not guidance-distilled: its guider runs classifier-free guidance at 6.0 on every step, and - `--cfg-parallel` (TP2 × CFG2) is its default layout. + `--cfg-parallel` (TP2 × CFG2) is its default layout on Trainium. On CUDA it runs both CFG + branches in sequence on one GPU, and `--cfg-parallel` is refused. - `sp_enabled` — Megatron-style sequence parallelism: shards the otherwise-replicated norm/modulation/residual regions along the sequence axis across the *existing* tensor-parallel group (reduce-scatter replaces the row-parallel all-reduce). Adds no mesh diff --git a/NOTICE b/NOTICE index 93aa0d6c..1467378f 100644 --- a/NOTICE +++ b/NOTICE @@ -52,7 +52,7 @@ This product includes software derived from the following third-party projects: `scripts/rebase_from_nxdi.py` for the rebase workflow. ================================================================================ -2. HuggingFace diffusers — pipeline extensions and LTX-2 attention +2. HuggingFace diffusers — pipeline extensions, LTX-2 attention and HunyuanVideo 1.5 ================================================================================ Source: https://github.com/huggingface/diffusers @@ -66,6 +66,12 @@ This product includes software derived from the following third-party projects: LTX2VideoTransformer3DModel attention processor under Apache License 2.0. Copyright 2025 The Lightricks team and The HuggingFace Team. The file retains the upstream license header and identifies the adaptation. + The HunyuanVideo 1.5 CUDA transformer, + difflet/backends/cuda/hunyuan_video_15/modeling.py, adapts portions of + diffusers' models/transformers/transformer_hunyuan_video15.py. That file keeps + the upstream copyright (2025 The Hunyuan Team and The HuggingFace Team) and its + Apache License 2.0 header; Difflet replaces its projections and attention and + keeps the checkpoint layout. ================================================================================ 3. xDiT (xfuser) — design inspiration only, no code diff --git a/README.md b/README.md index f52a2be6..dd0ce229 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ see the per-model limits below. | Wan 2.2 | ✅ video | ✅ video | ✅ | ❌ | ⚠️ ⁵ | ✅ ⁶ | ❌ | ❌ | | HunyuanVideo (original) | ⚠️ experimental ⁷ | ⚠️ experimental ⁷ | ⚠️ execution observed ⁷ | ❌ | Not measured | — ¹ | ❌ | ❌ | | LTX-2 | ⚠️ video, CPU checks ⁸ | ⚠️ video, CPU checks ⁸ | ⚠️ Python qualified ⁸ | ❌ | Not measured | ❌ | ❌ | ❌ | +| HunyuanVideo 1.5 | ⚠️ experimental ⁹ | ⚠️ experimental ⁹ | ✅ ⁹ | ❌ | Not measured | ❌ | ❌ | ❌ | An option the GPU backend does not support is refused with a message naming it, not ignored. TeaCache (adaptive), TAEF1, and `--host-vae` are refused on GPUs for these models. @@ -148,6 +149,12 @@ TeaCache (adaptive), TAEF1, and `--host-vae` are refused on GPUs for these model and performance checks on B200. CLI/HTTP video-only paths have CPU coverage; GPU frontend validation is not claimed. The opt-in profile requires explicit TP=1 and a pre-CUDA startup setting; see the [LTX-2 CUDA guide](docs/verification/ltx2-cuda.md). +9. Experimental single-GPU BF16 text-to-video for + `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, separate from the Trainium + HunyuanVideo 1.5 path. Classifier-free guidance runs both branches in sequence on one GPU. + On one RTX PRO 6000, eager and compiled output meet a relative accuracy bar against + vLLM-Omni and warm latency is within 10% of it; see the + [HunyuanVideo 1.5 CUDA guide](docs/verification/hunyuanvideo15-cuda.md). Measured timings against stock diffusers, and the full pass/limit matrix, are in [docs/verification/gpu-backend-results.md](docs/verification/gpu-backend-results.md). @@ -313,9 +320,9 @@ LTX-2 also has an experimental single-GPU CUDA path with joint audiovisual Pytho and an opt-in native-compatible BF16 profile. See the [LTX-2 CUDA guide](docs/verification/ltx2-cuda.md) for its separate runtime, explicit TP=1/startup requirements and qualified eager/compiled results. -For HunyuanVideo CUDA, pass `parallel=DiffletParallelConfig(tp_degree=1)` explicitly. -Omitting `parallel` preserves the model-specific registry default, including HunyuanVideo's -Trainium topology. GPU Python API paths have unit coverage; see the +For HunyuanVideo and HunyuanVideo 1.5 on CUDA, pass `parallel=DiffletParallelConfig(tp_degree=1)` +explicitly. Omitting `parallel` preserves the model-specific registry default, including the +HunyuanVideo models' Trainium topologies. GPU Python API paths have unit coverage; see the [per-model results](docs/verification/gpu-backend-results.md) for hardware coverage. ## Choose a topology diff --git a/difflet/backends/cuda/hunyuan_video_15/__init__.py b/difflet/backends/cuda/hunyuan_video_15/__init__.py new file mode 100644 index 00000000..7cc2a394 --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/__init__.py @@ -0,0 +1 @@ +"""Standard HunyuanVideo 1.5 on CUDA.""" diff --git a/difflet/backends/cuda/hunyuan_video_15/application.py b/difflet/backends/cuda/hunyuan_video_15/application.py new file mode 100644 index 00000000..8ae2499a --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/application.py @@ -0,0 +1,175 @@ +"""Standard 1.5 conditioning, CFG and scheduler semantics from the Diffusers driver.""" + +from __future__ import annotations + +from contextlib import contextmanager + +import torch +from diffusers.models.modeling_outputs import Transformer2DModelOutput + +from difflet.backends.cuda.hunyuan_video_15.conditioning import configure_conditioning +from difflet.backends.cuda.hunyuan_video_15.dit_config import HunyuanVideo15DiTConfig +from difflet.backends.gpu.application import GpuApplication +from difflet.models.hunyuan_video_15.spec import SPEC + + +class HunyuanVideo15DiTAdapter(torch.nn.Module): + def __init__(self, component): + super().__init__() + self.component = component + self.config = component.dit_config + + @property + def dtype(self): + return self.component.dtype + + @property + def device(self): + return self.component.device + + @contextmanager + def cache_context(self, name): + yield + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask, + encoder_hidden_states_2, + encoder_attention_mask_2, + image_embeds, + attention_kwargs=None, + return_dict=True, + ): + # Every input is cast to the transformer's dtype, as Diffusers' HunyuanVideo 1.5 driver + # does. Unlike the original HunyuanVideo adapter, the timestep is not kept in FP32. + values = ( + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask, + encoder_hidden_states_2, + encoder_attention_mask_2, + image_embeds, + ) + inputs = tuple( + t.to(device=self.device, dtype=torch.bool if i in (3, 5) else self.dtype) + for i, t in enumerate(values) + ) + result = self.component(*inputs, return_dict=False)[0] + return Transformer2DModelOutput(sample=result) if return_dict else (result,) + + +class CudaHunyuanVideo15Application(GpuApplication): + spec = SPEC + + def _configure(self, kwargs): + self.local_files_only = bool(kwargs.get("local_files_only", True)) + self.pipe = None + self.adapter = None + self._guider_config = None + + def _build_transformers(self): + from difflet.backends.cuda.hunyuan_video_15.transformer import ( + CudaHunyuanVideo15TransformerApplication, + ) + + self.config = HunyuanVideo15DiTConfig.from_pretrained( + self.transformer_path, **self.shape, torch_dtype=self.dtype + ) + self.transformer = CudaHunyuanVideo15TransformerApplication( + model_path=self.transformer_path, + config=self.config, + device=self.device, + torch_compile=self.torch_compile, + ) + self.adapter = HunyuanVideo15DiTAdapter(self.transformer) + + def load_companions(self): + from diffusers import HunyuanVideo15Pipeline + + base = HunyuanVideo15Pipeline.from_pretrained( + self.model_path, + transformer=None, + torch_dtype=self.dtype, + local_files_only=self.local_files_only, + ) + # Diffusers' native overlapping spatial tiles (256 px, 25% overlap) keep the decode + # within one GPU at 121 frames. + base.vae.enable_tiling() + for name in ("text_encoder", "text_encoder_2", "vae"): + getattr(base, name).to(self.device).requires_grad_(False).eval() + self.pipe = HunyuanVideo15Pipeline(**{**base.components, "transformer": self.adapter}) + configure_conditioning(self.pipe) + self._guider_config = dict(base.guider.config) + + def dit_input_contract(self): + c = self._require_transformer("input contract").dit_config + shapes = ( + ("hidden_states", (1, c.in_channels, c.latent_frames, c.latent_height, c.latent_width)), + ("timestep", (1,)), + ("encoder_hidden_states", (1, c.text_seq_len, c.text_embed_dim)), + ("encoder_attention_mask", (1, c.text_seq_len)), + ("encoder_hidden_states_2", (1, c.text_2_seq_len, c.text_embed_2_dim)), + ("encoder_attention_mask_2", (1, c.text_2_seq_len)), + ("image_embeds", (1, c.image_tokens, c.image_embed_dim)), + ) + return { + name: {"shape": shape, "dtype": torch.bool if "mask" in name else self.dtype} + for name, shape in shapes + } + + def generate( + self, + prompt, + *, + negative_prompt=None, + num_inference_steps=50, + guidance_scale=6.0, + seed=None, + latents=None, + output_type="pt", + height=None, + width=None, + num_frames=None, + ): + from diffusers import ClassifierFreeGuidance + + if self.pipe is None: + raise RuntimeError("application is not loaded; call load() first") + if not isinstance(prompt, str): + raise ValueError("HunyuanVideo 1.5 generates one prompt per request") + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be at least 1") + for key, value in (("height", height), ("width", width), ("num_frames", num_frames)): + if value is not None and value != self.shape[key]: + raise ValueError(f"{key} differs from the loaded shape {self.shape[key]}") + if output_type not in ("pt", "latent"): + raise ValueError("HunyuanVideo 1.5 output_type must be pt or latent") + generator = None + if latents is not None: + c = self.config + expected = (1, c.out_channels, c.latent_frames, c.latent_height, c.latent_width) + if tuple(latents.shape) != expected: + raise ValueError(f"latents must have shape {expected}") + elif seed is not None: + generator = torch.Generator(device=self.device).manual_seed(int(seed)) + self.pipe.guider = ClassifierFreeGuidance.from_config( + {**self._guider_config, "guidance_scale": float(guidance_scale)} + ) + output = self.pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=self.shape["height"], + width=self.shape["width"], + num_frames=self.shape["num_frames"], + num_inference_steps=num_inference_steps, + generator=generator, + latents=latents, + output_type=output_type, + ).frames + if output_type == "latent": + return output + return output.permute(0, 2, 1, 3, 4).to(device="cpu", dtype=torch.float32) * 2 - 1 diff --git a/difflet/backends/cuda/hunyuan_video_15/conditioning.py b/difflet/backends/cuda/hunyuan_video_15/conditioning.py new file mode 100644 index 00000000..1728480d --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/conditioning.py @@ -0,0 +1,42 @@ +"""Prompt conditioning as the reference computes it: device-local Qwen rotary +frequencies, and cuDNN SDPA off while prompts are encoded.""" + +from functools import wraps + +import torch + + +def initialize_qwen_rotary(qwen): + """Recompute the nonpersistent FP32 rotary frequencies on the encoder's device. + + Frequencies computed on the CPU and moved to CUDA differ in the last bits from the + device's own FP32 power arithmetic, which the reference conditioning uses. + """ + rotary = qwen.rotary_emb + if rotary.config.rope_parameters.get("rope_type", "default") != "default": + raise NotImplementedError("HunyuanVideo 1.5 requires default Qwen rotary embeddings") + with torch.no_grad(), torch.device(next(qwen.parameters()).device): + inv_freq, _ = rotary.compute_default_rope_parameters(rotary.config) + for name in ("inv_freq", "original_inv_freq"): + setattr(rotary, name, inv_freq.clone()) + + +def configure_conditioning(pipe): + """Encode prompts with cuDNN SDPA disabled, as the reference conditioning does. + + Only ``encode_prompt`` changes the policy, and it restores the caller's afterwards; the + denoiser keeps its own. No model class or learned weight changes. + """ + initialize_qwen_rotary(pipe.text_encoder) + original = pipe.encode_prompt + + @wraps(original) + def encode_prompt(*args, **kwargs): + previous = torch.backends.cuda.cudnn_sdp_enabled() + try: + torch.backends.cuda.enable_cudnn_sdp(False) + return original(*args, **kwargs) + finally: + torch.backends.cuda.enable_cudnn_sdp(previous) + + pipe.encode_prompt = encode_prompt diff --git a/difflet/backends/cuda/hunyuan_video_15/dit_config.py b/difflet/backends/cuda/hunyuan_video_15/dit_config.py new file mode 100644 index 00000000..1b390b41 --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/dit_config.py @@ -0,0 +1,78 @@ +"""The HunyuanVideo 1.5 text-to-video transformer config and its resident shape.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, fields +from pathlib import Path + +import torch + + +@dataclass +class HunyuanVideo15DiTConfig: + in_channels: int = 65 + out_channels: int = 32 + num_attention_heads: int = 16 + attention_head_dim: int = 128 + num_layers: int = 54 + num_refiner_layers: int = 2 + mlp_ratio: float = 4.0 + patch_size: int = 1 + patch_size_t: int = 1 + qk_norm: str = "rms_norm" + text_embed_dim: int = 3584 + text_embed_2_dim: int = 1472 + image_embed_dim: int = 1152 + rope_theta: float = 256.0 + rope_axes_dim: tuple[int, ...] = (16, 56, 56) + target_size: int = 640 + task_type: str = "t2v" + use_meanflow: bool = False + height: int = 480 + width: int = 848 + num_frames: int = 121 + text_seq_len: int = 1000 + text_2_seq_len: int = 256 + image_tokens: int = 729 + torch_dtype: torch.dtype = torch.bfloat16 + + def __post_init__(self): + # The one owner of the resident shape: the application builds this config before any + # weights load. + if min(self.height, self.width, self.num_frames) <= 0: + raise ValueError("HunyuanVideo 1.5 height, width and num_frames must be positive") + if self.height % 16 or self.width % 16: + raise ValueError("HunyuanVideo 1.5 height and width must be multiples of 16") + if (self.num_frames - 1) % 4: + raise ValueError("HunyuanVideo 1.5 num_frames must be 4n+1") + if self.task_type != "t2v" or self.use_meanflow: + raise NotImplementedError("HunyuanVideo 1.5 supports standard text-to-video only") + if self.patch_size != 1 or self.patch_size_t != 1 or self.qk_norm != "rms_norm": + raise NotImplementedError( + "HunyuanVideo 1.5 supports patch size 1 with RMSNorm Q/K normalization only" + ) + self.rope_axes_dim = tuple(self.rope_axes_dim) + + @property + def latent_height(self): + return self.height // 16 + + @property + def latent_width(self): + return self.width // 16 + + @property + def latent_frames(self): + return (self.num_frames - 1) // 4 + 1 + + @classmethod + def from_pretrained(cls, transformer_dir, **overrides): + body = json.loads((Path(transformer_dir) / "config.json").read_text()) + if ( + body.get("_class_name", "HunyuanVideo15Transformer3DModel") + != "HunyuanVideo15Transformer3DModel" + ): + raise NotImplementedError("Unsupported HunyuanVideo 1.5 transformer class") + known = {f.name for f in fields(cls)} + return cls(**{**{k: v for k, v in body.items() if k in known}, **overrides}) diff --git a/difflet/backends/cuda/hunyuan_video_15/modeling.py b/difflet/backends/cuda/hunyuan_video_15/modeling.py new file mode 100644 index 00000000..7227dc73 --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/modeling.py @@ -0,0 +1,407 @@ +# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# Copyright 2026 Difflet contributors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The HunyuanVideo 1.5 text-to-video DiT for the CUDA backend. + +Diffusers' building blocks keep the checkpoint layout, normalization and cast boundaries; +Difflet's projections and attention replace theirs. Eager single-sample blocks run the +joint attention with packed QKV, Difflet's fused Q/K RMSNorm + RoPE op (``qk_norm_rope``) +and cuDNN SDPA; compiled blocks attend each sample's valid token prefix. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.attention import SDPBackend, sdpa_kernel +from diffusers.models.embeddings import apply_rotary_emb +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import AdaLayerNormContinuous +from diffusers.models.transformers.transformer_hunyuan_video15 import ( + HunyuanVideo15ByT5TextProjection, + HunyuanVideo15ImageProjection, + HunyuanVideo15PatchEmbed, + HunyuanVideo15RotaryPosEmbed, + HunyuanVideo15TimeEmbedding, + HunyuanVideo15TokenRefiner, + HunyuanVideo15TransformerBlock, +) + +from difflet.backends.cuda.hunyuan_video_15.qk_norm_rope import ( + HAS_TRITON, + fused_joint_qk_norm_rope, + pack_rope_table, +) +from difflet.ops import ColumnParallelLinear, attention + + +def _bind_projections(module): + """Keep original state keys and fused linear bias arithmetic at TP=1.""" + for name, child in list(module.named_children()): + if type(child) is nn.Linear: + replacement = ColumnParallelLinear( + child.in_features, + child.out_features, + bias=child.bias is not None, + gather_output=True, + ) + replacement.weight = child.weight + replacement.bias = child.bias + setattr(module, name, replacement) + else: + _bind_projections(child) + + +def _prepare_compact_attention_bias(valid: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """The token refiner's additive key bias, ``[B, 1, 1, S]``: zero for valid keys and + -inf for padding, one value per key rather than an S-by-S mask.""" + bias = torch.zeros_like(valid, dtype=dtype).masked_fill(~valid, float("-inf")) + return bias[:, None, None, :] + + +@torch.library.custom_op("difflet_hv15::prefix_attention", mutates_args=()) +def prefix_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + lengths: torch.Tensor, + scale: float, +) -> torch.Tensor: + """Attend each sample over its valid joint-token prefix; padded query rows are zero. + + Equals valid-query AND valid-key masking when valid tokens lead the sequence, as the + DiT's packed [video | glyph | text | padding] layout guarantees. The prefix is + attended unmasked, so SDPA can pick its fastest kernel (cuDNN on recent GPUs, against a + slower kernel for the padded, biased call). Opaque to the compiler, so blocks keep one + static-shape graph for every prompt length. ``lengths`` is a CPU int64 vector. + """ + output = torch.zeros_like(query) + for index, length in enumerate(lengths.tolist()): + window = (slice(index, index + 1), slice(None), slice(0, length)) + output[window] = attention(query[window], key[window], value[window], scale=scale) + return output + + +@prefix_attention.register_fake +def _(query, key, value, lengths, scale): + return torch.empty_like(query) + + +class _AttentionProcessor: + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + image_rotary_emb=None, + ): + query = attn.to_q(hidden_states).unflatten(2, (attn.heads, -1)) + key = attn.to_k(hidden_states).unflatten(2, (attn.heads, -1)) + value = attn.to_v(hidden_states).unflatten(2, (attn.heads, -1)) + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + if encoder_hidden_states is not None: + cq = attn.add_q_proj(encoder_hidden_states).unflatten(2, (attn.heads, -1)) + ck = attn.add_k_proj(encoder_hidden_states).unflatten(2, (attn.heads, -1)) + cv = attn.add_v_proj(encoder_hidden_states).unflatten(2, (attn.heads, -1)) + cq = attn.norm_added_q(cq) + ck = attn.norm_added_k(ck) + query, key, value = ( + torch.cat(pair, dim=1) for pair in ((query, cq), (key, ck), (value, cv)) + ) + result = self._attend(query, key, value, attn.scale, attention_mask) + result = result.flatten(2, 3).to(query.dtype) + if encoder_hidden_states is None: + return attn.to_out[1](attn.to_out[0](result)) + result, context = result.split( + (hidden_states.shape[1], encoder_hidden_states.shape[1]), dim=1 + ) + return attn.to_out[1](attn.to_out[0](result)), attn.to_add_out(context) + + @staticmethod + def _attend(query, key, value, scale, attention_mask): + """[B, S, H, D] in and out; masked key bias, then zeroed invalid query rows.""" + result = attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + scale=scale, + attention_mask=attention_mask, + ).transpose(1, 2) + # Stock uses valid-query AND valid-key masking. Zero invalid query results + # before projections, reproducing it without ever building an S-by-S mask. + if attention_mask is not None: + keep = attention_mask == 0 + result = result.masked_fill(~keep[:, 0, 0, :, None, None], 0) + return result + + +class _PrefixAttentionProcessor(_AttentionProcessor): + """Joint DiT attention over each sample's valid prefix; ``attention_mask`` is lengths.""" + + @staticmethod + def _attend(query, key, value, scale, lengths): + return prefix_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), lengths, scale + ).transpose(1, 2) + + +def _cudnn_attention_available(device: torch.device) -> bool: + """cuDNN SDPA needs NVIDIA cuDNN 9 and sm_80+.""" + return bool( + torch.version.hip is None + and torch.backends.cudnn.is_available() + and (torch.backends.cudnn.version() or 0) >= 90000 + and torch.cuda.get_device_capability(device) >= (8, 0) + ) + + +def joint_attention_supported(weight: torch.Tensor, head_dim: int) -> bool: + """Whether a block whose projection weight is ``weight`` can take the joint path: CUDA + with Triton for the fused op (BF16, head_dim 128) and cuDNN SDPA. Decided once, at load; + otherwise the blocks keep the prefix path.""" + return ( + weight.is_cuda + and HAS_TRITON + and weight.dtype == torch.bfloat16 + and head_dim == 128 + and _cudnn_attention_available(weight.device) + ) + + +class _JointAttentionProcessor(_PrefixAttentionProcessor): + """Joint DiT attention for one sample whose context is trimmed to its valid tokens. + + The same arithmetic as vLLM-Omni's HunyuanVideo 1.5 attention: packed QKV projections, + a fused Q/K RMSNorm + RoPE op and cuDNN SDPA. + + ``image_rotary_emb`` is ``(cos, sin, table)`` with the packed table from + ``pack_rope_table`` and ``attention_mask`` is ``None`` (every context token kept is + valid). Any other call takes the prefix path. Sequence: one QKV GEMM per stream, + the fused joint RMSNorm/RoPE/cat op, cuDNN-only SDPA on ``[B, H, S, D]`` views, + flatten, cast, split, output projections. + """ + + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + image_rotary_emb=None, + ): + if attention_mask is not None: # prefix lengths: the prefix path + return super().__call__( + attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb + ) + heads, head_dim = attn.heads, attn.norm_q.weight.shape[0] + inner = heads * head_dim + qkv = F.linear(hidden_states, attn._packed_qkv_weight, attn._packed_qkv_bias) + query, key, value = (t.unflatten(-1, (heads, -1)) for t in qkv.split([inner] * 3, dim=-1)) + added = F.linear( + encoder_hidden_states, attn._packed_added_qkv_weight, attn._packed_added_qkv_bias + ) + cq, ck, cv = (t.unflatten(-1, (heads, -1)) for t in added.split([inner] * 3, dim=-1)) + query, key, value = fused_joint_qk_norm_rope( + query, + key, + value, + cq, + ck, + cv, + attn.norm_q.weight, + attn.norm_k.weight, + attn.norm_added_q.weight, + attn.norm_added_k.weight, + image_rotary_emb[2], + attn.norm_q.eps, + ) + with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): + result = F.scaled_dot_product_attention( + query.permute(0, 2, 1, 3), + key.permute(0, 2, 1, 3), + value.permute(0, 2, 1, 3), + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=1.0 / (head_dim**0.5), + enable_gqa=False, + ).permute(0, 2, 1, 3) + result = result.flatten(2, 3).to(query.dtype) + context_len = encoder_hidden_states.shape[1] + result, context = result.split_with_sizes( + [result.shape[1] - context_len, context_len], dim=1 + ) + return attn.to_out[0](result), attn.to_add_out(context) + + +class _CompactRefiner(nn.Module): + def __init__(self, refiner): + super().__init__() + self.refiner_blocks = refiner.refiner_blocks + for block in self.refiner_blocks: + block.attn.set_processor(_AttentionProcessor()) + + def forward(self, hidden_states, temb, attention_mask=None): + mask = ( + None + if attention_mask is None + else _prepare_compact_attention_bias(attention_mask.bool(), hidden_states.dtype) + ) + for block in self.refiner_blocks: + hidden_states = block(hidden_states, temb, mask) + return hidden_states + + +class HunyuanVideo15Transformer3DModel(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + inner = config.num_attention_heads * config.attention_head_dim + self.x_embedder = HunyuanVideo15PatchEmbed((1, 1, 1), config.in_channels, inner) + self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, inner) + self.context_embedder = HunyuanVideo15TokenRefiner( + config.text_embed_dim, + config.num_attention_heads, + config.attention_head_dim, + num_layers=config.num_refiner_layers, + ) + self.context_embedder.token_refiner = _CompactRefiner(self.context_embedder.token_refiner) + self.context_embedder_2 = HunyuanVideo15ByT5TextProjection( + config.text_embed_2_dim, 2048, inner + ) + self.time_embed = HunyuanVideo15TimeEmbedding(inner) + self.cond_type_embed = nn.Embedding(3, inner) + self.rope = HunyuanVideo15RotaryPosEmbed(1, 1, config.rope_axes_dim, config.rope_theta) + self.transformer_blocks = nn.ModuleList( + [ + HunyuanVideo15TransformerBlock( + config.num_attention_heads, + config.attention_head_dim, + config.mlp_ratio, + qk_norm=config.qk_norm, + ) + for _ in range(config.num_layers) + ] + ) + for block in self.transformer_blocks: + block.attn.set_processor(_JointAttentionProcessor()) + self.norm_out = AdaLayerNormContinuous(inner, inner, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner, config.out_channels) + _bind_projections(self) + + @torch.no_grad() + def pack_attention_projections(self): + """One Q/K/V weight per stream, so each stream's projection is one GEMM. + + Run once the weights are on their device: each block's ``to_q/to_k/to_v`` (and + ``add_*_proj``) weights and biases are concatenated into one contiguous tensor + and rebound as views of it, so the checkpoint keys and memory are unchanged and + the joint processor runs one GEMM per stream. + """ + for block in self.transformer_blocks: + attn = block.attn + for names, prefix in ( + (("to_q", "to_k", "to_v"), "_packed_qkv"), + (("add_q_proj", "add_k_proj", "add_v_proj"), "_packed_added_qkv"), + ): + modules = [getattr(attn, name) for name in names] + weight = torch.cat([m.weight for m in modules]).contiguous() + bias = torch.cat([m.bias for m in modules]).contiguous() + rows = [m.weight.shape[0] for m in modules] + for module, w, b in zip(modules, weight.split(rows), bias.split(rows)): + module.weight = nn.Parameter(w, requires_grad=False) + module.bias = nn.Parameter(b, requires_grad=False) + attn.register_buffer(prefix + "_weight", weight, persistent=False) + attn.register_buffer(prefix + "_bias", bias, persistent=False) + + def _joint_attention_ready(self, hidden_states): + """Projections packed at load (eager, on a supported device) and one sample.""" + return ( + hasattr(self.transformer_blocks[0].attn, "_packed_qkv_weight") + and hidden_states.shape[0] == 1 + ) + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask, + encoder_hidden_states_2, + encoder_attention_mask_2, + image_embeds, + *, + return_dict=True, + ): + batch, _, frames, height, width = hidden_states.shape + rope = self.rope(hidden_states) + temb = self.time_embed(timestep) + hidden_states = self.x_embedder(hidden_states) + text = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) + text = text + self.cond_type_embed(torch.zeros_like(text[:, :, 0], dtype=torch.long)) + glyph = self.context_embedder_2(encoder_hidden_states_2) + glyph = glyph + self.cond_type_embed(torch.ones_like(glyph[:, :, 0], dtype=torch.long)) + image = self.image_embedder(image_embeds) * 0.0 + image = image + self.cond_type_embed(torch.full_like(image[:, :, 0], 2, dtype=torch.long)) + text_mask, glyph_mask = encoder_attention_mask.bool(), encoder_attention_mask_2.bool() + image_mask = torch.zeros(image.shape[:2], device=image.device, dtype=torch.bool) + contexts, masks = [], [] + for txt, tm, gl, gm, im, imask in zip( + text, text_mask, glyph, glyph_mask, image, image_mask + ): + contexts.append( + torch.cat( + ( + im[imask], + gl[gm], + txt[tm], + im[~imask], + torch.zeros_like(gl[~gm]), + torch.zeros_like(txt[~tm]), + ) + ) + ) + masks.append(torch.cat((imask[imask], gm[gm], tm[tm], imask[~imask], gm[~gm], tm[~tm]))) + context = torch.stack(contexts) + valid = torch.cat( + ( + torch.ones(hidden_states.shape[:2], device=hidden_states.device, dtype=torch.bool), + torch.stack(masks), + ), + dim=1, + ) + # Valid tokens lead every sample by construction, so per-sample prefix lengths (one + # host sync per DiT call) replace an S-sized key bias. + lengths = valid.sum(dim=1).cpu() + if self._joint_attention_ready(hidden_states): + # Keep only the sample's valid context tokens (then all valid, so no mask) and + # pack one RoPE table per forward. + context = context[:, : int(lengths[0]) - hidden_states.shape[1]] + table = pack_rope_table(rope[0], rope[1], context.shape[1], hidden_states.dtype) + joint = (rope[0], rope[1], table) + for block in self.transformer_blocks: + hidden_states, context = block(hidden_states, context, temb, None, joint) + else: + for block in self.transformer_blocks: + hidden_states, context = block(hidden_states, context, temb, lengths, rope) + output = self.proj_out(self.norm_out(hidden_states, temb)) + output = output.reshape(batch, frames, height, width, -1).permute(0, 4, 1, 2, 3) + return Transformer2DModelOutput(sample=output) if return_dict else (output,) diff --git a/difflet/backends/cuda/hunyuan_video_15/qk_norm_rope.py b/difflet/backends/cuda/hunyuan_video_15/qk_norm_rope.py new file mode 100644 index 00000000..394433d4 --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/qk_norm_rope.py @@ -0,0 +1,251 @@ +"""Fused joint Q/K RMSNorm + RoPE (+ V gather) for HunyuanVideo 1.5 on the CUDA backend. + +A HunyuanVideo 1.5 DiT block attends over the joint sequence ``[video | context]``. Each +stream has its own Q/K RMSNorm weights; video Q/K are rotated by RoPE, context Q/K are not. +``fused_joint_qk_norm_rope`` writes the joint Q, K and V that attention consumes in one Triton +launch, instead of per-stream norm passes, RoPE passes and three ``cat`` copies. Inputs are +read through their strides (on CUDA they are views of one packed QKV projection output); the +outputs are new contiguous ``[B, S0 + S1, H, D]`` tensors. + +Numerics, per output token and head, in FP32 (V is copied unchanged): + +1. Sum of squares over ``D/8`` lanes of eight elements. Each lane is an in-order chain of + explicit FMAs from zero; the lanes are then added pairwise, adjacent lanes first, as a tree. +2. ``inv = rsqrt(ss / D + eps)``, the device's fast rsqrt. +3. ``n = (x * inv) * w``: two FP32 multiplies, then one rounding to BF16 (nearest-even). +4. Adjacent pairs ``(n[2i], n[2i+1])`` are rotated by the table row's ``cos[i]``/``sin[i]`` + with separate FP32 products and one add or subtract, stored as BF16. Context rows of the + table are identity rows (cos 1, sin 0), so their normalised values pass through. + +The kernel is compiled with ``enable_fp_fusion=False``: the step-1 FMAs are the only fused +operations. It takes BF16 q/k/v, weights and table with ``D == 128``, on CUDA with Triton. +""" + +import torch + +try: + import triton + import triton.language as tl +except ImportError: # CPU installs: the model keeps its prefix attention path + triton = tl = None + +HAS_TRITON = triton is not None + + +if HAS_TRITON: + + @triton.jit + def _copy_tile(src, dst, mask): + tl.store(dst, tl.load(src, mask=mask), mask=mask) + + @triton.jit + def _rms_rotate_heads( + src, weight, table_row, dst, mask, eps, HEAD_DIM: tl.constexpr, BLOCK_H: tl.constexpr + ): + """Steps 1-4 for every head of one token; ``src``/``dst`` are [BLOCK_H, D] pointers.""" + tl.static_assert(HEAD_DIM == 128) + x = tl.load(src, mask=mask, other=0.0).to(tl.float32) + + # 1. Lane j is x[8j .. 8j+7]: an in-order FMA chain from zero. Selecting element e of + # each lane is exact (one value summed with zeros). + lanes = tl.reshape(x, [BLOCK_H, HEAD_DIM // 8, 8]) + element = tl.arange(0, 8)[None, None, :] + acc = tl.zeros([BLOCK_H, HEAD_DIM // 8], dtype=tl.float32) + for e in tl.static_range(8): + x_e = tl.sum(tl.where(element == e, lanes, 0.0), axis=2) + acc = tl.fma(x_e, x_e, acc) + # Then adjacent lanes pairwise, 16 -> 8 -> 4 -> 2 -> 1. + lo, hi = tl.split(tl.reshape(acc, [BLOCK_H, 8, 2])) + acc = lo + hi + lo, hi = tl.split(tl.reshape(acc, [BLOCK_H, 4, 2])) + acc = lo + hi + lo, hi = tl.split(tl.reshape(acc, [BLOCK_H, 2, 2])) + acc = lo + hi + lo, hi = tl.split(acc) + ss = lo + hi + + # 2. D is a power of two, so ss * (1/D) equals the IEEE ss / D exactly; Triton's FP32 + # "/" is an approximate division. + inv = tl.rsqrt(ss * (1.0 / HEAD_DIM) + eps) + + # 3. Two FP32 multiplies, one BF16 rounding (nearest-even), back to FP32. + w = tl.load(weight + tl.arange(0, HEAD_DIM)).to(tl.float32) + n = (x * inv[:, None]) * w[None, :] + n = n.to(tl.bfloat16, fp_downcast_rounding="rtne").to(tl.float32) + + # 4. Rotate adjacent pairs by the row's [cos | sin]: separate products, no FMA. + a, b = tl.split(tl.reshape(n, [BLOCK_H, HEAD_DIM // 2, 2])) + half = tl.arange(0, HEAD_DIM // 2) + cos = tl.load(table_row + half).to(tl.float32)[None, :] + sin = tl.load(table_row + HEAD_DIM // 2 + half).to(tl.float32)[None, :] + even = a * cos - b * sin + odd = b * cos + a * sin + out = tl.reshape(tl.join(even, odd), [BLOCK_H, HEAD_DIM]) + tl.store(dst, out.to(tl.bfloat16, fp_downcast_rounding="rtne"), mask=mask) + + @triton.jit + def _fused_joint_qk_norm_rope_kernel( + q0, + k0, + v0, + q1, + k1, + v1, + wq0, + wk0, + wq1, + wk1, + table, + q_out, + k_out, + v_out, + s0, + s1, + heads, + stride0_b, + stride0_s, + stride0_h, + stride1_b, + stride1_s, + stride1_h, + eps, + HEAD_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + ): + """Axis 0: joint output row over all samples; axis 1: 0 = Q, 1 = K, 2 = V. + + A stream's q/k/v share one set of strides (unit last stride). Offsets are int64. + """ + row = tl.program_id(0).to(tl.int64) + tensor = tl.program_id(1) + sample = row // (s0 + s1) + token = row % (s0 + s1) + head = tl.arange(0, BLOCK_H).to(tl.int64)[:, None] + dim = tl.arange(0, HEAD_DIM)[None, :] + mask = (head < heads) & (dim < HEAD_DIM) + dst = (row * heads + head) * HEAD_DIM + dim + table_row = table + row * HEAD_DIM + if token < s0: + src0 = sample * stride0_b + token * stride0_s + head * stride0_h + dim + if tensor == 0: + _rms_rotate_heads( + q0 + src0, wq0, table_row, q_out + dst, mask, eps, HEAD_DIM, BLOCK_H + ) + elif tensor == 1: + _rms_rotate_heads( + k0 + src0, wk0, table_row, k_out + dst, mask, eps, HEAD_DIM, BLOCK_H + ) + else: + _copy_tile(v0 + src0, v_out + dst, mask) + else: + src1 = sample * stride1_b + (token - s0) * stride1_s + head * stride1_h + dim + if tensor == 0: + _rms_rotate_heads( + q1 + src1, wq1, table_row, q_out + dst, mask, eps, HEAD_DIM, BLOCK_H + ) + elif tensor == 1: + _rms_rotate_heads( + k1 + src1, wk1, table_row, k_out + dst, mask, eps, HEAD_DIM, BLOCK_H + ) + else: + _copy_tile(v1 + src1, v_out + dst, mask) + + +def _new_outputs(q0, q1): + joint = (q0.shape[0], q0.shape[1] + q1.shape[1], *q0.shape[2:]) + return tuple(torch.empty(joint, dtype=q0.dtype, device=q0.device) for _ in range(3)) + + +def _with_shared_strides(q, k, v): + """The kernel reads a stream's q/k/v through one set of strides with a unit last stride + (true of views split from one packed projection); anything else is made contiguous.""" + if q.stride() == k.stride() == v.stride() and q.stride(-1) == 1: + return q, k, v + return q.contiguous(), k.contiguous(), v.contiguous() + + +def _launch_kernel(q0, k0, v0, q1, k1, v1, wq0, wk0, wq1, wk1, table, eps): + batch, s0, heads, head_dim = q0.shape + s1 = q1.shape[1] + outputs = _new_outputs(q0, q1) + if outputs[0].numel() == 0: + return outputs + q0, k0, v0 = _with_shared_strides(q0, k0, v0) + q1, k1, v1 = _with_shared_strides(q1, k1, v1) + weights = [w.contiguous() for w in (wq0, wk0, wq1, wk1)] + with torch.cuda.device(q0.device): + _fused_joint_qk_norm_rope_kernel[(batch * (s0 + s1), 3)]( + q0, + k0, + v0, + q1, + k1, + v1, + *weights, + table.contiguous(), + *outputs, + s0, + s1, + heads, + *q0.stride()[:3], + *q1.stride()[:3], + float(eps), + HEAD_DIM=head_dim, + BLOCK_H=triton.next_power_of_2(heads), + num_warps=4, + enable_fp_fusion=False, + ) + return outputs + + +@torch.library.custom_op( + "difflet_hv15::fused_joint_qk_norm_rope", mutates_args=(), device_types="cuda" +) +def fused_joint_qk_norm_rope( + q0: torch.Tensor, + k0: torch.Tensor, + v0: torch.Tensor, + q1: torch.Tensor, + k1: torch.Tensor, + v1: torch.Tensor, + wq0: torch.Tensor, + wk0: torch.Tensor, + wq1: torch.Tensor, + wk1: torch.Tensor, + table: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Joint ``(q, k, v)`` for attention, each a new contiguous ``[B, S0 + S1, H, D]``. + + Stream 0 (video, ``[B, S0, H, D]``) rows come first in each sample, then stream 1 + (context). Q/K are RMS-normalised with their stream's weights (``wq0, wk0`` video, + ``wq1, wk1`` context) and rotated by ``table``'s row for that joint token + (``[B * (S0 + S1), D]``, ``[cos | sin]``, identity rows for context); V is copied. + Inputs are BF16 with ``D == 128``; the caller checks that before taking this path. + """ + return _launch_kernel(q0, k0, v0, q1, k1, v1, wq0, wk0, wq1, wk1, table, eps) + + +@fused_joint_qk_norm_rope.register_fake +def _(q0, k0, v0, q1, k1, v1, wq0, wk0, wq1, wk1, table, eps): + joint = (q0.shape[0], q0.shape[1] + q1.shape[1], *q0.shape[2:]) + return q0.new_empty(joint), q0.new_empty(joint), q0.new_empty(joint) + + +def pack_rope_table( + cos: torch.Tensor, sin: torch.Tensor, context_rows: int, dtype: torch.dtype +) -> torch.Tensor: + """One sample's ``[S_video + context_rows, D]`` table: rows ``[cos | sin]``, one value per + interleaved pair, cast to ``dtype``, then ``context_rows`` identity rows (cos 1, sin 0). + + ``cos``/``sin`` are ``[S_video, D]`` tables that repeat each angle for both members of a + pair, as Diffusers' HunyuanVideo 1.5 RoPE builds them. + """ + half = cos.shape[-1] // 2 + video = torch.cat((cos[:, 0::2], sin[:, 0::2]), dim=-1).to(dtype) + identity = torch.zeros((context_rows, 2 * half), dtype=dtype, device=cos.device) + identity[:, :half] = 1 + return torch.cat((video, identity)) + + +__all__ = ["HAS_TRITON", "fused_joint_qk_norm_rope", "pack_rope_table"] diff --git a/difflet/backends/cuda/hunyuan_video_15/transformer.py b/difflet/backends/cuda/hunyuan_video_15/transformer.py new file mode 100644 index 00000000..92966864 --- /dev/null +++ b/difflet/backends/cuda/hunyuan_video_15/transformer.py @@ -0,0 +1,73 @@ +"""HunyuanVideo 1.5 transformer loading through the shared GPU lifecycle.""" + +from __future__ import annotations + +import torch + +from difflet.backends._torch_megatron.checkpoint import build_weight_map +from difflet.backends.gpu.application import GpuComponent +from difflet.pipeline.parallel_mesh import MeshSpec + + +class CudaHunyuanVideo15TransformerApplication(GpuComponent): + # Timing-based reduction choices can change BF16 trajectories between processes; bind + # stable ones per compiled block, without changing Inductor's global policy. + compile_options = {"deterministic": True} + + def __init__(self, *, model_path, config, device=None, torch_compile=False): + super().__init__( + config=MeshSpec(tp=1), + device=device, + dtype=config.torch_dtype, + torch_compile=torch_compile, + ) + self.model_path = model_path + self.dit_config = config + + def build_module(self): + from accelerate import init_empty_weights + from difflet.backends.cuda.hunyuan_video_15.modeling import ( + HunyuanVideo15Transformer3DModel, + ) + + self._init_runtime() + with init_empty_weights(include_buffers=False): + return HunyuanVideo15Transformer3DModel(self.dit_config).to(dtype=self.dtype) + + def load_weights(self, module): + from safetensors.torch import load_file + + state = {} + for file in sorted(set(build_weight_map(self.model_path).values())): + state.update(load_file(file)) + module.load_state_dict(state, strict=True) + + def apply_load_time_swaps(self, module): + # On the device and before compile: one QKV GEMM per stream for the joint attention + # path (weights become views; memory and keys unchanged). Eager only: the joint path + # trims the context to its valid tokens, so static-shape compiled blocks would + # recompile per prompt length; they keep the prefix path, one graph for every length. + from difflet.backends.cuda.hunyuan_video_15.modeling import joint_attention_supported + + attn = module.transformer_blocks[0].attn + if not self.torch_compile and joint_attention_supported( + attn.to_q.weight, attn.norm_q.weight.shape[0] + ): + module.pack_attention_projections() + + def repeated_blocks(self, module): + return [module.transformer_blocks] + + def get_example_inputs(self): + c = self.dit_config + return ( + torch.zeros( + 1, c.in_channels, c.latent_frames, c.latent_height, c.latent_width, dtype=self.dtype + ), + torch.full((1,), 500.0, dtype=self.dtype), + torch.zeros(1, c.text_seq_len, c.text_embed_dim, dtype=self.dtype), + torch.arange(c.text_seq_len)[None] < min(3, c.text_seq_len), + torch.zeros(1, c.text_2_seq_len, c.text_embed_2_dim, dtype=self.dtype), + torch.zeros(1, c.text_2_seq_len, dtype=torch.bool), + torch.zeros(1, c.image_tokens, c.image_embed_dim, dtype=self.dtype), + ) diff --git a/difflet/cli/main.py b/difflet/cli/main.py index b4334a6b..9c040002 100644 --- a/difflet/cli/main.py +++ b/difflet/cli/main.py @@ -26,6 +26,7 @@ "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "Wan-AI/Wan2.1-T2V-14B-Diffusers", "hunyuanvideo-community/HunyuanVideo", + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", "Qwen/Qwen-Image", "Lightricks/LTX-2", } diff --git a/difflet/models/hunyuan_video_15/__init__.py b/difflet/models/hunyuan_video_15/__init__.py new file mode 100644 index 00000000..d281a709 --- /dev/null +++ b/difflet/models/hunyuan_video_15/__init__.py @@ -0,0 +1,2 @@ +"""HunyuanVideo 1.5 text-to-video: the GPU spec. The CUDA implementation is in +``difflet.backends.cuda.hunyuan_video_15``.""" diff --git a/difflet/models/hunyuan_video_15/spec.py b/difflet/models/hunyuan_video_15/spec.py new file mode 100644 index 00000000..cfd2ca7a --- /dev/null +++ b/difflet/models/hunyuan_video_15/spec.py @@ -0,0 +1,51 @@ +"""HunyuanVideo 1.5 text-to-video on the GPU flow.""" + +from difflet.models.spec import ModelSpec + +MODEL_ID = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + +SPEC = ModelSpec( + model_type="hunyuan_video_15", + display_name="HunyuanVideo 1.5", + short_name="hunyuan15", + cli_model_ids=(MODEL_ID,), + served_checkpoint_ids=(MODEL_ID,), + output="video", + default_height=480, + default_width=848, + default_num_frames=121, + default_steps=50, + default_fps=24, + guidance="cfg", + default_guidance_scale=6.0, + accepts_negative_prompt=True, + accepts_cfg_parallel=False, + refused_cli_flags=( + ("shapes", "--shapes"), + ("teacache_cadence", "--teacache-cadence"), + ("teacache_speedup", "--teacache-speedup"), + ("teacache_online_delta", "--teacache-online-delta"), + ("teacache_calibration", "--teacache-calibration"), + ("host_vae", "--host-vae"), + ("taef1", "--taef1"), + ("taef1_path", "--taef1-path"), + ("true_cfg_scale", "--true-cfg-scale"), + ("guidance_scale_2", "--guidance-scale-2"), + ), + refused_serving_options=( + ("teacache_speedup", "--teacache-speedup"), + ("host_vae", "--host-vae"), + ), + refused_video_fields=( + ("true_cfg_scale", "use guidance_scale for standard CFG"), + ("guidance_scale_2", "there is one transformer"), + ("boundary_ratio", "there is one transformer"), + ("flow_shift", "the scheduler configuration comes from the checkpoint"), + ), + refused_application_kwargs=("shapes", "host_vae", "taef1", "taef1_path"), + refused_application_kwarg_prefixes=("teacache_",), + refused_application_kwargs_reason="one resident shape with standard CFG", + # No prompt budgets: Diffusers' driver crops its own chat template and glyph text. + smoke_prompt="a cat walking across a sunlit garden", + application="difflet.backends.cuda.hunyuan_video_15.application:CudaHunyuanVideo15Application", +) diff --git a/difflet/models/spec.py b/difflet/models/spec.py index 771a6772..3e58a3ab 100644 --- a/difflet/models/spec.py +++ b/difflet/models/spec.py @@ -85,6 +85,7 @@ class ModelSpec: "wan": "difflet.models.wan.spec", "hunyuan_video": "difflet.models.hunyuan_video.spec", "ltx_2": "difflet.models.ltx_2.spec", + "hunyuan_video_15": "difflet.models.hunyuan_video_15.spec", } diff --git a/difflet/registry.py b/difflet/registry.py index 70487a7e..91a5e15b 100644 --- a/difflet/registry.py +++ b/difflet/registry.py @@ -369,7 +369,7 @@ def is_hunyuan_video_15(model_id: str) -> bool: supports_cp=True, supports_sp=True, ), - backends=("trainium", "neuron"), + backends=("trainium", "neuron", "cuda"), ) class _HunyuanVideo15Registration: pass diff --git a/difflet/serving/model_registry.py b/difflet/serving/model_registry.py index 7cf1f774..113188f6 100644 --- a/difflet/serving/model_registry.py +++ b/difflet/serving/model_registry.py @@ -41,27 +41,18 @@ def resolve_serving_model(options: ServeOptions) -> ResolvedServingModel: # TPU only: every other backend keeps the path it had. if current_backend() == "tpu": entry.require_backend("tpu") - try: - metadata = _SERVING_METADATA[entry.name] - except KeyError as exc: - raise ValueError(f"model type {entry.name!r} is not enabled for serving") from exc - normalized = options.model_id.rstrip("/") - if normalized not in metadata.checkpoint_ids: - allowed = ", ".join(metadata.checkpoint_ids) - raise ValueError( - f"checkpoint {options.model_id!r} is not enabled for P0 serving; " - f"allowed for {entry.name}: {allowed}" - ) from difflet.platforms import current_platform tp_degree = options.tp_degree platform = current_platform() on_gpu = platform.capabilities.flow == "gpu" - if on_gpu: - # The GPU flow's refusals are DiffletServingErrors, so `difflet serve` exits with the - # message. - from difflet.serving.errors import feature_not_supported + normalized = options.model_id.rstrip("/") + # The GPU flow's refusals are DiffletServingErrors, so `difflet serve` exits with the + # message. + from difflet.serving.errors import feature_not_supported + if on_gpu: + # A GPU-flow model's serving metadata comes from its GPU spec. metadata = _gpu_serving_metadata(entry.name, options.model_id, platform) if normalized not in metadata.checkpoint_ids: allowed = ", ".join(metadata.checkpoint_ids) @@ -79,11 +70,32 @@ def resolve_serving_model(options: ServeOptions) -> ResolvedServingModel: f"serving; on the {platform.display_name} backend TeaCache is a `difflet run` " "option (Qwen-Image's fixed cadence) and not served — drop the flag" ) + else: + try: + metadata = _SERVING_METADATA[entry.name] + except KeyError as exc: + # A model served only on GPUs (HunyuanVideo 1.5) is refused here, not with a + # traceback. + raise feature_not_supported( + f"model type {entry.name!r} is not enabled for serving on the " + f"{platform.display_name} backend" + ) from exc + if normalized not in metadata.checkpoint_ids: + allowed = ", ".join(metadata.checkpoint_ids) + raise ValueError( + f"checkpoint {options.model_id!r} is not enabled for P0 serving; " + f"allowed for {entry.name}: {allowed}" + ) _refuse_unsupported_replicas(int(options.replicas), platform) + cfg_parallel = options.cfg_parallel if on_gpu and tp_degree is None: # The registry default (tp=4) is a Trainium topology; the GPU CLI uses # `args.tp_degree or 1`. tp_degree = platform.default_tp_degree() + if on_gpu and cfg_parallel is None: + # Likewise a registry CFG-parallel layout is Trainium's; `--cfg-parallel` is opt-in on + # the GPU CLI, so it defaults off here too. + cfg_parallel = False profile = build_serving_profile( model_id=options.model_id, model_type=metadata.model_type, @@ -97,7 +109,7 @@ def resolve_serving_model(options: ServeOptions) -> ResolvedServingModel: tp_degree=tp_degree, cp_degree=options.cp_degree, cp_mode=options.cp_mode, - cfg_parallel=options.cfg_parallel, + cfg_parallel=cfg_parallel, sp_enabled=options.sp_enabled, height=options.height, width=options.width, diff --git a/docs/verification/hunyuanvideo15-cuda.md b/docs/verification/hunyuanvideo15-cuda.md new file mode 100644 index 00000000..a100d9f6 --- /dev/null +++ b/docs/verification/hunyuanvideo15-cuda.md @@ -0,0 +1,170 @@ +# HunyuanVideo 1.5 on CUDA + +Experimental single-GPU BF16 text-to-video for +`hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, through Difflet's GPU +application, Python API, CLI and resident HTTP server. It is separate from the Trainium +HunyuanVideo 1.5 path and from the original HunyuanVideo; support does not carry over +between them. + +## Usage + +Install the CUDA stack as in the [README](../../README.md#install), then: + +```bash +DIFFLET_BACKEND=cuda difflet run \ + --model-id hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --height 480 --width 848 --num-frames 121 \ + --steps 50 --guidance-scale 6 --seed 0 \ + --prompt "a cat walking across a sunlit garden" --output video.mp4 + +DIFFLET_BACKEND=cuda difflet serve \ + --model-id hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --height 480 --width 848 --num-frames 121 --request-timeout 1800 +``` + +A 121-frame request takes about seven minutes on an RTX PRO 6000, longer than the server's +default 300 s request timeout, hence `--request-timeout`. + +Add `--torch-compile` to compile the 54 transformer blocks; the token refiner, text encoders +and VAE stay eager. In Python, pass `parallel=DiffletParallelConfig(tp_degree=1)`: without it +the registry's Trainium layout (TP2 × CFG2) is used and refused. Set +`application_kwargs={"torch_compile": True}` to compile. +`pipe(prompt=..., num_inference_steps=50, guidance_scale=6, seed=0)` returns CPU FP32 video in +`(B, C, F, H, W)` layout, scaled to `[-1, 1]`. + +Each loaded application has one shape and generates one prompt at a time. Height and width +must be multiples of 16, and frames `4n+1`. Defaults are 480×848×121, 50 steps, guidance 6 +and 24 fps. Classifier-free guidance runs its two branches one after the other on the same +GPU. Negative prompts work through `difflet run --negative-prompt`, `--requests` batches, +`--dp` runs, the Python API and the server. + +Image-to-video, the 720p checkpoints, CFG-parallel, TP/CP/SP, TeaCache, host VAE, TAEF1 and +shape buckets are not supported. DP replicas have not been measured. + +## How it runs + +- **Transformer.** Diffusers' HunyuanVideo 1.5 building blocks keep the checkpoint layout, + normalization and cast boundaries; Difflet replaces the projections and attention. The + checkpoint loads strictly. +- **Eager attention.** When the model loads eagerly in BF16 on a GPU with cuDNN attention and + Triton, each block of a one-sample request projects Q, K and V with one GEMM per stream, + trims the text context to its valid tokens, applies Q/K RMSNorm and RoPE in one fused Triton + op (`difflet_hv15::fused_joint_qk_norm_rope`, written for Difflet) and attends with cuDNN + SDPA. Otherwise the blocks use the prefix attention below. +- **Prefix attention.** Compiled blocks, and eager blocks that cannot take the path above, + attend each sample's valid token prefix, so one static-shape graph serves every prompt + length. In compiled blocks, Inductor's reductions are deterministic. +- **Conditioning.** Diffusers' driver encodes the prompts. Difflet recomputes the Qwen rotary + frequencies on the GPU and turns cuDNN SDPA off while prompts are encoded, which matches + vLLM-Omni's conditioning bit for bit. +- **Decoding.** Diffusers' native VAE tiles: 256 px with 25% overlap, which keeps a 121-frame + decode on one GPU. Tiled and untiled decodes differ slightly; compare outputs with the same + decoder settings. + +## Matched comparison with vLLM-Omni + +### Setup + +- **Code:** the HunyuanVideo 1.5 source as committed with this guide. +- **Reference:** vLLM-Omni `08177216` (vLLM 0.32.0), unmodified. +- **Runtime, shared by both:** Python 3.12, Torch 2.13.0+cu130, Triton 3.7.1, + Diffusers 0.40.0 and Transformers 5.19.0. +- **Hardware:** one NVIDIA RTX PRO 6000 Blackwell Server Edition (96 GB). +- **Procedure:** every engine and mode ran in a fresh process. The whole set ran twice, the + second time in reverse order. + +The two requests: + +- **G**, vLLM-Omni's own test settings: 832×480, 33 frames, 50 steps, guidance 6, flow shift 5, + seed 42, prompt "A cat walks on the grass, realistic style." +- **P**, the request in [Usage](#usage): 848×480, 121 frames, 50 steps, guidance 6, seed 0. + +Both engines start from bit-identical initial latents. The decoders differ: Difflet decodes +BF16 latents with a BF16 VAE in native tiles; vLLM-Omni decodes FP32 latents with an FP32 VAE, +untiled for G (its default; the decode peaks at 92.9 GiB of the card's 96 GB) and with its VAE +tiling on for P. + +### Output agreement + +Videos were written with Diffusers `export_to_video` at 24 fps and scored with ffmpeg SSIM and +PSNR through vLLM-Omni's own accuracy-test helper. + +The acceptance bar is relative. On G, for each mode, Difflet against vLLM-Omni must score at +least what vLLM-Omni's own eager and compiled outputs score against each other. P is reported +but does not gate. vLLM-Omni's own accuracy test uses a fixed bar of SSIM ≥ 0.94 and PSNR ≥ +28 dB: Difflet's eager G meets it, and compiled G (0.9181) does not. + +| Request | Mode | Difflet vs vLLM-Omni (SSIM / PSNR) | vLLM-Omni eager vs compiled | Gate | +|---|---|---|---|---| +| G | eager | 0.9711 / 39.95 dB | 0.8412 / 27.51 dB | pass | +| G | compiled | 0.9181 / 31.65 dB | 0.8412 / 27.51 dB | pass | +| P | eager | 0.9555 / 35.46 dB | 0.4999 / 18.04 dB | reported | +| P | compiled | 0.6465 / 21.16 dB | 0.4999 / 18.04 dB | reported | + +**Eager.** The denoiser is bit-identical to vLLM-Omni's: for the first step, in both guidance +branches, all 54 blocks and the final output match exactly. The remaining difference comes from +the VAE decode. Decoding vLLM-Omni's latents with Difflet's decoder reproduces Difflet's video +exactly, and rounding vLLM-Omni's FP32 latents to BF16 before its own decoder already scores +SSIM 0.9747 / 40.59 dB against its own video. + +**Compiled.** The two engines compile differently, so their compiled trajectories differ from +block 0: vLLM-Omni compiles regions with dynamic shapes, Difflet compiles each block with static +shapes and deterministic reductions. Each engine's eager and compiled outputs differ from each +other by a similar amount: Difflet 0.8615 / 28.35 dB, vLLM-Omni 0.8412 / 27.51 dB on G. + +An earlier run of the same code, on the same card model, gave +byte-identical Difflet videos in both modes and for both requests. Only vLLM-Omni's compiled +videos changed, which moved its eager-vs-compiled G score to 0.8343 / 27.46 dB and Difflet's +compiled G score against it to 0.9041 / 30.20 dB: the bar passed there too. + +**Repeatability.** Each engine and mode reproduced its first-pass denoiser digests exactly in +the second pass, in both guidance branches, and gave the same video scores. + +**Fused op.** Difflet's fused Q/K op was compared directly with vLLM-Omni's fused op on the GPU, +at the G and P sequence lengths, with 13 and 6 context tokens and three seeds each: Q, K and V +were bit-identical in 12 of 12 cases. + +### Latency, start-up and memory + +Warm latency is CUDA-synchronized from prompt to decoded frames: the median of three G +requests, and one P request after a warm-up. The speed bar is Difflet / vLLM-Omni ≤ 1.10. + +| Mode | Request | Difflet (s), pass 1 / 2 | vLLM-Omni (s), pass 1 / 2 | Ratio, pass 1 / 2 | +|---|---|---|---|---| +| eager | G | 58.83 / 58.85 | 58.81 / 58.85 | 1.000 / 1.000 | +| eager | P | 412.2 / 412.4 | 418.5 / 418.8 | 0.985 / 0.985 | +| compiled | G | 58.58 / 58.65 | 54.11 / 54.14 | 1.083 / 1.083 | +| compiled | P | 408.1 / 407.8 | 398.4 / 397.5 | 1.024 / 1.026 | + +Start-up and memory are recorded but do not gate: + +- **Load time** (process start to loaded model): Difflet 15–18 s eager and 17–26 s compiled, + vLLM-Omni 25–28 s. +- **Compilation** (time inside TorchDynamo, pass 1 with empty compile caches / pass 2): Difflet + 5.4 / 0.8 s, while loading; vLLM-Omni 9.4 / 3.3 s, during its first request. +- **Peak device memory**, over the G and P requests: + +| Peak memory | Difflet | vLLM-Omni | +|---|---|---| +| Device (NVML) | 62.9 GiB | 92.9 GiB | +| Torch allocated | 45.3 GiB | 72.5 GiB | + +### The pinned CUDA environment + +vLLM 0.32.0 needs Torch 2.13, so the comparison above uses that runtime. On the same card, +Difflet was also run in the pinned environment of `requirements-cuda.lock` (the README's Torch +2.11/cu128 stack): Python 3.12, Torch 2.11.0+cu128, Triton 3.6.0, Diffusers 0.38.0 and +Transformers 5.20.1. + +- `tests/unit/backends/test_cuda_hunyuan15_joint_attention.py` passed on the GPU, 20 of 20, + including the fused kernel against its reference. +- In eager mode, `difflet run` wrote a 33-frame G video at 24 fps (59.4 s to generate), and + `difflet serve` returned a G video over HTTP (startup 26 s, inference 59.4 s). +- Warm latency was within 3% of the Torch 2.13 runs: eager G 59.13 s and P 413.6 s, compiled G + 60.24 s and P 420.3 s. +- Initial latents were identical to the Torch 2.13 run, but the Qwen prompt encoder's output + differs between the two stacks, so the videos differ: eager G scores 0.6880 / 22.32 dB and + compiled G 0.8140 / 27.00 dB against Difflet's Torch 2.13 videos. + +These results cover this checkpoint, these two requests and this card. They do not establish +broader model support. diff --git a/tests/unit/backends/test_cuda_hunyuan15_conditioning.py b/tests/unit/backends/test_cuda_hunyuan15_conditioning.py new file mode 100644 index 00000000..94df3d0a --- /dev/null +++ b/tests/unit/backends/test_cuda_hunyuan15_conditioning.py @@ -0,0 +1,184 @@ +"""Device-local rotary initialization and request-scoped Qwen attention policy.""" + +from types import SimpleNamespace +import importlib + +import pytest +import torch +import transformers + +# The CUDA stack's Transformers (5.x) has the Qwen2.5-VL rotary API the conditioning uses; +# the Trainium stacks pin 4.x. +pytestmark = pytest.mark.skipif( + int(transformers.__version__.split(".")[0]) < 5, + reason="HunyuanVideo 1.5 CUDA conditioning needs Transformers 5", +) + + +def conditioning(): + return importlib.import_module("difflet.backends.cuda.hunyuan_video_15.conditioning") + + +def tiny_qwen(): + from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLTextConfig + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLTextModel + + config = Qwen2_5_VLTextConfig( + vocab_size=16, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + rope_parameters={ + "rope_type": "default", + "rope_theta": 1000000.0, + "mrope_section": [1, 1, 2], + }, + ) + return Qwen2_5_VLTextModel(config).eval() + + +def test_rotary_initialization_uses_model_device_without_changing_parameters(): + m = conditioning() + qwen = tiny_qwen() + before = {k: p.detach().clone() for k, p in qwen.named_parameters()} + config = qwen.config.to_dict() + expected, scale = qwen.rotary_emb.compute_default_rope_parameters(qwen.rotary_emb.config) + m.initialize_qwen_rotary(qwen) + for name in ("inv_freq", "original_inv_freq"): + value = getattr(qwen.rotary_emb, name) + assert value.dtype == torch.float32 and value.device == next(qwen.parameters()).device + assert torch.equal(value, expected) + assert qwen.rotary_emb._buffers[name] is value + assert name in qwen.rotary_emb._non_persistent_buffers_set + assert qwen.rotary_emb.attention_scaling == scale + assert qwen.config.to_dict() == config + assert all(torch.equal(p, before[k]) for k, p in qwen.named_parameters()) + + +def test_actual_qwen_encoding_uses_policy_and_restores_it(monkeypatch): + m = conditioning() + qwen = tiny_qwen() + seen = [] + + def encode_prompt(input_ids): + seen.append(torch.backends.cuda.cudnn_sdp_enabled()) + return qwen( + input_ids=input_ids, attention_mask=torch.ones_like(input_ids), use_cache=False + ).last_hidden_state + + pipe = SimpleNamespace(text_encoder=qwen, encode_prompt=encode_prompt) + old = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(True) + try: + m.configure_conditioning(pipe) + got = pipe.encode_prompt(torch.tensor([[1, 2, 3]])) + assert got.shape == (1, 3, 16) and torch.isfinite(got).all() + assert seen == [False] and torch.backends.cuda.cudnn_sdp_enabled() + finally: + torch.backends.cuda.enable_cudnn_sdp(old) + + +@pytest.mark.parametrize("initial", [False, True]) +def test_conditioning_policy_restores_on_exception(initial): + m = conditioning() + qwen = tiny_qwen() + + def broken(): + assert not torch.backends.cuda.cudnn_sdp_enabled() + raise RuntimeError("encoding failed") + + pipe = SimpleNamespace(text_encoder=qwen, encode_prompt=broken) + old = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(initial) + try: + m.configure_conditioning(pipe) + with pytest.raises(RuntimeError, match="encoding failed"): + pipe.encode_prompt() + assert torch.backends.cuda.cudnn_sdp_enabled() == initial + finally: + torch.backends.cuda.enable_cudnn_sdp(old) + + +def test_application_load_wires_policy_for_direct_pipe_and_generate(monkeypatch): + from difflet.backends.cuda.hunyuan_video_15.application import CudaHunyuanVideo15Application + import diffusers + + seen = [] + qwen = tiny_qwen() + + class VAE(torch.nn.Identity): + def enable_tiling(self): + pass + + class Pipe: + def __init__(self, **components): + self.components = components + self.__dict__.update(components) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + return cls( + text_encoder=qwen, + text_encoder_2=torch.nn.Identity(), + vae=VAE(), + tokenizer=object(), + guider=SimpleNamespace(config={}), + ) + + def encode_prompt(self, *args, **kwargs): + seen.append(("encode", torch.backends.cuda.cudnn_sdp_enabled())) + + def __call__(self, **kwargs): + self.encode_prompt(kwargs["prompt"]) + seen.append(("after_encode", torch.backends.cuda.cudnn_sdp_enabled())) + return SimpleNamespace(frames=torch.zeros(1)) + + monkeypatch.setattr(diffusers, "HunyuanVideo15Pipeline", Pipe) + app = object.__new__(CudaHunyuanVideo15Application) + app.model_path = "unused" + app.dtype = torch.float32 + app.device = torch.device("cpu") + app.local_files_only = True + app.config = SimpleNamespace(out_channels=2) + app.adapter = object() + app.shape = dict(height=480, width=848, num_frames=121) + old = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(True) + try: + app.load_companions() + app.pipe.encode_prompt("a cat") + app.generate("a cat", output_type="latent") + assert seen == [("encode", False), ("encode", False), ("after_encode", True)] + assert torch.backends.cuda.cudnn_sdp_enabled() + finally: + torch.backends.cuda.enable_cudnn_sdp(old) + + +def test_rotary_initializer_uses_encoder_device_instead_of_ambient_context(): + m = conditioning() + qwen = tiny_qwen() + with torch.device("meta"): + m.initialize_qwen_rotary(qwen) + assert qwen.rotary_emb.inv_freq.device.type == "cpu" + assert torch.isfinite(qwen.rotary_emb.inv_freq).all() + + +def test_conditioning_preserves_other_sdpa_flags_and_other_instances(): + m = conditioning() + states = lambda: ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + expected = states() + seen = [] + configured = SimpleNamespace( + text_encoder=tiny_qwen(), encode_prompt=lambda: seen.append(states()) + ) + untouched = SimpleNamespace(encode_prompt=lambda: torch.backends.cuda.cudnn_sdp_enabled()) + before = torch.backends.cuda.cudnn_sdp_enabled() + m.configure_conditioning(configured) + configured.encode_prompt() + assert seen == [expected] and states() == expected and untouched.encode_prompt() == before diff --git a/tests/unit/backends/test_cuda_hunyuan15_joint_attention.py b/tests/unit/backends/test_cuda_hunyuan15_joint_attention.py new file mode 100644 index 00000000..d92dae71 --- /dev/null +++ b/tests/unit/backends/test_cuda_hunyuan15_joint_attention.py @@ -0,0 +1,449 @@ +"""HunyuanVideo 1.5 joint attention and its fused Q/K RMSNorm + RoPE op. + +The op is CUDA-only. These tests register a PyTorch reference of the kernel's arithmetic as +its CPU kernel, test that reference and the model wiring on the CPU, and compare the Triton +kernel with the reference on CUDA. +""" + +import contextlib +import importlib +import os + +import pytest +import torch +from diffusers.models.embeddings import apply_rotary_emb + +from tests.unit.backends.test_cuda_hunyuan15_transformer import TINY, _inputs, _pair + +joint = modeling = None +_LANE_WIDTH = 8 + + +# The kernel's arithmetic in PyTorch, step by step as the op's docstring states it. + + +def _fma_fp32(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + """FP32 ``fma(a, b, c)`` with a single rounding, on any device. + + The FP32 product is exact in FP64 (48 significant bits). The FP64 sum is made + round-to-odd from its exact TwoSum error, and rounding a round-to-odd value with + 53 >= 24 + 2 bits to FP32 equals rounding the exact ``a * b + c`` once. + """ + p = a.double() * b.double() + c = c.double() + s = p + c + t = s - p + err = (p - (s - t)) + (c - t) + bits = s.view(torch.int64) + even_and_inexact = (err != 0) & ((bits & 1) == 0) & torch.isfinite(s) + toward_exact = torch.where((err > 0) == (s > 0), 1, -1) # one ulp up or down in magnitude + return torch.where(even_and_inexact, bits + toward_exact, bits).view(torch.float64).float() + + +def _sum_of_squares(x: torch.Tensor) -> torch.Tensor: + """Step 1 over the last dim of FP32 ``x``: lanes of eight FMA chains, then the pair tree.""" + lanes = x.unflatten(-1, (x.shape[-1] // _LANE_WIDTH, _LANE_WIDTH)) + acc = torch.zeros_like(lanes[..., 0]) + for e in range(_LANE_WIDTH): + acc = _fma_fp32(lanes[..., e], lanes[..., e], acc) + while acc.shape[-1] > 1: + acc = acc[..., 0::2] + acc[..., 1::2] + return acc[..., 0] + + +def _normalize(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Steps 1-3: FP32 values of ``(x * inv) * w`` rounded once to ``x.dtype``.""" + xf = x.float() + inv = torch.rsqrt(_sum_of_squares(xf) / xf.shape[-1] + eps) + return ((xf * inv[..., None]) * weight.float()).to(x.dtype).float() + + +def _rotate_pairs(n: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Step 4 in FP32: ``(a, b) -> (a*c - b*s, b*c + a*s)`` on adjacent pairs.""" + a, b = n[..., 0::2], n[..., 1::2] + even = a * cos - b * sin + odd = b * cos + a * sin + return torch.stack((even, odd), dim=-1).flatten(-2) + + +def _reference(q0, k0, v0, q1, k1, v1, wq0, wk0, wq1, wk1, table, eps): + """The op in PyTorch, rounding to the inputs' dtype where the kernel rounds to BF16: + identical for BF16 inputs; for FP32 inputs, the Diffusers chain to FP32 rounding.""" + joint_shape = (q0.shape[0], q0.shape[1] + q1.shape[1], *q0.shape[2:]) + outputs = tuple(torch.empty(joint_shape, dtype=q0.dtype, device=q0.device) for _ in range(3)) + batch, joint_len, _, head_dim = joint_shape + rows = table.float().reshape(batch, joint_len, 1, head_dim) + cos, sin = rows[..., : head_dim // 2], rows[..., head_dim // 2 :] + for out, x0, x1, w0, w1 in zip(outputs, (q0, k0), (q1, k1), (wq0, wk0), (wq1, wk1)): + normed = torch.cat((_normalize(x0, w0, eps), _normalize(x1, w1, eps)), dim=1) + out.copy_(_rotate_pairs(normed, cos, sin)) # one rounding to the output dtype + s0 = q0.shape[1] + outputs[2][:, :s0].copy_(v0) + outputs[2][:, s0:].copy_(v1) + return outputs + + +_registered = False + + +@pytest.fixture(autouse=True) +def _cpu_ops(monkeypatch): + # Bind difflet.ops to the CPU ops before the modules import it (a host with no + # backend set may default to Trainium), then import them lazily. + global joint, modeling, _registered + monkeypatch.setenv("DIFFLET_BACKEND", os.environ.get("DIFFLET_BACKEND") or "cpu") + from difflet.ops import init_parallel_mesh + from difflet.pipeline.parallel_mesh import MeshSpec + + init_parallel_mesh(MeshSpec(tp=1)) + joint = importlib.import_module("difflet.backends.cuda.hunyuan_video_15.qk_norm_rope") + modeling = importlib.import_module("difflet.backends.cuda.hunyuan_video_15.modeling") + if not _registered: + joint.fused_joint_qk_norm_rope.register_kernel("cpu")(_reference) + _registered = True + # The joint path pins cuDNN SDPA, which CPU tensors cannot use. + monkeypatch.setattr(modeling, "sdpa_kernel", lambda backends: contextlib.nullcontext()) + + +def _tables(s, d, dtype=torch.float32, seed=5): + g = torch.Generator().manual_seed(seed) + angles = torch.rand(s, d // 2, generator=g, dtype=torch.float64) * 6.28 + cos = angles.cos().to(dtype).repeat_interleave(2, dim=-1) + sin = angles.sin().to(dtype).repeat_interleave(2, dim=-1) + return cos, sin + + +def test_pack_rope_table_is_the_native_layout(): + cos, sin = _tables(6, 8) + table = joint.pack_rope_table(cos, sin, 3, torch.bfloat16) + assert table.shape == (9, 8) and table.dtype == torch.bfloat16 + assert torch.equal(table[:6, :4], cos[:, 0::2].bfloat16()) + assert torch.equal(table[:6, 4:], sin[:, 0::2].bfloat16()) + assert torch.equal(table[6:, :4], torch.ones(3, 4, dtype=torch.bfloat16)) + assert torch.equal(table[6:, 4:], torch.zeros(3, 4, dtype=torch.bfloat16)) + assert joint.pack_rope_table(cos, sin, 0, torch.float32).shape == (6, 8) + + +def test_joint_op_reference_equals_per_stream_norm_cat_and_rope_in_fp32(): + # FP32 table: the op's eager reference is RMSNorm per stream, cat, interleaved RoPE + # on the video rows and identity on the context rows (the Diffusers chain's math). + g = torch.Generator().manual_seed(11) + s0, s1, h, d = 6, 3, 2, 8 + q0, k0, v0 = (torch.randn(1, s0, h, d, generator=g) for _ in range(3)) + q1, k1, v1 = (torch.randn(1, s1, h, d, generator=g) for _ in range(3)) + w = [torch.rand(d, generator=g) + 0.5 for _ in range(4)] + cos, sin = _tables(s0, d) + table = joint.pack_rope_table(cos, sin, s1, torch.float32) + q, k, v = joint.fused_joint_qk_norm_rope(q0, k0, v0, q1, k1, v1, *w, table, 1e-6) + assert q.shape == k.shape == v.shape == (1, s0 + s1, h, d) + norm = lambda x, weight: torch.nn.functional.rms_norm(x, (d,), weight, 1e-6) # noqa: E731 + expected_q = torch.cat( + (apply_rotary_emb(norm(q0, w[0]), (cos, sin), sequence_dim=1), norm(q1, w[2])), dim=1 + ) + expected_k = torch.cat( + (apply_rotary_emb(norm(k0, w[1]), (cos, sin), sequence_dim=1), norm(k1, w[3])), dim=1 + ) + torch.testing.assert_close(q, expected_q, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(k, expected_k, rtol=1e-6, atol=1e-6) + assert torch.equal(v, torch.cat((v0, v1), dim=1)) + + +def test_joint_op_is_one_opaque_op_under_compile(): + s0, s1, h, d = 6, 3, 2, 8 + streams = [torch.randn(1, s, h, d) for s in (s0, s0, s0, s1, s1, s1)] + weights = [torch.ones(d) for _ in range(4)] + cos, sin = _tables(s0, d) + table = joint.pack_rope_table(cos, sin, s1, torch.float32) + + def call(*args): + q, k, v = joint.fused_joint_qk_norm_rope(*args, table, 1e-6) + return q + k + v + + torch._dynamo.reset() + compiled = torch.compile(call, backend="eager", fullgraph=True) + torch.testing.assert_close(compiled(*streams, *weights), call(*streams, *weights)) + # The fake (new_empty joint-shaped tensors) agrees with the real op, and the schema + # declares no mutation or aliasing. + torch.library.opcheck( + joint.fused_joint_qk_norm_rope, + (*streams, *weights, table, 1e-6), + test_utils=("test_schema", "test_faketensor"), + ) + + +def _packed(model): + model.pack_attention_projections() + return model + + +def test_packing_keeps_keys_values_and_memory(): + _, candidate = _pair() + before = {k: v.clone() for k, v in candidate.state_dict().items()} + _packed(candidate) + after = candidate.state_dict() + assert after.keys() == before.keys() # packed buffers are not persistent + assert all(torch.equal(after[k], before[k]) for k in before) + attn = candidate.transformer_blocks[0].attn + packed = attn._packed_qkv_weight + assert packed.shape == (3 * attn.to_q.weight.shape[0], attn.to_q.weight.shape[1]) + # The per-projection weights are views of the packed tensor: no second copy. + base = packed.untyped_storage().data_ptr() + for name in ("to_q", "to_k", "to_v"): + assert getattr(attn, name).weight.untyped_storage().data_ptr() == base + assert attn.add_v_proj.bias.untyped_storage().data_ptr() == ( + attn._packed_added_qkv_bias.untyped_storage().data_ptr() + ) + + +def _random_norms(*models): + # Diffusers initialises RMSNorm weights to ones; randomise them so the test sees the + # four norm weights' wiring (and eps), identically in both models. + g = torch.Generator().manual_seed(23) + for blocks in zip(*(m.transformer_blocks for m in models)): + for name in ("norm_q", "norm_k", "norm_added_q", "norm_added_k"): + w = torch.rand(getattr(blocks[0].attn, name).weight.shape, generator=g) + 0.5 + for block in blocks: + getattr(block.attn, name).weight.data.copy_(w) + + +@pytest.mark.parametrize("valid", [(5, 4), (2, 0), (3, 2)]) +def test_joint_path_equals_the_prefix_path_in_fp32(valid, monkeypatch): + _, prefix = _pair() + _, candidate = _pair() + _random_norms(prefix, candidate) + assert all(b.attn.norm_q.eps == 1e-6 for b in candidate.transformer_blocks) + _packed(candidate) # packed projections take the joint path, here on the CPU + calls = [] + real = modeling.fused_joint_qk_norm_rope + monkeypatch.setattr( + modeling, "fused_joint_qk_norm_rope", lambda *a: calls.append(a[3].shape[1]) or real(*a) + ) + inputs = _inputs(text_valid=valid[0], glyph_valid=valid[1]) + with torch.no_grad(): + expected = prefix(**inputs, return_dict=False)[0] + actual = candidate(**inputs, return_dict=False)[0] + # One call per block, each with exactly the valid context rows (the native trim). + assert calls == [valid[0] + valid[1]] * TINY["num_layers"] + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_joint_path_needs_packed_projections_and_one_sample(monkeypatch): + _, candidate = _pair() + calls = [] + monkeypatch.setattr(modeling, "fused_joint_qk_norm_rope", lambda *a: calls.append(1)) + with torch.no_grad(): + candidate(**_inputs(), return_dict=False) # not packed: prefix path + _packed(candidate) + candidate(**_inputs(batch=2), return_dict=False) # two samples: prefix path + assert calls == [] + + +@pytest.mark.parametrize("torch_compile", [False, True]) +@pytest.mark.parametrize("supported", [False, True]) +def test_load_packs_only_eager_blocks_on_a_supported_device(monkeypatch, torch_compile, supported): + # The load-time swap hook runs on the device, before compile (GpuComponent). Compiled + # blocks keep the prefix path: a trimmed context would recompile per prompt length. + from difflet.backends.cuda.hunyuan_video_15.transformer import ( + CudaHunyuanVideo15TransformerApplication, + ) + + _, candidate = _pair() + monkeypatch.setattr(modeling, "joint_attention_supported", lambda weight, d: supported) + app = CudaHunyuanVideo15TransformerApplication.__new__(CudaHunyuanVideo15TransformerApplication) + app.torch_compile = torch_compile + app.apply_load_time_swaps(candidate) + packed = hasattr(candidate.transformer_blocks[0].attn, "_packed_qkv_weight") + assert packed == (supported and not torch_compile) + + +def _f32(value): + """Round a Python float (a double) to FP32, nearest-even.""" + return torch.tensor(value, dtype=torch.float32).item() + + +def _explicit_sum_of_squares(row): + # Every double operation is exact for the rows used here, so each _f32 is the single + # rounding of one lane FMA or one tree add. + lanes = [] + for j in range(len(row) // 8): + acc = 0.0 + for x in row[8 * j : 8 * j + 8]: + acc = _f32(acc + x * x) + lanes.append(acc) + while len(lanes) > 1: + lanes = [_f32(lanes[i] + lanes[i + 1]) for i in range(0, len(lanes), 2)] + return lanes[0] + + +def test_sum_of_squares_is_lane_fmas_then_a_pairwise_tree(): + d = 32 # four lanes + # Lane 0 is 2^24 then seven ties that round back to 2^24; the other lanes are 1 each, so + # the tree gives (2^24 + 1) + (1 + 1) = 2^24 + 2 (any other order or lane split differs). + order = [2.0**12] + [1.0] * 7 + ([1.0] + [0.0] * 7) * 3 + # Lane 0: 2^-24, 2^-23, then fma(1 + 2^-12, 1 + 2^-12, 2^-23): one rounding of + # 1 + 2^-11 + 3 * 2^-24 (a tie, to even). Multiply-then-add would give 1 + 2^-11 + 2^-23. + fused = [2.0**-12, 2.0**-12, 1 + 2.0**-12] + [0.0] * 29 + g = torch.Generator().manual_seed(17) + small = (torch.randint(-8, 9, (2, d), generator=g) / 4).tolist() # exact everywhere + # Lanes [2^24, 1, 0, 1]: adjacent-first gives (2^24 + 1) + (0 + 1) = 2^24, while a + # halves-first (shuffle-down) pairing gives (2^24 + 0) + (1 + 1) = 2^24 + 2. + halves = [2.0**12] + [1.0] * 7 + [1.0] + [0.0] * 7 + [0.0] * 8 + [1.0] + [0.0] * 7 + rows = [order, fused, halves, *small] + expected = [_explicit_sum_of_squares(row) for row in rows] + assert expected[:3] == [2.0**24 + 2, 1 + 2.0**-11 + 2.0**-22, 2.0**24] + assert _f32(sum(x * x for x in order)) == 2.0**24 + 10 # the exact sum, for contrast + + x = torch.tensor(rows)[None, :, None, :] # [B=1, S, H=1, D] + assert _sum_of_squares(x)[0, :, 0].tolist() == expected + # Through the op: context rows are identity rows, so with FP32 inputs the output is + # (x * rsqrt(ss / D + eps)) * w with exactly this ss, multiplied in that order. + video = torch.randn(1, 2, 1, d, generator=g) + table = joint.pack_rope_table(*_tables(2, d), len(rows), torch.float32) + wq1, wk1 = (torch.rand(d, generator=g) * 3 + 0.25 for _ in range(2)) + weights = [torch.ones(d), torch.ones(d), wq1, wk1] + q, k, _ = joint.fused_joint_qk_norm_rope(video, video, video, x, x, x, *weights, table, 1e-6) + inv = torch.rsqrt(torch.tensor(expected) / d + 1e-6) + scaled = x * inv[None, :, None, None] + assert torch.equal(q[:, 2:], scaled * wq1) and torch.equal(k[:, 2:], scaled * wk1) + assert not torch.equal(q[:, 2:], x * (inv[None, :, None, None] * wq1)) # order is pinned + + +def test_fma_is_emulated_with_one_rounding(): + t = lambda *v: torch.tensor(v, dtype=torch.float32) # noqa: E731 + a, c = t(1 + 2.0**-12), t(-(1 + 2.0**-11)) + assert _fma_fp32(a, a, c).item() == 2.0**-24 + assert (a * a + c).item() == 0.0 # two roundings lose it + # a * b = 1 + 2^-33 sits just above the FP32 midpoint 2^24 + 1; plain FP64 rounds onto + # the midpoint and then to even (2^24), the single rounding goes up. + a, b, c = t(1 + 2.0**-11), t(1 - 2.0**-11 + 2.0**-22), t(2.0**24) + assert (a.double() * b.double() + c.double()).float().item() == 2.0**24 + assert _fma_fp32(a, b, c).item() == 2.0**24 + 2 + special = _fma_fp32(t(float("inf"), float("nan")), t(1.0, 1.0), t(0.0, 0.0)) + assert special[0].item() == float("inf") and special[1].isnan() + + +def _rotate(n, cos, sin): + a, b = n[..., 0::2], n[..., 1::2] + return torch.stack((a * cos - b * sin, b * cos + a * sin), dim=-1).flatten(-2) + + +def test_one_bf16_rounding_before_the_rotation(): + g = torch.Generator().manual_seed(29) + s0, s1, h, d = 4, 2, 2, 8 + + def signs(s): # +-1: ss = D exactly, so with eps = 0 the inverse RMS is exactly 1 + return (torch.randint(0, 2, (1, s, h, d), generator=g) * 2 - 1).bfloat16() + + q0, k0, v0, q1, k1, v1 = (signs(s) for s in (s0, s0, s0, s1, s1, s1)) + w = [torch.rand(d, generator=g) + 0.5 for _ in range(4)] # FP32: n has bits BF16 drops + table = joint.pack_rope_table(*_tables(s0, d), s1, torch.float32) + q, k, _ = joint.fused_joint_qk_norm_rope(q0, k0, v0, q1, k1, v1, *w, table, 0.0) + cos, sin = table[None, :, None, : d // 2], table[None, :, None, d // 2 :] + for out, x0, x1, w0, w1 in ((q, q0, q1, w[0], w[2]), (k, k0, k1, w[1], w[3])): + n = torch.cat((x0.float() * w0, x1.float() * w1), dim=1) + expected = _rotate(n.bfloat16().float(), cos, sin).bfloat16() + assert torch.equal(out, expected) + assert not torch.equal(out, _rotate(n, cos, sin).bfloat16()) # rounding is observable + + +def test_bf16_rounding_is_nearest_even(): + # x = 1 and eps = 0 make the inverse RMS exactly 1, so n is the weight rounded once to + # BF16 (ulp 2^-7 on [1, 2)); context rows are identity rows, so the output is n. + d = 8 + ones = torch.ones(1, 1, 1, d, dtype=torch.bfloat16) + w = torch.tensor([1 + 2**-8, 1 + 3 * 2**-8, -(1 + 3 * 2**-8), 1 + 2**-8 + 2**-20, 1, 1, 1, 1]) + table = joint.pack_rope_table(*_tables(1, d), 1, torch.float32) + q, _, _ = joint.fused_joint_qk_norm_rope( + ones, ones, ones, ones, ones, ones, *[w] * 4, table, 0.0 + ) + assert q[0, 1, 0].float().tolist() == [1.0, 1 + 2**-6, -(1 + 2**-6), 1 + 2**-7, 1, 1, 1, 1] + + +def _packed_streams(batch, s, h, d, generator, device="cpu"): + """BF16 q, k, v views split from one packed [B, S, 3*H*D] projection, as on CUDA.""" + packed = (torch.randn(batch, s, 3 * h * d, generator=generator) * 2).to(device, torch.bfloat16) + return [t.unflatten(-1, (h, d)) for t in packed.split(h * d, dim=-1)] + + +def test_strided_batched_inputs_identity_rows_and_v_gather(): + g = torch.Generator().manual_seed(31) + batch, s0, s1, h, d, eps = 2, 5, 3, 3, 16, 1e-6 + streams = (*_packed_streams(batch, s0, h, d, g), *_packed_streams(batch, s1, h, d, g)) + q0, k0, v0, q1, k1, v1 = streams + assert not q0.is_contiguous() and q0.stride() == k0.stride() == v0.stride() + w = [(torch.rand(d, generator=g) + 0.5).bfloat16() for _ in range(4)] + tables = [ + joint.pack_rope_table(*_tables(s0, d, seed=b), s1, torch.bfloat16) for b in range(batch) + ] + q, k, v = outputs = joint.fused_joint_qk_norm_rope(*streams, *w, torch.cat(tables), eps) + for out in outputs: + assert out.shape == (batch, s0 + s1, h, d) and out.dtype == torch.bfloat16 + assert out.is_contiguous() + # Reading through the strides equals reading contiguous copies. + dense = joint.fused_joint_qk_norm_rope( + *(t.contiguous() for t in streams), *w, torch.cat(tables), eps + ) + assert all(torch.equal(a, b) for a, b in zip(outputs, dense)) + assert torch.equal(v, torch.cat((v0, v1), dim=1)) # V gathered bit for bit + # Identity context rows leave the normalised values unrotated; video rows are rotated. + assert torch.equal(q[:, s0:], _normalize(q1, w[2], eps).bfloat16()) + assert torch.equal(k[:, s0:], _normalize(k1, w[3], eps).bfloat16()) + assert not torch.equal(q[:, :s0], _normalize(q0, w[0], eps).bfloat16()) + # Each sample reads its own table rows. + for b in range(batch): + one = joint.fused_joint_qk_norm_rope(*(t[b : b + 1] for t in streams), *w, tables[b], eps) + assert all(torch.equal(full[b : b + 1], part) for full, part in zip(outputs, one)) + + +def _ulps(x): + """BF16 bit patterns as integers that count ulps through zero.""" + bits = x.view(torch.int16).int() + return torch.where(bits < 0, -(bits & 0x7FFF), bits) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the Triton kernel needs CUDA") +@pytest.mark.parametrize("batch, s0, s1, heads", [(1, 37, 11, 16), (2, 20, 5, 12)]) +def test_cuda_kernel_matches_the_reference(batch, s0, s1, heads): + pytest.importorskip("triton") + g = torch.Generator().manual_seed(41) + d = 128 + streams = ( + *_packed_streams(batch, s0, heads, d, g, "cuda"), + *_packed_streams(batch, s1, heads, d, g, "cuda"), + ) + w = [(torch.rand(d, generator=g) + 0.5).to("cuda", torch.bfloat16) for _ in range(4)] + tables = [] + for b in range(batch): + cos, sin = (t.cuda() for t in _tables(s0, d, seed=b)) + tables.append(joint.pack_rope_table(cos, sin, s1, torch.bfloat16)) + args = (*streams, *w, torch.cat(tables), 1e-6) + kernel = joint.fused_joint_qk_norm_rope(*args) + reference = _reference(*args) + for out in kernel: + assert out.shape == (batch, s0 + s1, heads, d) and out.dtype == torch.bfloat16 + assert out.is_contiguous() + assert torch.equal(kernel[2].view(torch.int16), reference[2].view(torch.int16)) + for got, want in zip(kernel[:2], reference[:2]): + # Identity (context) rows are the normalised values: bitwise. + assert torch.equal(got[:, s0:].view(torch.int16), want[:, s0:].view(torch.int16)) + assert (_ulps(got) - _ulps(want)).abs().max().item() <= 1 + + +def test_joint_path_needs_cuda_triton_bf16_head_128_and_cudnn(monkeypatch): + from types import SimpleNamespace as NS + + cuda = torch.device("cuda", 0) + bf16 = NS(is_cuda=True, dtype=torch.bfloat16, device=cuda) + monkeypatch.setattr(modeling, "_cudnn_attention_available", lambda device: True) + monkeypatch.setattr(modeling, "HAS_TRITON", True) + assert modeling.joint_attention_supported(bf16, 128) + cpu = NS(is_cuda=False, dtype=torch.bfloat16, device=torch.device("cpu")) + assert not modeling.joint_attention_supported(cpu, 128) + fp32 = NS(is_cuda=True, dtype=torch.float32, device=cuda) + assert not modeling.joint_attention_supported(fp32, 128) + assert not modeling.joint_attention_supported(bf16, 64) + monkeypatch.setattr(modeling, "HAS_TRITON", False) + assert not modeling.joint_attention_supported(bf16, 128) + monkeypatch.setattr(modeling, "HAS_TRITON", True) + monkeypatch.setattr(modeling, "_cudnn_attention_available", lambda device: False) + assert not modeling.joint_attention_supported(bf16, 128) diff --git a/tests/unit/backends/test_cuda_hunyuan15_transformer.py b/tests/unit/backends/test_cuda_hunyuan15_transformer.py new file mode 100644 index 00000000..f09feea9 --- /dev/null +++ b/tests/unit/backends/test_cuda_hunyuan15_transformer.py @@ -0,0 +1,393 @@ +"""HunyuanVideo 1.5 reference oracles; all tests run without CUDA hardware.""" + +import importlib +import json +import os +import subprocess +import sys + +import pytest +import torch +from diffusers import HunyuanVideo15Transformer3DModel as Reference + + +@pytest.fixture(autouse=True) +def _tensor_parallel_mesh(monkeypatch): + # Direct model construction binds TP=1 projections; do not depend on an + # earlier test module having initialized the process mesh. With no backend set, + # bind difflet.ops to the CPU ops (a Trainium-defaulting host would bind Neuron's). + monkeypatch.setenv("DIFFLET_BACKEND", os.environ.get("DIFFLET_BACKEND") or "cpu") + from difflet.ops import init_parallel_mesh + from difflet.pipeline.parallel_mesh import MeshSpec + + init_parallel_mesh(MeshSpec(tp=1)) + + +TINY = dict( + in_channels=9, + out_channels=4, + num_attention_heads=2, + attention_head_dim=8, + num_layers=2, + num_refiner_layers=1, + mlp_ratio=2.0, + patch_size=1, + patch_size_t=1, + text_embed_dim=12, + text_embed_2_dim=8, + image_embed_dim=6, + rope_axes_dim=(2, 2, 4), + task_type="t2v", + use_meanflow=False, +) + + +def _config(**overrides): + cls = importlib.import_module( + "difflet.backends.cuda.hunyuan_video_15.dit_config" + ).HunyuanVideo15DiTConfig + return cls( + **{ + **TINY, + "height": 32, + "width": 32, + "num_frames": 9, + "text_seq_len": 5, + "text_2_seq_len": 4, + "image_tokens": 3, + "torch_dtype": torch.float32, + **overrides, + } + ) + + +def _inputs(dtype=torch.float32, *, text_valid=3, glyph_valid=2, batch=1): + g = torch.Generator().manual_seed(7) + return dict( + hidden_states=torch.randn((batch, 9, 3, 2, 2), generator=g, dtype=dtype), + timestep=torch.full((batch,), 500.0, dtype=dtype), + encoder_hidden_states=torch.randn((batch, 5, 12), generator=g, dtype=dtype), + encoder_attention_mask=(torch.arange(5)[None].expand(batch, -1) < text_valid).clone(), + encoder_hidden_states_2=torch.randn((batch, 4, 8), generator=g, dtype=dtype), + encoder_attention_mask_2=(torch.arange(4)[None].expand(batch, -1) < glyph_valid).clone(), + image_embeds=torch.zeros((batch, 3, 6), dtype=dtype), + ) + + +def _pair(dtype=torch.float32): + model_cls = importlib.import_module( + "difflet.backends.cuda.hunyuan_video_15.modeling" + ).HunyuanVideo15Transformer3DModel + torch.manual_seed(42) + reference = Reference(**TINY).to(dtype).eval() + candidate = model_cls(_config(torch_dtype=dtype)).to(dtype).eval() + assert candidate.state_dict().keys() == reference.state_dict().keys() + candidate.load_state_dict(reference.state_dict(), strict=True) + return reference, candidate + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("valid", [(5, 4), (2, 0), (3, 2)]) +def test_tiny_hv15_matches_diffusers(dtype, valid, monkeypatch): + reference, candidate = _pair(dtype) + if dtype == torch.bfloat16: + # The CPU reference op rounds scores/probabilities to BF16; SDPA keeps + # intermediates in FP32. Match attention math, without widening bounds. + from difflet.backends.cuda.ops_impl.attention import attention + + modeling = importlib.import_module("difflet.backends.cuda.hunyuan_video_15.modeling") + monkeypatch.setattr(modeling, "attention", attention) + inputs = _inputs(dtype, text_valid=valid[0], glyph_valid=valid[1]) + with torch.no_grad(): + expected = reference(**inputs, return_dict=False)[0] + actual = candidate(**inputs, return_dict=False)[0] + assert actual.shape == (1, 4, 3, 2, 2) + assert torch.isfinite(actual).all() + # Same matmul/attention math and cast boundaries, including BF16. + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + else: + # Some x86 BF16 GEMM kernels round an element one ulp differently between the two + # models; a difference in attention math or cast boundaries moves elements further. + ulps = _bf16_ulps(actual, expected) + assert ulps.max() <= 1, f"{int((ulps > 1).sum())} elements off by up to {ulps.max()} ulps" + + +def _bf16_ulps(actual, expected): + """Elementwise distance in BF16 ulps: adjacent BF16 values, across zero too, are 1 apart.""" + + def ordered(x): + bits = x.view(torch.int16).to(torch.int32) + return torch.where(bits < 0, -(bits & 0x7FFF), bits) + + return (ordered(actual) - ordered(expected)).abs() + + +def test_hv15_masks_and_glyph_ordering(): + reference, candidate = _pair() + inputs = _inputs(batch=2) + inputs["encoder_attention_mask"][1] = torch.tensor([1, 0, 1, 0, 1]) + inputs["encoder_attention_mask_2"][1] = torch.tensor([0, 1, 0, 1]) + # Noncontiguous features and large values in padding must not enter valid attention. + for key, mask in [ + ("encoder_hidden_states", "encoder_attention_mask"), + ("encoder_hidden_states_2", "encoder_attention_mask_2"), + ]: + inputs[key] = inputs[key].transpose(1, 2).contiguous().transpose(1, 2) + inputs[key][~inputs[mask]] = 100.0 + with torch.no_grad(): + expected = reference(**inputs, return_dict=False)[0] + actual = candidate(**inputs, return_dict=False)[0] + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +def test_hv15_empty_glyphs_are_finite_and_blocks_get_prefix_lengths(): + _, candidate = _pair() + seen = [] + for block in candidate.transformer_blocks: + block.register_forward_pre_hook(lambda module, args: seen.append(args[3])) + with torch.no_grad(): + actual = candidate(**_inputs(glyph_valid=0), return_dict=False)[0] + assert torch.isfinite(actual).all() + assert len(seen) == 2 and seen[0] is seen[1] + # 12 video + 0 glyph + 3 text valid of 24 joint tokens; CPU int64, no S-sized mask. + assert seen[0].dtype == torch.int64 and seen[0].device.type == "cpu" + assert seen[0].tolist() == [15] + + +def test_hv15_invalid_config_and_geometry(tmp_path): + cls = type(_config()) + directory = tmp_path / "transformer" + directory.mkdir() + (directory / "config.json").write_text(json.dumps(TINY)) + cfg = cls.from_pretrained(directory, height=480, width=848, num_frames=121) + assert (cfg.latent_frames, cfg.latent_height, cfg.latent_width) == (31, 30, 53) + # The config owns the resident shape; checkpoint fields come from the checkpoint. + for overrides in [dict(height=481), dict(width=0), dict(num_frames=120)]: + with pytest.raises(ValueError): + _config(**overrides) + for overrides in [ + dict(task_type="i2v"), + dict(use_meanflow=True), + dict(patch_size=2), + dict(qk_norm="layer_norm"), + ]: + with pytest.raises(NotImplementedError): + _config(**overrides) + (directory / "config.json").write_text(json.dumps({**TINY, "_class_name": "OtherTransformer"})) + with pytest.raises(NotImplementedError, match="class"): + cls.from_pretrained(directory) + + +def test_hv15_cuda_ops_binding_runs_on_cpu_in_fresh_process(): + code = """ +import torch +from difflet.ops import init_parallel_mesh +from difflet.pipeline.parallel_mesh import MeshSpec +from difflet.backends.cuda.hunyuan_video_15.dit_config import HunyuanVideo15DiTConfig +from difflet.backends.cuda.hunyuan_video_15.modeling import HunyuanVideo15Transformer3DModel +init_parallel_mesh(MeshSpec(tp=1)) +cfg=HunyuanVideo15DiTConfig(in_channels=9,out_channels=4,num_attention_heads=2, +attention_head_dim=8,num_layers=1,num_refiner_layers=1,mlp_ratio=2,text_embed_dim=12, +text_embed_2_dim=8,image_embed_dim=6,rope_axes_dim=(2,2,4),height=32,width=32, +num_frames=9,text_seq_len=5,text_2_seq_len=4,image_tokens=3,torch_dtype=torch.float32) +m=HunyuanVideo15Transformer3DModel(cfg).eval() +assert "_torch_megatron" in type(m.proj_out).__module__ +with torch.no_grad(): + y=m(torch.randn(1,9,3,2,2),torch.tensor([500.]),torch.randn(1,5,12), + torch.tensor([[1,1,1,0,0]]),torch.randn(1,4,8),torch.zeros(1,4,dtype=torch.bool), + torch.zeros(1,3,6),return_dict=False)[0] +assert y.shape==(1,4,3,2,2) and torch.isfinite(y).all() +""" + result = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DIFFLET_BACKEND": "cuda"}, + text=True, + capture_output=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _component(tmp_path, *, sharded=False, torch_compile=False, dtype=torch.float32): + from difflet.backends.cuda.hunyuan_video_15.transformer import ( + CudaHunyuanVideo15TransformerApplication, + ) + + torch.manual_seed(42) + reference = Reference(**TINY).eval() + directory = tmp_path / "transformer" + reference.save_pretrained(directory, max_shard_size="10MB" if sharded else "10GB") + if sharded: + assert (directory / "diffusion_pytorch_model.safetensors.index.json").exists() + component = CudaHunyuanVideo15TransformerApplication( + model_path=str(directory), + config=_config(torch_dtype=dtype), + device="cpu", + torch_compile=torch_compile, + ) + artifact = str(tmp_path / "compiled") + component.compile(artifact) + assert component.has_compiled_artifacts(artifact) + component.load(artifact, skip_warmup=True) + return reference, component + + +@pytest.mark.parametrize("sharded", [False, True]) +def test_hv15_strict_component_matches_real_checkpoint(tmp_path, sharded): + reference, component = _component(tmp_path, sharded=sharded) + with torch.no_grad(): + expected = reference(**_inputs(), return_dict=False)[0] + actual = component(**_inputs(), return_dict=False)[0] + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + assert all(p.device.type == "cpu" and not p.is_meta for p in component.module.parameters()) + + +@pytest.mark.parametrize("error", ["missing", "unexpected"]) +def test_hv15_checkpoint_key_errors_are_fatal(tmp_path, error): + from difflet.backends.cuda.hunyuan_video_15.transformer import ( + CudaHunyuanVideo15TransformerApplication, + ) + from safetensors.torch import load_file, save_file + + _, component = _component(tmp_path) + file = tmp_path / "transformer/diffusion_pytorch_model.safetensors" + state = load_file(file) + if error == "missing": + state.pop("proj_out.bias") + else: + state["unexpected.weight"] = torch.zeros(1) + save_file(state, file) + fresh = CudaHunyuanVideo15TransformerApplication( + model_path=component.model_path, config=_config(), device="cpu" + ) + with pytest.raises(RuntimeError, match="Missing|Unexpected"): + fresh.load(str(tmp_path / "compiled"), skip_warmup=True) + + +def test_hv15_compile_scope_and_example_contract(tmp_path, monkeypatch): + seen = [] + + def compile_block(block, **kwargs): + seen.append((block, kwargs)) + return block + + monkeypatch.setattr(torch, "compile", compile_block) + _, component = _component(tmp_path, torch_compile=True, dtype=torch.bfloat16) + assert [row[0] for row in seen] == list(component.module.transformer_blocks) + inputs = component.get_example_inputs() + assert len(inputs) == 7 and inputs[0].shape == (1, 9, 3, 2, 2) + assert inputs[1].dtype == torch.bfloat16 + assert inputs[3].dtype == inputs[5].dtype == torch.bool + assert inputs[6].shape == (1, 3, 6) and torch.count_nonzero(inputs[6]) == 0 + + +def test_hv15_production_scope_contains_exactly_54_compiled_blocks(monkeypatch, tmp_path): + from difflet.backends.cuda.hunyuan_video_15.dit_config import HunyuanVideo15DiTConfig + from difflet.backends.cuda.hunyuan_video_15.modeling import ( + HunyuanVideo15Transformer3DModel, + ) + from difflet.backends.cuda.hunyuan_video_15.transformer import ( + CudaHunyuanVideo15TransformerApplication, + ) + from difflet.backends.gpu.application import compile_blocks + + cfg = HunyuanVideo15DiTConfig() + with torch.device("meta"): + model = HunyuanVideo15Transformer3DModel(cfg) + component = CudaHunyuanVideo15TransformerApplication( + model_path=str(tmp_path), config=cfg, device="cpu" + ) + seen = [] + monkeypatch.setattr( + torch, "compile", lambda block, **kwargs: seen.append((block, kwargs)) or block + ) + assert compile_blocks(component.repeated_blocks(model), dynamic=False) == 54 + assert [x[0] for x in seen] == list(model.transformer_blocks) + assert not any(x[0] in list(model.context_embedder.token_refiner.refiner_blocks) for x in seen) + assert sum(p.numel() for p in model.parameters()) == 8_326_608_160 + + +@pytest.mark.parametrize("valid", [(5, 4), (2, 0), (3, 2)]) +def test_tiny_hv15_unmasked_prefix_path_matches_diffusers_fp32(valid): + reference, candidate = _pair(torch.float32) + inputs = _inputs(torch.float32, text_valid=valid[0], glyph_valid=valid[1]) + with torch.no_grad(): + expected = reference(**inputs, return_dict=False)[0] + actual = candidate(**inputs, return_dict=False)[0] + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_hv15_blocks_attend_unmasked_valid_prefix(dtype, monkeypatch): + modeling = importlib.import_module("difflet.backends.cuda.hunyuan_video_15.modeling") + _, candidate = _pair(dtype) + calls = [] + original = modeling.attention + + def record(q, k, v, **kw): + calls.append((q.shape[2], kw.get("attention_mask"))) + return original(q, k, v, **kw) + + monkeypatch.setattr(modeling, "attention", record) + with torch.no_grad(): + candidate(**_inputs(dtype, text_valid=3, glyph_valid=2), return_dict=False) + dit = [c for c in calls if c[1] is None] + refiner = [c for c in calls if c[1] is not None] + # Two DiT blocks: 12 video + 2 glyph + 3 text = 17 of 24 joint tokens, unmasked. + assert [length for length, _ in dit] == [17, 17] + # The text refiner keeps its compact additive key bias over the 5 Qwen slots. + assert refiner and all(length == 5 for length, _ in refiner) + + +def test_hv15_prefix_attention_equals_masked_and_zeroes_padding(): + from difflet.backends.cuda.hunyuan_video_15.modeling import prefix_attention + + g = torch.Generator().manual_seed(3) + q, k, v = (torch.randn(2, 9, 2, 8, generator=g).transpose(1, 2) for _ in range(3)) + lengths = torch.tensor([5, 9]) + out = prefix_attention(q, k, v, lengths, 0.3) + valid = torch.arange(9)[None] < lengths[:, None] + bias = torch.zeros(2, 1, 1, 9).masked_fill(~valid[:, None, None], float("-inf")) + expected = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=bias, scale=0.3) + expected = expected.masked_fill(~valid[:, None, :, None], 0) + torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-6) + assert torch.count_nonzero(out[0, :, 5:]) == 0 + assert out.stride() == q.stride() # keeps the projection layout for the next transpose + + +def test_hv15_prefix_attention_is_one_opaque_op_under_compile(): + from difflet.backends.cuda.hunyuan_video_15.modeling import prefix_attention + + q, k, v = (torch.randn(1, 11, 2, 8).transpose(1, 2) for _ in range(3)) + lengths = torch.tensor([7]) + + def call(q, k, v, lengths): + return prefix_attention(q, k, v, lengths, 0.25) * 2 + + torch._dynamo.reset() + compiled = torch.compile(call, backend="eager", fullgraph=True, dynamic=False) + torch.testing.assert_close(compiled(q, k, v, lengths), call(q, k, v, lengths)) + # A different valid length reuses the same static-shape graph: no value guard. + with torch._dynamo.config.patch(error_on_recompile=True): + torch.testing.assert_close( + compiled(q, k, v, torch.tensor([4])), call(q, k, v, torch.tensor([4])) + ) + + +@pytest.mark.skipif( + torch.__version__ < (2, 10), reason="Inductor's `deterministic` option is new in torch 2.10" +) +def test_hv15_compiler_policy_stabilizes_reductions_without_global_changes(tmp_path): + # Inspect the real lazy Inductor wrappers: a CUDA-less host can verify that + # every block binds the policy, while CUDA repeatability remains a device gate. + from torch._inductor import config as inductor_config + + with inductor_config.patch(deterministic=False): + _, component = _component(tmp_path, torch_compile=True) + for block in component.module.transformer_blocks: + policy = block.get_compiler_config() + assert policy["deterministic"] is True + assert policy["max_autotune"] is True + assert policy["triton.cudagraphs"] is False + assert inductor_config.deterministic is False diff --git a/tests/unit/backends/test_gpu_application.py b/tests/unit/backends/test_gpu_application.py index 4386bf1c..5764203d 100644 --- a/tests/unit/backends/test_gpu_application.py +++ b/tests/unit/backends/test_gpu_application.py @@ -24,6 +24,20 @@ SHAPE = {"height": 64, "width": 64} # Tiny transformer/config.json per model: enough to build the components, never loaded. TINY_CONFIGS = { + "hunyuan_video_15": dict( + in_channels=9, + out_channels=4, + num_attention_heads=2, + attention_head_dim=8, + num_layers=1, + num_refiner_layers=1, + mlp_ratio=2.0, + text_embed_dim=12, + text_embed_2_dim=8, + image_embed_dim=6, + rope_axes_dim=[2, 2, 4], + task_type="t2v", + ), "hunyuan_video": dict( in_channels=2, out_channels=2, diff --git a/tests/unit/backends/test_gpu_shape_policy.py b/tests/unit/backends/test_gpu_shape_policy.py index 6a273055..12ba4360 100644 --- a/tests/unit/backends/test_gpu_shape_policy.py +++ b/tests/unit/backends/test_gpu_shape_policy.py @@ -13,10 +13,11 @@ def test_original_hunyuan_is_a_gpu_spec_but_15_remains_separate(): assert gpu_model_spec("hunyuan_video") is SPEC assert gpu_spec_for_cli_model_id(SPEC.cli_model_ids[0]) is SPEC assert "cuda" in resolve_model(SPEC.cli_model_ids[0]).backends - assert gpu_model_spec("hunyuan_video_15") is None + assert gpu_model_spec("hunyuan_video_15").model_type == "hunyuan_video_15" + assert gpu_model_spec("hunyuan_video_15") is not SPEC assert ( "cuda" - not in resolve_model("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v").backends + in resolve_model("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v").backends ) diff --git a/tests/unit/cli/test_cli_main_cuda.py b/tests/unit/cli/test_cli_main_cuda.py index ed460fd1..33aee0b9 100644 --- a/tests/unit/cli/test_cli_main_cuda.py +++ b/tests/unit/cli/test_cli_main_cuda.py @@ -24,6 +24,7 @@ def _cli(): ("black-forest-labs/FLUX.1-schnell", "flux"), ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "wan"), ("Wan-AI/Wan2.1-T2V-14B-Diffusers", "wan"), + ("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", "hunyuan_video_15"), ], ) def test_cuda_backend_routes_to_the_gpu_orchestrator(monkeypatch, model_id, model_type): @@ -44,13 +45,13 @@ def test_other_models_on_cuda_point_at_the_trainium_path(monkeypatch): with pytest.raises(SystemExit) as exc: _cli()._get_orchestrator( argparse.Namespace( - model_id="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + model_id="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v" ) ) message = str(exc.value) assert ( "DIFFLET_BACKEND=trainium" in message - and "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" in message + and "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v" in message ) # Wan runs on the GPU flow: the fallback does not name it. assert "Wan" not in message diff --git a/tests/unit/models/hunyuan_video_15/__init__.py b/tests/unit/models/hunyuan_video_15/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/models/hunyuan_video_15/test_cuda_application.py b/tests/unit/models/hunyuan_video_15/test_cuda_application.py new file mode 100644 index 00000000..cffab1ec --- /dev/null +++ b/tests/unit/models/hunyuan_video_15/test_cuda_application.py @@ -0,0 +1,344 @@ +"""HunyuanVideo 1.5 GPU spec, public factory, and the application on the real Diffusers driver.""" + +import os +import subprocess +import sys +from types import MethodType, SimpleNamespace + +import pytest +import torch +from diffusers import ( + AutoencoderKLHunyuanVideo15, + ClassifierFreeGuidance, + FlowMatchEulerDiscreteScheduler, + HunyuanVideo15Pipeline, +) +from diffusers.models.modeling_utils import ModelMixin + +from tests.unit.backends.test_cuda_hunyuan15_transformer import TINY, _component, _inputs + +MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + + +def _application(tmp_path): + from difflet.backends.cuda.hunyuan_video_15.application import CudaHunyuanVideo15Application + from difflet.pipeline.parallel_config import DiffletParallelConfig + + return CudaHunyuanVideo15Application( + model_path=str(tmp_path), + parallel=DiffletParallelConfig(tp_degree=1), + dtype="fp32", + shape=dict(height=32, width=32, num_frames=9), + device="cpu", + ) + + +class TinyVAE(ModelMixin): + temporal_compression_ratio = 4 + spatial_compression_ratio = 16 + + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(1)) + self.config = SimpleNamespace(latent_channels=4) + + +def _pipeline(transformer): + pipe = HunyuanVideo15Pipeline( + transformer=transformer, + vae=TinyVAE(), + scheduler=FlowMatchEulerDiscreteScheduler(shift=5), + text_encoder=None, + tokenizer=None, + text_encoder_2=None, + tokenizer_2=None, + guider=ClassifierFreeGuidance(guidance_scale=6.0), + ) + pipe.vision_num_semantic_tokens = 3 + + def encode(self, prompt=None, **kwargs): + inputs = _inputs() + # Deterministic conditioning distinguishes prompt, glyph and negative branches. + offset = sum(ord(c) for c in (prompt or "")) % 17 / 17 + glyph = '"' in (prompt or "") + return ( + (inputs["encoder_hidden_states"] + offset).to(kwargs["dtype"]), + inputs["encoder_attention_mask"], + (inputs["encoder_hidden_states_2"] - offset).to(kwargs["dtype"]), + ( + inputs["encoder_attention_mask_2"] + if glyph + else torch.zeros_like(inputs["encoder_attention_mask_2"]) + ), + ) + + pipe.encode_prompt = MethodType(encode, pipe) + pipe.set_progress_bar_config(disable=True) + return pipe + + +def _assert_driver_close(actual, expected, dtype, *, bf16_fraction): + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + else: + bound = bf16_fraction * expected.abs().max().item() + torch.testing.assert_close(actual, expected, atol=bound, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_hv15_real_driver_cfg_aba_and_negative_prompts(tmp_path, monkeypatch, dtype): + from difflet.backends.cuda.hunyuan_video_15.application import HunyuanVideo15DiTAdapter + + reference, component = _component(tmp_path, dtype=dtype) + reference.to(dtype) + if dtype == torch.bfloat16: + import difflet.backends.cuda.hunyuan_video_15.modeling as modeling + from difflet.backends.cuda.ops_impl.attention import attention + + monkeypatch.setattr(modeling, "attention", attention) + app = _application(tmp_path) + app.config = component.dit_config + app.transformer = component + app.adapter = HunyuanVideo15DiTAdapter(component) + app.pipe = _pipeline(app.adapter) + app.load_text_encoders = True + # Encoding is substituted only at the text-encoder boundary above. + app.pipe.tokenizer = object() + app._guider_config = dict(app.pipe.guider.config) + noise = torch.randn(1, 4, 3, 2, 2, generator=torch.Generator().manual_seed(8)) + outputs = [] + for prompt, negative, scale in [ + ("a cat", "", 6.0), + ('a sign "你好"', "blur", 1.0), + ("a cat", "", 6.0), + ("a cat", "motion", 6.0), + ]: + stock = _pipeline(reference) + stock.guider = ClassifierFreeGuidance(guidance_scale=scale) + # Capture each actual denoiser input/output, hence every scheduler transition. + stock_calls = [] + candidate_calls = [] + ref_hook = reference.register_forward_hook( + lambda m, a, k, y: stock_calls.append(y[0].clone()), with_kwargs=True + ) + candidate_hook = component.module.register_forward_hook( + lambda m, a, k, y: candidate_calls.append(y[0].clone()), with_kwargs=True + ) + try: + with torch.no_grad(): + expected = stock( + prompt=prompt, + negative_prompt=negative, + latents=noise.clone(), + height=32, + width=32, + num_frames=9, + num_inference_steps=3, + output_type="latent", + ).frames + actual = app.generate( + prompt, + negative_prompt=negative, + guidance_scale=scale, + latents=noise.clone(), + num_inference_steps=3, + output_type="latent", + ) + finally: + ref_hook.remove() + candidate_hook.remove() + assert len(stock_calls) == len(candidate_calls) == (3 if scale == 1 else 6) + # FP32 checks the driver, its CFG scale and negative prompts included, to FP32 + # rounding. In BF16 the valid-prefix attention rounds differently from stock's masked + # attention, by an amount that depends on the CPU's BF16 kernels and scales with the + # largest values rather than each element; on this tiny model that can exceed what a + # wrong scale or prompt changes, so BF16 checks only that the path stays within BF16 + # resolution of stock: each denoiser call to 2^-6 of its peak magnitude, and the + # final latents, where CFG at scale s can multiply a per-call difference by up to + # 2s - 1, to 2^-4. + for a, b in zip(candidate_calls, stock_calls): + _assert_driver_close(a, b, dtype, bf16_fraction=2**-6) + _assert_driver_close(actual, expected, dtype, bf16_fraction=2**-4) + outputs.append(actual) + torch.testing.assert_close(outputs[0], outputs[2], atol=0, rtol=0) + assert not torch.equal(outputs[0], outputs[3]) + + +class Pipe: + tokenizer = object() + + def __init__(self): + self.guider = ClassifierFreeGuidance(guidance_scale=6) + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append((kwargs, self.guider)) + return SimpleNamespace(frames=torch.full((1, 9, 3, 32, 32), 0.5)) + + +def test_hv15_generation_contract_and_guider_reset(tmp_path): + app = _application(tmp_path) + app.config = SimpleNamespace(out_channels=4, latent_frames=3, latent_height=2, latent_width=2) + app.pipe = Pipe() + app._guider_config = dict(app.pipe.guider.config) + output = app.generate('a sign "你好"', negative_prompt="blur", seed=7, num_inference_steps=3) + assert output.shape == (1, 3, 9, 32, 32) and output.dtype == torch.float32 + assert torch.count_nonzero(output) == 0 + first = app.pipe.calls[0][1] + app.generate("a cat", guidance_scale=1) + assert app.pipe.calls[1][1] is not first + assert app.pipe.calls[1][1].config.guidance_scale == 1 + assert app.pipe.calls[0][0]["negative_prompt"] == "blur" + assert app.pipe.calls[0][0]["prompt"] == 'a sign "你好"' + + +@pytest.mark.parametrize( + "kwargs", + [ + dict(num_inference_steps=0), + dict(width=48), + dict(num_frames=8), + dict(output_type="pil"), + dict(latents=torch.zeros(1, 4, 3, 2, 3)), + ], +) +def test_hv15_generation_refusals_precede_driver(tmp_path, kwargs): + app = _application(tmp_path) + app.pipe = Pipe() + app.config = SimpleNamespace(out_channels=4, latent_frames=3, latent_height=2, latent_width=2) + with pytest.raises(ValueError): + app.generate("a cat", **kwargs) + with pytest.raises(ValueError, match="one prompt"): + app.generate(["a cat", "a dog"]) + assert not app.pipe.calls + + +def test_hv15_load_companions_preserves_reference_text_contract(tmp_path, monkeypatch): + from difflet.backends.cuda.hunyuan_video_15 import application + + _, component = _component(tmp_path) + app = _application(tmp_path) + app.config = component.dit_config + app.transformer = component + app.adapter = application.HunyuanVideo15DiTAdapter(component) + monkeypatch.setattr(application, "configure_conditioning", lambda pipe: None) + + def load(*args, **kwargs): + assert kwargs["transformer"] is None + # The real constructor, so the reference's text contract is what is checked. + return HunyuanVideo15Pipeline( + transformer=None, + vae=AutoencoderKLHunyuanVideo15( + latent_channels=4, block_out_channels=(8,) * 5, layers_per_block=1 + ), + text_encoder=torch.nn.Identity(), + tokenizer=None, + text_encoder_2=torch.nn.Identity(), + tokenizer_2=None, + scheduler=FlowMatchEulerDiscreteScheduler(shift=5), + guider=ClassifierFreeGuidance(guidance_scale=6), + ) + + monkeypatch.setattr(HunyuanVideo15Pipeline, "from_pretrained", load) + app.load_companions() + assert app.pipe.prompt_template_encode_start_idx == 108 + assert app.pipe.tokenizer_max_length == 1000 and app.pipe.tokenizer_2_max_length == 256 + assert app.pipe.vision_num_semantic_tokens == 729 + assert app.pipe.vae.use_tiling + + +def test_hv15_gpu_spec_preserves_legacy_backend_contract(): + from difflet.models.spec import gpu_model_spec, gpu_spec_for_cli_model_id + from difflet.registry import resolve_model + + spec = gpu_model_spec("hunyuan_video_15") + assert spec is not None and gpu_spec_for_cli_model_id(MODEL) is spec + assert (spec.guidance, spec.accepts_negative_prompt, spec.accepts_cfg_parallel) == ( + "cfg", + True, + False, + ) + assert ( + spec.default_height, + spec.default_width, + spec.default_num_frames, + spec.default_steps, + spec.default_guidance_scale, + spec.default_fps, + ) == (480, 848, 121, 50, 6, 24) + entry = resolve_model(MODEL) + assert entry.backends == ("trainium", "neuron", "cuda") + # The registry keeps the Trainium layout (#139: TP2 x CFG2, standard CFG); the GPU + # spec above does not change it. + assert entry.default_parallel.tp_degree == 2 and entry.default_parallel.cfg_parallel_enabled + assert not entry.require_capabilities().is_distilled + assert spec.served_checkpoint_ids == (MODEL,) and not spec.prompt_budgets + + +@pytest.mark.parametrize("load_in_factory", [True, False]) +def test_hv15_public_factory_loads_real_transformer(tmp_path, monkeypatch, load_in_factory): + from diffusers import HunyuanVideo15Transformer3DModel + from difflet.backends.cuda.hunyuan_video_15.application import CudaHunyuanVideo15Application + from difflet.pipeline import difflet_pipeline as factory + from difflet.pipeline.parallel_config import DiffletParallelConfig + from difflet.platforms.cuda import CudaPlatform + + ref = HunyuanVideo15Transformer3DModel(**TINY).eval() + ref.save_pretrained(tmp_path / "transformer") + monkeypatch.setattr(CudaPlatform, "prepare_runtime", lambda *a: None) + monkeypatch.setattr(factory, "resolve_model_path", lambda *a, **kw: str(tmp_path)) + monkeypatch.setattr(CudaHunyuanVideo15Application, "load_companions", lambda *a: None) + pipe = factory.DiffletPipeline.from_pretrained( + MODEL, + backend="cuda", + # The registry's default layout is Trainium's (TP2 x CFG2); GPU callers pass TP=1. + parallel=DiffletParallelConfig(tp_degree=1), + dtype=torch.float32, + height=32, + width=32, + num_frames=9, + load=load_in_factory, + skip_compile=not load_in_factory, + skip_warmup=True, + compile_cache_dir=str(tmp_path / "cache"), + application_kwargs={"device": "cpu"}, + ) + if not load_in_factory: + pipe.compile() + pipe.load(skip_warmup=True) + assert pipe.parallel.tp_degree == 1 and pipe.app.transformer.is_loaded + with torch.no_grad(): + torch.testing.assert_close( + pipe.app.transformer(**_inputs(), return_dict=False)[0], + ref(**_inputs(), return_dict=False)[0], + atol=1e-4, + rtol=1e-4, + ) + + +def test_hv15_cli_cfg_refusal_names_the_missing_feature(): + from difflet.models.spec import gpu_model_spec + from tests.unit.cli.test_gpu_cli import _orch + + with pytest.raises(SystemExit, match="not implemented") as refused: + _orch(gpu_model_spec("hunyuan_video_15"), cfg_parallel=True).refuse_unsupported() + assert "no conditional/unconditional pair" not in str(refused.value) + + +def test_hv15_spec_and_cli_import_without_torch(): + code = """ +import sys +from difflet.models.spec import gpu_model_spec +from difflet.backends.gpu.cli import GpuOrchestrator +spec=gpu_model_spec("hunyuan_video_15") +assert spec.guidance=="cfg" +assert "torch" not in sys.modules +""" + result = subprocess.run( + [sys.executable, "-c", code], + env={**os.environ, "DIFFLET_BACKEND": "cpu"}, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/unit/serving/test_gpu_serving.py b/tests/unit/serving/test_gpu_serving.py index a99c02f6..dd86e88a 100644 --- a/tests/unit/serving/test_gpu_serving.py +++ b/tests/unit/serving/test_gpu_serving.py @@ -196,6 +196,7 @@ def test_request_shape_must_equal_the_profile_shape(spec, source): assert exc.value.code == "profile_mismatch" +@_specs(lambda s: bool(s.prompt_budgets)) def test_prompt_over_the_text_encoder_budget_is_refused(spec, source): (budget,) = spec.prompt_budgets limit = budget.drop_prefix_tokens + budget.max_tokens @@ -206,6 +207,13 @@ def test_prompt_over_the_text_encoder_budget_is_refused(spec, source): assert exc.value.code == "prompt_too_long" +@_specs(lambda s: not s.prompt_budgets) +def test_a_model_without_prompt_budgets_leaves_long_prompts_to_its_driver(spec, source): + # HunyuanVideo 1.5's driver crops its own chat template and glyph text; the serving + # validator must not apply a raw-string limit of its own. + _validator(_prepare(spec, _profile(spec))).validate(_request(spec, prompt="cat " * 2000)) + + def test_preload_loads_the_tokenizers_from_the_pinned_snapshot(spec, source, monkeypatch): seen = [] import transformers @@ -607,3 +615,33 @@ def test_videos_sync_serves_a_clip_through_the_gpu_adapter( in refused.text ) assert len(adapter.pipe.app.calls) == n + + +def test_gpu_serving_defaults_cfg_parallel_off_under_a_trainium_cfg_layout(monkeypatch, tmp_path): + """HunyuanVideo 1.5's registry layout (TP2 x CFG2) is Trainium's. Like the registry tp, + GPU serving must not inherit it: CFG-parallel is opt-in, as on the GPU CLI.""" + from difflet.registry import resolve_model + from difflet.serving.model_registry import resolve_serving_model + from difflet.serving.options import ServeOptions + + model = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + assert resolve_model(model).default_parallel.cfg_parallel_enabled # the Trainium layout + monkeypatch.setenv("DIFFLET_BACKEND", "cuda") + resolved = resolve_serving_model(ServeOptions(model_id=model, cache_dir=str(tmp_path))) + assert resolved.profile.parallel.cfg_parallel_enabled is False + assert resolved.profile.parallel.tp_degree == 1 + explicit = ServeOptions(model_id=model, cache_dir=str(tmp_path), cfg_parallel=True) + with pytest.raises(DiffletServingError, match="cfg-parallel"): + resolve_serving_model(explicit) + + +def test_a_model_served_only_on_gpus_is_refused_cleanly_elsewhere(monkeypatch, tmp_path): + """HunyuanVideo 1.5 is served on GPUs only. Elsewhere `difflet serve` gets a serving + error it prints in one line, not a traceback.""" + from difflet.serving.model_registry import resolve_serving_model + from difflet.serving.options import ServeOptions + + monkeypatch.setenv("DIFFLET_BACKEND", "cpu") + model = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + with pytest.raises(DiffletServingError, match="not enabled for serving"): + resolve_serving_model(ServeOptions(model_id=model, cache_dir=str(tmp_path)))