From c5c149cd8246ca4e0f07ace150d35bfa15c6bcaf Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Wed, 26 Aug 2026 15:37:58 +0800
Subject: [PATCH 01/10] [Perf] Add exact MiniMax-H3 VAE decoder ops

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_contract.py    |  43 +++
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 246 ++++++++++++++++++
 .../models/minimax_h3/ops/__init__.py         |   7 +
 .../models/minimax_h3/ops/vae/__init__.py     | 182 +++++++++++++
 .../models/minimax_h3/ops/vae/qk_norm_rope.py | 191 ++++++++++++++
 .../minimax_h3/ops/vae/scaled_residual.py     | 101 +++++++
 vllm_omni/diffusion/models/minimax_h3/vae.py  |  14 +-
 7 files changed, 781 insertions(+), 3 deletions(-)
 create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/__init__.py
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py

diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 76adf7471bc..5a2c1233063 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -1417,6 +1417,49 @@ def __init__(self):
     assert next(video_vae.parameters()).dtype == torch.float32
 
 
+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_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 
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..f8590216c79
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+
+def _sm90_available() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)
+
+
+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,
+    )
+
+
+@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
+def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
+        try_qk_norm_rope_exact,
+    )
+
+    torch.manual_seed(17)
+    qkv = torch.randn(
+        batch,
+        sequence,
+        32,
+        192,
+        device="cuda",
+        dtype=torch.float16,
+    )
+    q, k, _ = qkv.chunk(3, dim=-1)
+    cos = torch.randn(batch, sequence, 1, 48, device="cuda", 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 = try_qk_norm_rope_exact(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)
+
+
+@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_h3_vae_scaled_residual_is_bit_exact():
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import (
+        try_scaled_residual_exact,
+    )
+
+    torch.manual_seed(29)
+    residual = torch.randn(195, 2048, device="cuda", dtype=torch.float32)
+    branch = torch.randn(195, 2048, device="cuda", dtype=torch.float16)
+    scale = torch.randn(2048, device="cuda", dtype=torch.float32)
+    expected = residual + branch * scale
+
+    with torch.inference_mode():
+        actual = try_scaled_residual_exact(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()
+
+        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_h3_vae_install_precasts_only_block_linears(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(vae_ops, "_supports_h3_vae_optimizations", lambda _device: True)
+    decoder = _make_decoder()
+
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("cuda:0"),
+    )
+
+    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 decoder.proj_out.weight.dtype == torch.float32
+
+    # Repeated installation is idempotent.
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("cuda:0"),
+    )
+
+
+@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
+def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setattr(vae_ops, "_supports_h3_vae_optimizations", lambda _device: True)
+    decoder = _make_decoder().cuda()
+    feed_forward = decoder.transformer_blocks[0].ff
+    reference_forward = type(feed_forward).forward
+    assert vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("cuda:0"),
+    )
+
+    hidden_states = torch.randn(4, 8, device="cuda", dtype=torch.float32)
+    with torch.inference_mode(), torch.autocast("cuda", 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, "_supports_h3_vae_optimizations", lambda _device: False)
+    decoder = _make_decoder()
+
+    assert not vae_ops.install_h3_vae_optimizations(
+        decoder,
+        device=torch.device("cuda:0"),
+    )
+    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"
+
+
+def test_h3_vae_support_is_platform_and_capability_scoped(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    class FakePlatform:
+        def __init__(self, cuda, capability):
+            self.cuda = cuda
+            self.capability = capability
+
+        def is_cuda(self):
+            return self.cuda
+
+        def is_available(self):
+            return True
+
+        def get_device_capability(self, _device_index):
+            capability = self.capability
+
+            class Capability:
+                def to_int(self):
+                    return capability
+
+            return Capability()
+
+    platform = FakePlatform(True, 90)
+    monkeypatch.setattr(vae_ops, "current_omni_platform", platform)
+    monkeypatch.setattr(vae_ops, "HAS_TRITON", True)
+
+    assert vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
+    platform.capability = 89
+    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
+    platform.cuda = False
+    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
+    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cpu"))
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..909c5b7878e
--- /dev/null
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -0,0 +1,182 @@
+# 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.triton_utils import HAS_TRITON
+
+from vllm_omni.diffusion.layers.activation import SiluAndMul
+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
+
+_H3_VAE_COMPUTE_CAPABILITY = 90
+
+
+def _supports_h3_vae_optimizations(device: torch.device) -> bool:
+    """Return whether the exact operators were validated for ``device``."""
+
+    if (
+        not HAS_TRITON
+        or device.type != "cuda"
+        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()) == _H3_VAE_COMPUTE_CAPABILITY
+
+
+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 self.use_gated:
+        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
+    else:
+        hidden_states = self.act_fn(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 = try_qk_norm_rope_exact(
+        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 = try_scaled_residual_exact(
+        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 = try_scaled_residual_exact(
+        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)
+        candidates = (
+            getattr(attention, "to_qkv", None),
+            getattr(attention, "to_out", None),
+            getattr(feed_forward, "w1", None),
+            getattr(feed_forward, "w2", None),
+        )
+        if (
+            not all(isinstance(linear, nn.Linear) for linear in candidates)
+            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 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 getattr(feed_forward, "use_gated", False)
+            or not isinstance(getattr(feed_forward, "act_fn", None), nn.SiLU)
+        ):
+            return None
+        linears.extend(candidates)
+    return tuple(linears)
+
+
+def install_h3_vae_optimizations(
+    decoder: nn.Module,
+    *,
+    device: torch.device,
+) -> bool:
+    """Install validated SM90 optimizations once, leaving other targets intact."""
+
+    if getattr(decoder, "_omni_h3_vae_optimizations_installed", False):
+        return True
+    if not _supports_h3_vae_optimizations(device):
+        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.forward = MethodType(_optimized_attention, block.attn)
+        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/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py
new file mode 100644
index 00000000000..95570acb34d
--- /dev/null
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py
@@ -0,0 +1,191 @@
+# 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 SM90 launch layout is part of the exactness contract: 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(

From 9e66411481abc8b1a92f44e48229bff44e7d2fce Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Thu, 27 Aug 2026 16:22:58 +0800
Subject: [PATCH 02/10] [Refactor] Add H3 VAE hardware dispatch

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 118 ++++++++----------
 .../models/minimax_h3/ops/vae/__init__.py     |  35 ++----
 .../models/minimax_h3/ops/vae/dispatch.py     |  78 ++++++++++++
 3 files changed, 142 insertions(+), 89 deletions(-)
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py

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
index f8590216c79..b29b4f280f0 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -5,13 +5,34 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from vllm.triton_utils import HAS_TRITON
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion]
 
-def _sm90_available() -> bool:
-    return torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)
+
+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):
@@ -25,13 +46,9 @@ def _qk_reference(x, cos, sin):
     )
 
 
-@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
 @pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
 def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
-    from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
-        try_qk_norm_rope_exact,
-    )
+    device, operators = _selected_operators()
 
     torch.manual_seed(17)
     qkv = torch.randn(
@@ -39,17 +56,17 @@ def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
         sequence,
         32,
         192,
-        device="cuda",
+        device=device,
         dtype=torch.float16,
     )
     q, k, _ = qkv.chunk(3, dim=-1)
-    cos = torch.randn(batch, sequence, 1, 48, device="cuda", dtype=torch.float16)
+    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 = try_qk_norm_rope_exact(q, k, (cos, sin), 1e-5)
+        actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5)
 
     assert actual is not None
     actual_q, actual_k = actual
@@ -57,21 +74,17 @@ def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
     assert torch.equal(actual_k, expected_k)
 
 
-@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
-@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
 def test_h3_vae_scaled_residual_is_bit_exact():
-    from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import (
-        try_scaled_residual_exact,
-    )
+    device, operators = _selected_operators()
 
     torch.manual_seed(29)
-    residual = torch.randn(195, 2048, device="cuda", dtype=torch.float32)
-    branch = torch.randn(195, 2048, device="cuda", dtype=torch.float16)
-    scale = torch.randn(2048, device="cuda", dtype=torch.float32)
+    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 = try_scaled_residual_exact(residual, branch, scale)
+        actual = operators.scaled_residual(residual, branch, scale)
 
     assert actual is not None
     assert torch.equal(actual, expected)
@@ -146,12 +159,13 @@ def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
 def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
 
-    monkeypatch.setattr(vae_ops, "_supports_h3_vae_optimizations", lambda _device: True)
+    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("cuda:0"),
+        device=torch.device("meta"),
     )
 
     for block in decoder.transformer_blocks:
@@ -162,30 +176,32 @@ def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
         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("cuda:0"),
+        device=torch.device("meta"),
     )
 
 
-@pytest.mark.skipif(not _sm90_available(), reason="CUDA SM90 required")
 def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
 
-    monkeypatch.setattr(vae_ops, "_supports_h3_vae_optimizations", lambda _device: True)
-    decoder = _make_decoder().cuda()
+    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=torch.device("cuda:0"),
+        device=device,
     )
 
-    hidden_states = torch.randn(4, 8, device="cuda", dtype=torch.float32)
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
+    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)
 
@@ -196,12 +212,12 @@ def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
 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, "_supports_h3_vae_optimizations", lambda _device: False)
+    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("cuda:0"),
+        device=torch.device("meta"),
     )
     assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
     for block in decoder.transformer_blocks:
@@ -211,36 +227,12 @@ def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
         assert block.attn.forward.__func__.__name__ == "forward"
 
 
-def test_h3_vae_support_is_platform_and_capability_scoped(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
-
-    class FakePlatform:
-        def __init__(self, cuda, capability):
-            self.cuda = cuda
-            self.capability = capability
-
-        def is_cuda(self):
-            return self.cuda
-
-        def is_available(self):
-            return True
-
-        def get_device_capability(self, _device_index):
-            capability = self.capability
-
-            class Capability:
-                def to_int(self):
-                    return capability
-
-            return Capability()
+def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
 
-    platform = FakePlatform(True, 90)
-    monkeypatch.setattr(vae_ops, "current_omni_platform", platform)
-    monkeypatch.setattr(vae_ops, "HAS_TRITON", True)
+    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 vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
-    platform.capability = 89
-    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
-    platform.cuda = False
-    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cuda:0"))
-    assert not vae_ops._supports_h3_vae_optimizations(torch.device("cpu"))
+    assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added
+    assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index 909c5b7878e..007779e4866 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -9,30 +9,10 @@
 
 import torch
 import torch.nn as nn
-from vllm.triton_utils import HAS_TRITON
 
 from vllm_omni.diffusion.layers.activation import SiluAndMul
-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
-
-_H3_VAE_COMPUTE_CAPABILITY = 90
-
-
-def _supports_h3_vae_optimizations(device: torch.device) -> bool:
-    """Return whether the exact operators were validated for ``device``."""
-
-    if (
-        not HAS_TRITON
-        or device.type != "cuda"
-        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()) == _H3_VAE_COMPUTE_CAPABILITY
+from .dispatch import resolve_h3_vae_operators
 
 
 def _optimized_feed_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -69,7 +49,7 @@ def _optimized_attention(
         3 * self.dim_head,
     )
     query, key, value = qkv.chunk(3, dim=-1)
-    optimized_qk = try_qk_norm_rope_exact(
+    optimized_qk = self._omni_qk_norm_rope(
         query,
         key,
         rotary_pos_emb,
@@ -96,7 +76,7 @@ def _optimized_transformer_block(
 
     normalized = self.norm1(hidden_states.float()).to(hidden_states.dtype)
     attention_output = self.attn(normalized, rotary_pos_emb, pack_info)
-    updated = try_scaled_residual_exact(
+    updated = self._omni_scaled_residual(
         hidden_states,
         attention_output,
         self.scale1,
@@ -105,7 +85,7 @@ def _optimized_transformer_block(
 
     normalized = self.norm2(hidden_states.float()).to(hidden_states.dtype)
     feed_forward_output = self.ff(normalized)
-    updated = try_scaled_residual_exact(
+    updated = self._omni_scaled_residual(
         hidden_states,
         feed_forward_output,
         self.scale2,
@@ -153,11 +133,12 @@ def install_h3_vae_optimizations(
     *,
     device: torch.device,
 ) -> bool:
-    """Install validated SM90 optimizations once, leaving other targets intact."""
+    """Install the operators selected for ``device`` once."""
 
     if getattr(decoder, "_omni_h3_vae_optimizations_installed", False):
         return True
-    if not _supports_h3_vae_optimizations(device):
+    operators = resolve_h3_vae_operators(device)
+    if operators is None:
         return False
 
     linears = _decoder_block_linears(decoder)
@@ -172,7 +153,9 @@ def install_h3_vae_optimizations(
     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
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..4d890951f7c
--- /dev/null
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
@@ -0,0 +1,78 @@
+# 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_sm90(device: torch.device) -> 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()) == 90
+
+
+# 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,
+    ),
+)
+
+
+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",
+]

From 85edb464c21481690f023f5365f61e98974ddc65 Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Thu, 27 Aug 2026 16:57:34 +0800
Subject: [PATCH 03/10] [Bugfix] Guard H3 VAE eager operator installation

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 57 +++++++++++++++
 .../models/minimax_h3/ops/vae/__init__.py     | 70 +++++++++++++++----
 2 files changed, 114 insertions(+), 13 deletions(-)

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
index b29b4f280f0..f0c49e0a20f 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -130,6 +130,9 @@ def __init__(self):
             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)
@@ -227,6 +230,60 @@ def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
         assert block.attn.forward.__func__.__name__ == "forward"
 
 
+@pytest.mark.parametrize(
+    ("name", "value"),
+    [
+        ("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "0"),
+        ("MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "1"),
+        ("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "true"),
+    ],
+)
+def test_h3_vae_install_preserves_nondefault_official_semantics(monkeypatch, name, value):
+    from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+    monkeypatch.setenv(name, value)
+    monkeypatch.setattr(
+        vae_ops,
+        "resolve_h3_vae_operators",
+        lambda _device: _operator_set(),
+    )
+    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")
+    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, "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
 
diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index 007779e4866..f705a0ee433 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import os
 from types import MethodType
 from typing import Any
 
@@ -15,19 +16,31 @@
 from .dispatch import resolve_h3_vae_operators
 
 
+def _env_flag(name: str, default: str) -> bool:
+    value = os.environ.get(name, default)
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _official_execution_semantics_supported() -> bool:
+    """Return whether replacing the official eager forwards is semantics-safe."""
+
+    return (
+        _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1")
+        and not _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "0")
+        and not _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0")
+    )
+
+
 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 self.use_gated:
-        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
+    if hidden_states.is_cuda and hidden_states.dtype == torch.float16:
+        hidden_states = self._omni_silu_and_mul(hidden_states)
     else:
-        hidden_states = self.act_fn(hidden_states)
+        gate, hidden_states = hidden_states.chunk(2, dim=-1)
+        hidden_states = self.act_fn(gate) * hidden_states
     return self.w2(hidden_states)
 
 
@@ -102,14 +115,20 @@ def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None:
     for block in blocks:
         attention = getattr(block, "attn", None)
         feed_forward = getattr(block, "ff", None)
-        candidates = (
-            getattr(attention, "to_qkv", None),
-            getattr(attention, "to_out", None),
-            getattr(feed_forward, "w1", None),
-            getattr(feed_forward, "w2", 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 isinstance(getattr(attention, "spatial_parallel", None), bool)
+            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
@@ -120,8 +139,31 @@ def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None:
             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)
@@ -137,6 +179,8 @@ def install_h3_vae_optimizations(
 
     if getattr(decoder, "_omni_h3_vae_optimizations_installed", False):
         return True
+    if not _official_execution_semantics_supported():
+        return False
     operators = resolve_h3_vae_operators(device)
     if operators is None:
         return False

From f4703e6ed8b22576696144182799c059e5818399 Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Thu, 27 Aug 2026 17:05:59 +0800
Subject: [PATCH 04/10] [Test] Keep H3 VAE ops coverage self-contained

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_contract.py    | 43 ------------------
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 45 +++++++++++++++++++
 2 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 5a2c1233063..76adf7471bc 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -1417,49 +1417,6 @@ def __init__(self):
     assert next(video_vae.parameters()).dtype == torch.float32
 
 
-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_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
 
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
index f0c49e0a20f..417dfda946e 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -1,6 +1,8 @@
 # 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
@@ -159,6 +161,49 @@ def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
     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
 

From f6775966619c4a7a3c4e88793b7bce3a023fb476 Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 10:32:12 +0800
Subject: [PATCH 05/10] [Bugfix] Accept H3 VAE integer parallel flags

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 21 +++++++++++++++++++
 .../models/minimax_h3/ops/vae/__init__.py     |  8 ++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

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
index 417dfda946e..9988bc71a8e 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -235,6 +235,26 @@ def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
     )
 
 
+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
 
@@ -306,6 +326,7 @@ def test_h3_vae_install_preserves_nondefault_official_semantics(monkeypatch, nam
     "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),
diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index f705a0ee433..28848b848ba 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -31,6 +31,12 @@ def _official_execution_semantics_supported() -> bool:
     )
 
 
+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 _optimized_feed_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
     if torch.compiler.is_compiling():
         return type(self).forward(self, hidden_states)
@@ -125,7 +131,7 @@ def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None:
         scale2 = getattr(block, "scale2", None)
         if (
             not all(isinstance(linear, nn.Linear) for linear in candidates)
-            or not isinstance(getattr(attention, "spatial_parallel", None), bool)
+            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))

From 19f3c3608a8ae76c4a1aa1c7df9a6749d2886aed Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 11:37:14 +0800
Subject: [PATCH 06/10] [Perf] Enable MiniMax-H3 VAE operators on SM103

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_vae_ops.py      | 18 ++++++++++++++++++
 .../models/minimax_h3/ops/vae/dispatch.py      | 17 +++++++++++++++--
 .../models/minimax_h3/ops/vae/qk_norm_rope.py  |  5 +++--
 3 files changed, 36 insertions(+), 4 deletions(-)

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
index 9988bc71a8e..e30d54226bc 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -359,3 +359,21 @@ def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
 
     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_only_validated_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, 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, 100, 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/vae/dispatch.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
index 4d890951f7c..66658900596 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
@@ -39,7 +39,7 @@ class H3VAEOperatorSet:
     scaled_residual: ScaledResidualOp
 
 
-def _supports_cuda_sm90(device: torch.device) -> bool:
+def _supports_cuda_capability(device: torch.device, expected: int) -> bool:
     if (
         not HAS_TRITON
         or device.type != "cuda"
@@ -50,7 +50,15 @@ def _supports_cuda_sm90(device: torch.device) -> bool:
         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()) == 90
+    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_sm103(device: torch.device) -> bool:
+    return _supports_cuda_capability(device, 103)
 
 
 # Keep hardware selection flat: adding a backend means adding one operator set,
@@ -61,6 +69,11 @@ def _supports_cuda_sm90(device: torch.device) -> bool:
         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,
+    ),
 )
 
 
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
index 95570acb34d..40b7069f61c 100644
--- 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
@@ -143,8 +143,9 @@ def try_qk_norm_rope_exact(
     sin = sin.reshape(tokens, sin.shape[-1])
     q_output = torch.empty_like(q)
     k_output = torch.empty_like(k)
-    # This SM90 launch layout is part of the exactness contract: changing the
-    # warp-to-row mapping changes the FP32 RMS reduction order.
+    # 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 = {

From f682aa84207ef936022ebce782a8f6b783cf18fd Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 11:45:42 +0800
Subject: [PATCH 07/10] [Perf] Enable MiniMax-H3 VAE operators on SM100

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../models/minimax_h3/test_minimax_h3_vae_ops.py         | 6 +++---
 .../diffusion/models/minimax_h3/ops/vae/dispatch.py      | 9 +++++++++
 2 files changed, 12 insertions(+), 3 deletions(-)

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
index e30d54226bc..7ae0538e2c3 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -361,7 +361,7 @@ def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
     assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
 
 
-def test_h3_vae_dispatch_selects_only_validated_cuda_capabilities(monkeypatch):
+def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
 
     platform = Mock()
@@ -370,10 +370,10 @@ def test_h3_vae_dispatch_selects_only_validated_cuda_capabilities(monkeypatch):
     monkeypatch.setattr(dispatch, "HAS_TRITON", True)
     monkeypatch.setattr(dispatch, "current_omni_platform", platform)
 
-    for capability in (90, 103):
+    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, 100, 110):
+    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/vae/dispatch.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
index 66658900596..d413ef85008 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py
@@ -57,6 +57,10 @@ 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)
 
@@ -69,6 +73,11 @@ def _supports_cuda_sm103(device: torch.device) -> bool:
         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,

From 108ab753c07e97c951f5c76250549379bad8b1de Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 12:33:28 +0800
Subject: [PATCH 08/10] [Docs] Explain MiniMax-H3 VAE operator ownership

Signed-off-by: mglyn <1203789601@qq.com>
---
 vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index 28848b848ba..00e41171e4d 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -1,6 +1,12 @@
 # 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."""
+"""Install bit-exact eager optimizations for the MiniMax H3 video VAE.
+
+These operations stay model-owned because correctness and performance depend
+on H3's exact math, dtype, layout, execution contract, and per-platform
+validation. Dispatch and fallbacks therefore live beside the model integration;
+migration to a shared operator tier can follow once #6305 defines one.
+"""
 
 from __future__ import annotations
 

From d34b75bd403bc8058e2dd24cc444342efd172161 Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 12:41:25 +0800
Subject: [PATCH 09/10] [Docs] Document MiniMax-H3 operator ownership

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../diffusion/models/minimax_h3/ops/README.md | 50 +++++++++++++++++++
 .../models/minimax_h3/ops/vae/__init__.py     |  8 +--
 2 files changed, 51 insertions(+), 7 deletions(-)
 create mode 100644 vllm_omni/diffusion/models/minimax_h3/ops/README.md

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/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index 00e41171e4d..28848b848ba 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -1,12 +1,6 @@
 # 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.
-
-These operations stay model-owned because correctness and performance depend
-on H3's exact math, dtype, layout, execution contract, and per-platform
-validation. Dispatch and fallbacks therefore live beside the model integration;
-migration to a shared operator tier can follow once #6305 defines one.
-"""
+"""Install bit-exact eager optimizations for the MiniMax H3 video VAE."""
 
 from __future__ import annotations
 

From 6030e5976a38d71fab755cf4a482828cb5a35117 Mon Sep 17 00:00:00 2001
From: mglyn <1203789601@qq.com>
Date: Fri, 28 Aug 2026 14:15:49 +0800
Subject: [PATCH 10/10] [Fix] Avoid model-specific H3 environment access

Signed-off-by: mglyn <1203789601@qq.com>
---
 .../minimax_h3/test_minimax_h3_vae_ops.py     | 26 ++++++++------
 .../models/minimax_h3/ops/vae/__init__.py     | 36 ++++++++++---------
 2 files changed, 35 insertions(+), 27 deletions(-)

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
index 7ae0538e2c3..7a1c54f7b31 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -48,6 +48,18 @@ def _qk_reference(x, cos, sin):
     )
 
 
+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()
@@ -295,24 +307,18 @@ def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
         assert block.attn.forward.__func__.__name__ == "forward"
 
 
-@pytest.mark.parametrize(
-    ("name", "value"),
-    [
-        ("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "0"),
-        ("MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "1"),
-        ("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "true"),
-    ],
-)
-def test_h3_vae_install_preserves_nondefault_official_semantics(monkeypatch, name, value):
+@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.setenv(name, value)
     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,
diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
index 28848b848ba..247e072c82b 100644
--- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
+++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py
@@ -4,7 +4,6 @@
 
 from __future__ import annotations
 
-import os
 from types import MethodType
 from typing import Any
 
@@ -16,25 +15,29 @@
 from .dispatch import resolve_h3_vae_operators
 
 
-def _env_flag(name: str, default: str) -> bool:
-    value = os.environ.get(name, default)
-    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+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 _official_execution_semantics_supported() -> bool:
-    """Return whether replacing the official eager forwards is semantics-safe."""
-
-    return (
-        _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1")
-        and not _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "0")
-        and not _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0")
-    )
 
+def _uses_fp32_attention_norm(attention: nn.Module) -> bool:
+    """Probe the loaded remote module's Q/K normalization semantics."""
 
-def _is_boolean_flag(value: Any) -> bool:
-    """Accept bools and the JSON-derived 0/1 flags used by remote VAE code."""
+    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
 
-    return isinstance(value, bool) or type(value) is int and value in (0, 1)
+    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:
@@ -140,6 +143,7 @@ def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None:
             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
@@ -185,8 +189,6 @@ def install_h3_vae_optimizations(
 
     if getattr(decoder, "_omni_h3_vae_optimizations_installed", False):
         return True
-    if not _official_execution_semantics_supported():
-        return False
     operators = resolve_h3_vae_operators(device)
     if operators is None:
         return False