diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py new file mode 100644 index 00000000000..7a1c54f7b31 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion] + + +def _selected_operators(): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import ( + resolve_h3_vae_operators, + ) + + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + device = current_omni_platform.get_torch_device() + operators = resolve_h3_vae_operators(device) + if operators is None: + pytest.skip("No H3 VAE operator implementation is registered for this target") + return device, operators + + +def _operator_set(*, supports=lambda _device: True): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet + + return H3VAEOperatorSet( + supports=supports, + qk_norm_rope=lambda *_args: None, + scaled_residual=lambda *_args: None, + ) + + +def _qk_reference(x, cos, sin): + normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype) + rotary = normalized[..., :48] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return torch.cat( + (rotary * cos + rotated * sin, normalized[..., 48:]), + dim=-1, + ) + + +def _vit_norm_input(_module, hidden_states): + return hidden_states.float() + + +def _same_dtype_norm_input(_module, hidden_states): + return hidden_states + + +def _failing_norm_input(_module, _hidden_states): + raise RuntimeError("unsupported remote normalization semantics") + + +@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)]) +def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence): + device, operators = _selected_operators() + + torch.manual_seed(17) + qkv = torch.randn( + batch, + sequence, + 32, + 192, + device=device, + dtype=torch.float16, + ) + q, k, _ = qkv.chunk(3, dim=-1) + cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16) + sin = torch.randn_like(cos) + + expected_q = _qk_reference(q, cos, sin) + expected_k = _qk_reference(k, cos, sin) + with torch.inference_mode(): + actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5) + + assert actual is not None + actual_q, actual_k = actual + assert torch.equal(actual_q, expected_q) + assert torch.equal(actual_k, expected_k) + + +def test_h3_vae_scaled_residual_is_bit_exact(): + device, operators = _selected_operators() + + torch.manual_seed(29) + residual = torch.randn(195, 2048, device=device, dtype=torch.float32) + branch = torch.randn(195, 2048, device=device, dtype=torch.float16) + scale = torch.randn(2048, device=device, dtype=torch.float32) + expected = residual + branch * scale + + with torch.inference_mode(): + actual = operators.scaled_residual(residual, branch, scale) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_h3_vae_exact_ops_reject_unsupported_inputs(): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import ( + try_qk_norm_rope_exact, + ) + from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact + + q = torch.randn(1, 2, 32, 64, dtype=torch.float32) + cos = torch.randn(1, 2, 1, 48, dtype=torch.float32) + residual = torch.randn(2, 2048) + branch = torch.randn(2, 2048, dtype=torch.float16) + scale = torch.randn(2048) + with torch.inference_mode(): + assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None + assert try_scaled_residual_exact(residual, branch, scale) is None + + +def _make_decoder(): + class Attention(nn.Module): + def __init__(self): + super().__init__() + self.to_qkv = nn.Linear(8, 24) + self.to_out = nn.Linear(8, 8) + self.norm_q = nn.RMSNorm(8, elementwise_affine=False) + self.norm_k = nn.RMSNorm(8, elementwise_affine=False) + self.spatial_parallel = False + self.dim_head = 8 + + def perform_attention(self, query, _key, _value, _pack_info): + return query + + def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None): + return hidden_states + + class FeedForward(nn.Module): + def __init__(self): + super().__init__() + self.w1 = nn.Linear(8, 32) + self.w2 = nn.Linear(16, 8) + self.use_gated = True + self.act_fn = nn.SiLU() + self._compile_forward_enabled = False + self._compile_forward_fatal = False + self._compiled_forward = None + + def forward(self, hidden_states): + hidden_states = self.w1(hidden_states) + gate, hidden_states = hidden_states.chunk(2, dim=-1) + return self.w2(self.act_fn(gate) * hidden_states) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.attn = Attention() + self.ff = FeedForward() + self.norm1 = nn.RMSNorm(8) + self.norm2 = nn.RMSNorm(8) + self.scale1 = nn.Parameter(torch.zeros(8)) + self.scale2 = nn.Parameter(torch.zeros(8)) + self.use_scale = True + + def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None): + return hidden_states + + decoder = nn.Module() + decoder.transformer_blocks = nn.ModuleList([Block(), Block()]) + decoder.proj_out = nn.Linear(8, 8) + return decoder + + +def test_video_vae_installs_exact_optimizations(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.decoder = torch.nn.Module() + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = FakeModel() + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + install = Mock(return_value=True) + monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install) + monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock()) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("cuda"), + load_device=torch.device("cpu"), + ) + + install.assert_called_once_with( + video_vae.model.decoder, + device=torch.device("cuda"), + ) + + +def test_h3_vae_install_precasts_only_block_linears(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + operators = _operator_set() + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators) + decoder = _make_decoder() + + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + + for block in decoder.transformer_blocks: + assert block.attn.to_qkv.weight.dtype == torch.float16 + assert block.attn.to_out.weight.dtype == torch.float16 + assert block.ff.w1.weight.dtype == torch.float16 + assert block.ff.w2.weight.dtype == torch.float16 + assert block.forward.__func__.__name__ == "_optimized_transformer_block" + assert block.attn.forward.__func__.__name__ == "_optimized_attention" + assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward" + assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope + assert block._omni_scaled_residual is operators.scaled_residual + assert decoder.proj_out.weight.dtype == torch.float32 + + # Repeated installation is idempotent. + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + + +def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + for block in decoder.transformer_blocks: + # The official checkpoint loads this JSON boolean-like field as int 0. + block.attn.spatial_parallel = 0 + + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False) + + +def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + device, operators = _selected_operators() + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators) + decoder = _make_decoder().to(device) + feed_forward = decoder.transformer_blocks[0].ff + reference_forward = type(feed_forward).forward + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=device, + ) + + hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32) + with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16): + expected = reference_forward(feed_forward, hidden_states) + actual = feed_forward(hidden_states) + + assert expected.dtype == torch.float16 + assert torch.equal(actual, expected) + + +def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None) + decoder = _make_decoder() + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed") + for block in decoder.transformer_blocks: + assert block.attn.to_qkv.weight.dtype == torch.float32 + assert block.ff.w1.weight.dtype == torch.float32 + assert block.forward.__func__.__name__ == "forward" + assert block.attn.forward.__func__.__name__ == "forward" + + +@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input]) +def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__ + monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input) + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed") + assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32 + + +@pytest.mark.parametrize( + "break_contract", + [ + lambda block: delattr(block.attn, "spatial_parallel"), + lambda block: setattr(block.attn, "spatial_parallel", 2), + lambda block: setattr(block.attn, "perform_attention", None), + lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))), + lambda block: setattr(block.ff, "_compile_forward_enabled", True), + ], +) +def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + break_contract(decoder.transformer_blocks[0]) + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32 + + +def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch + + first = _operator_set(supports=lambda _device: False) + added = _operator_set(supports=lambda device: device.type == "meta") + monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added)) + + assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added + assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None + + +def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch + + platform = Mock() + platform.is_cuda.return_value = True + platform.is_available.return_value = True + monkeypatch.setattr(dispatch, "HAS_TRITON", True) + monkeypatch.setattr(dispatch, "current_omni_platform", platform) + + for capability in (90, 100, 103): + platform.get_device_capability.return_value.to_int.return_value = capability + assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None + + for capability in (89, 101, 110): + platform.get_device_capability.return_value.to_int.return_value = capability + assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/README.md b/vllm_omni/diffusion/models/minimax_h3/ops/README.md new file mode 100644 index 00000000000..e882ee78e28 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/README.md @@ -0,0 +1,50 @@ +# MiniMax-H3 optimized operators + +This package contains eager optimizations whose contracts are specific to the +MiniMax-H3 model implementation. Keeping them beside the model is intentional, +not an assumption that their underlying operations can never be shared. + +## Why model-owned + +The optimized path must preserve the H3 VAE's exact operation order, dtype and +rounding behavior, tensor layout, execution mode, and remote-code contract. A +kernel that implements a mathematically similar expression may still change the +decoded video, and a kernel that is beneficial on one platform may regress on +another. The model package therefore owns the capability checks, per-operation +input guards, reference fallbacks, numerical evidence, and performance evidence. + +This placement is one point in the design space being discussed in the +[diffusion operator-boundary RFC](https://github.com/vllm-project/vllm-omni/issues/6305). +If that RFC establishes a suitable shared tier for any of these operations, +migration can be handled separately without weakening the current contract. + +## Current boundary + +The VAE installer validates the official model structure before making any +change. It then binds the selected operator set, reuses Omni's existing +`SiluAndMul`, and materializes only the decoder-block Linear weights in the FP16 +dtype already selected by decode autocast. Unsupported model contracts, tensor +inputs, execution modes, and devices retain the original implementation. + +Hardware dispatch is an explicit allowlist. SM90, SM100, and SM103 are enabled; +other capabilities fall back to the reference path. Bit-exact full-decode and +operator evidence has been collected on SM90, with independent full-decode and +stress validation on SM103. Enabling a target and claiming it as validated are +kept separate so the evidence remains clear. + +## Extending platform support + +Validation and improvements for other platforms are welcome. A contribution +should keep the model-facing installation path unchanged and add a complete +operator-set entry with conservative fallback behavior. Please include: + +- the GPU product, compute capability, driver, CUDA, PyTorch, and Triton versions; +- representative production shapes plus edge-shape and unsupported-input tests; +- complete decoded-output comparison against the reference path; +- direct numerical checks for every optimized operation; +- warmed latency measurements for both complete decode and individual operations; +- tests proving that unsupported targets and contracts still use the reference path. + +Bitwise equality is the current numerical contract. If a platform cannot meet +it, please discuss the proposed quality contract and validation criteria before +changing the default path. diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py new file mode 100644 index 00000000000..6164781288b --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 model-specific optimized operators.""" + +from .vae import install_h3_vae_optimizations + +__all__ = ["install_h3_vae_optimizations"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py new file mode 100644 index 00000000000..247e072c82b --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Install bit-exact eager optimizations for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +import torch.nn as nn + +from vllm_omni.diffusion.layers.activation import SiluAndMul + +from .dispatch import resolve_h3_vae_operators + + +def _is_boolean_flag(value: Any) -> bool: + """Accept bools and the JSON-derived 0/1 flags used by remote VAE code.""" + + return isinstance(value, bool) or type(value) is int and value in (0, 1) + + +def _uses_fp32_attention_norm(attention: nn.Module) -> bool: + """Probe the loaded remote module's Q/K normalization semantics.""" + + forward = getattr(type(attention), "forward", None) + namespace = getattr(forward, "__globals__", None) + norm_input = namespace.get("_vit_norm_input") if isinstance(namespace, dict) else None + if not callable(norm_input): + return False + + probe = torch.empty(1, dtype=torch.float16) + try: + normalized = norm_input(attention.norm_q, probe) + except Exception: + return False + return ( + isinstance(normalized, torch.Tensor) and normalized.shape == probe.shape and normalized.dtype == torch.float32 + ) + + +def _optimized_feed_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return type(self).forward(self, hidden_states) + + hidden_states = self.w1(hidden_states) + if hidden_states.is_cuda and hidden_states.dtype == torch.float16: + hidden_states = self._omni_silu_and_mul(hidden_states) + else: + gate, hidden_states = hidden_states.chunk(2, dim=-1) + hidden_states = self.act_fn(gate) * hidden_states + return self.w2(hidden_states) + + +def _optimized_attention( + self: nn.Module, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + pack_info: dict[str, Any] | None = None, +) -> torch.Tensor: + pack_info = {} if pack_info is None else pack_info + if torch.compiler.is_compiling() or self.spatial_parallel or rotary_pos_emb is None: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + batch_size, sequence, _ = hidden_states.shape + qkv = self.to_qkv(hidden_states).view( + batch_size, + sequence, + -1, + 3 * self.dim_head, + ) + query, key, value = qkv.chunk(3, dim=-1) + optimized_qk = self._omni_qk_norm_rope( + query, + key, + rotary_pos_emb, + float(self.norm_q.eps), + ) + if optimized_qk is None: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + query, key = optimized_qk + hidden_states = self.perform_attention(query, key, value, pack_info) + hidden_states = hidden_states.reshape(batch_size, sequence, -1) + return self.to_out(hidden_states) + + +def _optimized_transformer_block( + self: nn.Module, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + pack_info: dict[str, Any] | None = None, +) -> torch.Tensor: + pack_info = {} if pack_info is None else pack_info + if torch.compiler.is_compiling() or hidden_states.dtype != torch.float32: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + normalized = self.norm1(hidden_states.float()).to(hidden_states.dtype) + attention_output = self.attn(normalized, rotary_pos_emb, pack_info) + updated = self._omni_scaled_residual( + hidden_states, + attention_output, + self.scale1, + ) + hidden_states = hidden_states + attention_output * self.scale1 if updated is None else updated + + normalized = self.norm2(hidden_states.float()).to(hidden_states.dtype) + feed_forward_output = self.ff(normalized) + updated = self._omni_scaled_residual( + hidden_states, + feed_forward_output, + self.scale2, + ) + return hidden_states + feed_forward_output * self.scale2 if updated is None else updated + + +def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None: + blocks = getattr(decoder, "transformer_blocks", None) + if not isinstance(blocks, nn.ModuleList) or not blocks: + return None + + linears: list[nn.Linear] = [] + for block in blocks: + attention = getattr(block, "attn", None) + feed_forward = getattr(block, "ff", None) + to_qkv = getattr(attention, "to_qkv", None) + to_out = getattr(attention, "to_out", None) + w1 = getattr(feed_forward, "w1", None) + w2 = getattr(feed_forward, "w2", None) + candidates = (to_qkv, to_out, w1, w2) + dim_head = getattr(attention, "dim_head", None) + scale1 = getattr(block, "scale1", None) + scale2 = getattr(block, "scale2", None) + if ( + not all(isinstance(linear, nn.Linear) for linear in candidates) + or not _is_boolean_flag(getattr(attention, "spatial_parallel", None)) + or not isinstance(dim_head, int) + or dim_head <= 0 + or not callable(getattr(attention, "perform_attention", None)) + or not isinstance(getattr(attention, "norm_q", None), nn.RMSNorm) + or not isinstance(getattr(attention, "norm_k", None), nn.RMSNorm) + or attention.norm_q.weight is not None + or attention.norm_k.weight is not None + or attention.norm_q.eps != attention.norm_k.eps + or not _uses_fp32_attention_norm(attention) + or not isinstance(getattr(block, "norm1", None), nn.RMSNorm) + or not isinstance(getattr(block, "norm2", None), nn.RMSNorm) + or block.norm1.weight is None + or block.norm2.weight is None + or not getattr(block, "use_scale", False) + or not isinstance(scale1, torch.Tensor) + or not isinstance(scale2, torch.Tensor) + or not getattr(feed_forward, "use_gated", False) + or not isinstance(getattr(feed_forward, "act_fn", None), nn.SiLU) + or not hasattr(feed_forward, "_compile_forward_enabled") + or getattr(feed_forward, "_compile_forward_enabled") + or not hasattr(feed_forward, "_compile_forward_fatal") + or not hasattr(feed_forward, "_compiled_forward") + ): + return None + + hidden_size = to_qkv.in_features + if ( + to_qkv.out_features != 3 * to_out.in_features + or to_out.in_features % dim_head != 0 + or to_out.out_features != hidden_size + or w1.in_features != hidden_size + or w1.out_features != 2 * w2.in_features + or w2.out_features != hidden_size + or attention.norm_q.normalized_shape != (dim_head,) + or attention.norm_k.normalized_shape != (dim_head,) + or block.norm1.normalized_shape != (hidden_size,) + or block.norm2.normalized_shape != (hidden_size,) + or scale1.shape != (hidden_size,) + or scale2.shape != (hidden_size,) + ): + return None + linears.extend(candidates) + return tuple(linears) + + +def install_h3_vae_optimizations( + decoder: nn.Module, + *, + device: torch.device, +) -> bool: + """Install the operators selected for ``device`` once.""" + + if getattr(decoder, "_omni_h3_vae_optimizations_installed", False): + return True + operators = resolve_h3_vae_operators(device) + if operators is None: + return False + + linears = _decoder_block_linears(decoder) + if linears is None or any(linear.weight.dtype != torch.float32 for linear in linears): + return False + + # The H3 decode path always uses FP16 CUDA autocast. Persisting these + # rounded decoder-block weights avoids rebuilding the same casts per tile. + for linear in linears: + linear.to(dtype=torch.float16) + + for block in decoder.transformer_blocks: + block.ff._omni_silu_and_mul = SiluAndMul() + block.ff.forward = MethodType(_optimized_feed_forward, block.ff) + block.attn._omni_qk_norm_rope = operators.qk_norm_rope + block.attn.forward = MethodType(_optimized_attention, block.attn) + block._omni_scaled_residual = operators.scaled_residual + block.forward = MethodType(_optimized_transformer_block, block) + + decoder._omni_h3_vae_optimizations_installed = True + return True + + +__all__ = ["install_h3_vae_optimizations"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py new file mode 100644 index 00000000000..d413ef85008 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Hardware dispatch for MiniMax H3 VAE operators.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.platforms import current_omni_platform + +from .qk_norm_rope import try_qk_norm_rope_exact +from .scaled_residual import try_scaled_residual_exact + +QKNormRopeOp = Callable[ + [ + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, torch.Tensor], + float, + ], + tuple[torch.Tensor, torch.Tensor] | None, +] +ScaledResidualOp = Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor], + torch.Tensor | None, +] + + +@dataclass(frozen=True) +class H3VAEOperatorSet: + """One validated hardware implementation of the H3 VAE operations.""" + + supports: Callable[[torch.device], bool] + qk_norm_rope: QKNormRopeOp + scaled_residual: ScaledResidualOp + + +def _supports_cuda_capability(device: torch.device, expected: int) -> bool: + if ( + not HAS_TRITON + or device.type != "cuda" + or torch.version.hip is not None + or not current_omni_platform.is_cuda() + or not current_omni_platform.is_available() + ): + return False + device_index = 0 if device.index is None else int(device.index) + capability = current_omni_platform.get_device_capability(device_index) + return capability is not None and int(capability.to_int()) == expected + + +def _supports_cuda_sm90(device: torch.device) -> bool: + return _supports_cuda_capability(device, 90) + + +def _supports_cuda_sm100(device: torch.device) -> bool: + return _supports_cuda_capability(device, 100) + + +def _supports_cuda_sm103(device: torch.device) -> bool: + return _supports_cuda_capability(device, 103) + + +# Keep hardware selection flat: adding a backend means adding one operator set, +# without changing the installer or the model execution path. +H3_VAE_OPERATOR_TABLE: tuple[H3VAEOperatorSet, ...] = ( + H3VAEOperatorSet( + supports=_supports_cuda_sm90, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), + H3VAEOperatorSet( + supports=_supports_cuda_sm100, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), + H3VAEOperatorSet( + supports=_supports_cuda_sm103, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), +) + + +def resolve_h3_vae_operators(device: torch.device) -> H3VAEOperatorSet | None: + for operators in H3_VAE_OPERATOR_TABLE: + if operators.supports(device): + return operators + return None + + +__all__ = [ + "H3VAEOperatorSet", + "H3_VAE_OPERATOR_TABLE", + "resolve_h3_vae_operators", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py new file mode 100644 index 00000000000..40b7069f61c --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bit-exact fused Q/K RMSNorm and RoPE for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _qk_rms_norm_rope_exact_kernel( + x_ptr, + cos_ptr, + sin_ptr, + output_ptr, + x_stride_token, + x_stride_head, + x_stride_dim, + rope_stride_token, + rope_stride_dim, + output_stride_token, + output_stride_head, + output_stride_dim, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + rotary_dim: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + ): + token = tl.program_id(0) + head_group = tl.program_id(1) + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + dims = tl.arange(0, head_dim) + mask = heads[:, None] < num_heads + + offsets = token * x_stride_token + heads[:, None] * x_stride_head + dims[None, :] * x_stride_dim + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None]).to(tl.float16) + + rotary_half = rotary_dim // 2 + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + pair_offsets = token * x_stride_token + heads[:, None] * x_stride_head + pair_dims[None, :] * x_stride_dim + pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None]).to(tl.float16) + + rope_mask = dims < rotary_dim + rope_offsets = token * rope_stride_token + dims * rope_stride_dim + cos = tl.load(cos_ptr + rope_offsets, mask=rope_mask, other=0.0).to(tl.float16) + sin = tl.load(sin_ptr + rope_offsets, mask=rope_mask, other=0.0).to(tl.float16) + signed_pair = tl.where( + dims[None, :] < rotary_half, + -pair_normalized, + pair_normalized, + ) + # Preserve multiply-round, multiply-round, add-round. Ordinary Triton + # arithmetic may contract this expression into an FMA. + rotated = tl.inline_asm_elementwise( + """ + { + .reg .b32 first; + .reg .b32 second; + mul.rn.f16x2 first, $1, $2; + mul.rn.f16x2 second, $3, $4; + add.rn.f16x2 $0, first, second; + } + """, + constraints="=r,r,r,r,r", + args=[normalized, cos[None, :], signed_pair, sin[None, :]], + dtype=tl.float16, + is_pure=True, + pack=2, + ) + output = tl.where(rope_mask[None, :], rotated, normalized) + + output_offsets = ( + token * output_stride_token + heads[:, None] * output_stride_head + dims[None, :] * output_stride_dim + ) + tl.store(output_ptr + output_offsets, output, mask=mask) + + +def _supported_inputs( + q: torch.Tensor, + k: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor], +) -> bool: + cos, sin = rotary_pos_emb + return ( + HAS_TRITON + and torch.version.hip is None + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and q.is_cuda + and k.is_cuda + and q.device == k.device + and q.dtype == torch.float16 + and k.dtype == q.dtype + and q.ndim == 4 + and k.shape == q.shape + and q.numel() > 0 + and q.shape[-1] == 64 + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and cos.shape == sin.shape + and cos.ndim == 4 + and cos.shape[:2] == q.shape[:2] + and cos.shape[2] == 1 + and cos.shape[-1] == 48 + and cos.device == q.device + and sin.device == q.device + and cos.dtype == q.dtype + and sin.dtype == q.dtype + and cos.is_contiguous() + and sin.is_contiguous() + and cos.stride() == sin.stride() + ) + + +def try_qk_norm_rope_exact( + q: torch.Tensor, + k: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor], + eps: float, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return the fused bit-exact result, or ``None`` outside its contract.""" + + if not _supported_inputs(q, k, rotary_pos_emb): + return None + + output_shape = q.shape + batch, sequence, heads, head_dim = output_shape + tokens = batch * sequence + cos, sin = rotary_pos_emb + q = q.reshape(tokens, heads, head_dim) + k = k.reshape(tokens, heads, head_dim) + cos = cos.reshape(tokens, cos.shape[-1]) + sin = sin.reshape(tokens, sin.shape[-1]) + q_output = torch.empty_like(q) + k_output = torch.empty_like(k) + # This launch layout is part of the exactness contract on the validated + # SM90 and SM103 targets: changing the warp-to-row mapping changes the + # FP32 RMS reduction order. + heads_per_program = 8 + grid = (tokens, triton.cdiv(heads, heads_per_program)) + launch_args = { + "num_heads": heads, + "head_dim": head_dim, + "rotary_dim": cos.shape[-1], + "eps": eps, + "heads_per_program": heads_per_program, + "num_warps": 4, + } + _qk_rms_norm_rope_exact_kernel[grid]( + q, + cos, + sin, + q_output, + q.stride(0), + q.stride(1), + q.stride(2), + cos.stride(0), + cos.stride(1), + q_output.stride(0), + q_output.stride(1), + q_output.stride(2), + **launch_args, + ) + _qk_rms_norm_rope_exact_kernel[grid]( + k, + cos, + sin, + k_output, + k.stride(0), + k.stride(1), + k.stride(2), + cos.stride(0), + cos.stride(1), + k_output.stride(0), + k_output.stride(1), + k_output.stride(2), + **launch_args, + ) + return q_output.reshape(output_shape), k_output.reshape(output_shape) + + +__all__ = ["try_qk_norm_rope_exact"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py new file mode 100644 index 00000000000..290e10e5fbe --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bit-exact scaled residual update for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _mul_rn_f32(x, y): + return tl.inline_asm_elementwise( + "mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + @triton.jit + def _scaled_residual_exact_kernel( + output_ptr, + residual_ptr, + branch_ptr, + scale_ptr, + residual_stride_row, + branch_stride_row, + hidden_size: tl.constexpr, + block_n: tl.constexpr, + ): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + residual = tl.load( + residual_ptr + row * residual_stride_row + columns, + mask=mask, + other=0.0, + ).to(tl.float32) + branch = tl.load( + branch_ptr + row * branch_stride_row + columns, + mask=mask, + other=0.0, + ).to(tl.float32) + scale = tl.load(scale_ptr + columns, mask=mask, other=0.0).to(tl.float32) + # Prevent contraction into FMA and retain the eager rounding boundary. + updated = residual + _mul_rn_f32(branch, scale) + tl.store(output_ptr + row * hidden_size + columns, updated, mask=mask) + + +def try_scaled_residual_exact( + residual: torch.Tensor, + branch: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor | None: + """Return the optimized result, or ``None`` outside its exact contract.""" + + if not ( + HAS_TRITON + and torch.version.hip is None + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and residual.is_cuda + and branch.is_cuda + and residual.device == branch.device == scale.device + and residual.dtype == torch.float32 + and branch.dtype == torch.float16 + and scale.dtype == torch.float32 + and residual.shape == branch.shape + and residual.ndim >= 1 + and residual.shape[-1] == 2048 + and scale.shape == residual.shape[-1:] + and residual.numel() > 0 + and residual.is_contiguous() + and branch.is_contiguous() + and scale.is_contiguous() + ): + return None + + hidden_size = residual.shape[-1] + residual_2d = residual.reshape(-1, hidden_size) + branch_2d = branch.reshape(-1, hidden_size) + output = torch.empty_like(residual_2d) + rows = residual_2d.shape[0] + _scaled_residual_exact_kernel[(rows,)]( + output, + residual_2d, + branch_2d, + scale, + residual_2d.stride(0), + branch_2d.stride(0), + hidden_size=hidden_size, + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return output.reshape_as(residual) + + +__all__ = ["try_scaled_residual_exact"] diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index b52417f06b3..2e830fa0481 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -27,6 +27,7 @@ PinnedModuleStager, ) +from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 @@ -128,11 +129,18 @@ def __init__( component_path, self.config_dict, ) - # Match the reference loader contract: video VAE weights stay FP32. - # Keyframe encoding is numerically sensitive to first casting the - # checkpoint through FP16; decode still runs under FP16 autocast. + # Match the reference loader contract before installing inference-only + # decoder fast paths. Keyframe encoding remains FP32; decoder Linear + # weights may be materialized in FP16 because reference decode casts + # those same tensors through CUDA autocast on every tile. initial_device = load_device or device self.remote.eval().to(device=initial_device, dtype=torch.float32) + decoder = getattr(self.remote.model, "decoder", None) + if decoder is not None: + install_h3_vae_optimizations( + decoder, + device=device, + ) self._stager = None if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): self._stager = PinnedModuleStager(