[Bugfix][MiniMax-H3] Use Hopper-safe modulation precision

Signed-off-by: david6666666 <530634352@qq.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py

@@ -43,7 +43,7 @@ def _collective_worker(rank, rendezvous, backend):
                 assert value.item() == 2
 
             def encode_image(self, image):
-                assert rank == 0
+                self.collective()
                 assert image.size == (32, 32)
                 return torch.ones(1, 96)
 
@@ -115,23 +115,6 @@ def _collective_worker(rank, rendezvous, backend):
         torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
         assert conditioning.audio_condition_lengths == (80,)
 
-        image_media = (
-            MiniMaxH3EncoderMediaInput(
-                task="fl2va",
-                height=32,
-                width=32,
-                num_frames=96,
-                latent_t=29,
-                audio_t=160,
-                images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
-            )
-            if rank == 0
-            else None
-        )
-        image_conditioning = model._encode_local_media(image_media)
-        torch.testing.assert_close(image_conditioning.visual_condition.cpu(), torch.ones(1, 96))
-        assert image_conditioning.visual_condition_shapes == ((1, 2, 2),)
-
         model.video_vae.fail = True
         with pytest.raises(OmniClientError, match="codec failed on rank 1"):
             model._encode_local_media(media)

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -10,6 +10,7 @@ from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,
+    _use_hopper_bf16_affine_semantics,
     indexed_gate,
     indexed_gate_rms_norm_scale_shift,
     indexed_scale_shift_,
@@ -190,3 +191,55 @@ def test_fused_modulation_preserves_bf16_residual_boundary() -> None:
     )
 
     assert torch.equal(modulated_out, expected)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(
+    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] != 9,
+    reason="SM90-specific precision path",
+)
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_hopper_fused_modulation_matches_pytorch_reference() -> None:
+    torch.manual_seed(42)
+    rows, hidden_size, conditions = 128, 3072, 4
+    residual = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    branch = torch.randn(rows, hidden_size, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(conditions, hidden_size, device="cuda", dtype=torch.bfloat16)
+    indices = torch.arange(rows, device="cuda") % conditions
+    eps = 1e-6
+
+    reference_residual = (residual + gate.index_select(0, indices) * branch).to(torch.bfloat16)
+    normalized = reference_residual.float()
+    variance = normalized.pow(2).mean(-1, keepdim=True)
+    normalized = (weight.float() * (normalized * torch.rsqrt(variance + eps))).to(torch.bfloat16)
+    reference_modulated = (
+        normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)
+    ).to(torch.bfloat16)
+
+    actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift(
+        residual,
+        gate,
+        branch,
+        weight,
+        shift,
+        scale,
+        indices,
+        eps,
+    )
+
+    torch.testing.assert_close(actual_residual, reference_residual, atol=5e-2, rtol=5e-2)
+    torch.testing.assert_close(actual_modulated, reference_modulated, atol=5e-2, rtol=5e-2)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_hopper_precision_dispatch_is_architecture_scoped(monkeypatch) -> None:
+    tensor = torch.empty(0, device="cuda")
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0))
+    assert _use_hopper_bf16_affine_semantics(tensor)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
+    assert not _use_hopper_bf16_affine_semantics(tensor)

vllm_omni/diffusion/layers/indexed_modulation.py

@@ -24,6 +24,18 @@ def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) ->
         kernel[(chunk_rows,)](*args, row_offset, **kwargs)
 
 
+def _use_hopper_bf16_affine_semantics(x: torch.Tensor) -> bool:
+    """Use the unfused AdaLN dtype boundaries on Hopper.
+
+    The fused FP32 modulation path is stable on Blackwell, but the SM90
+    kernel can accumulate enough extra precision in the repeated RMSNorm and
+    AdaLN operations to move the 50-step I2VA result below its CI threshold.
+    Keep the existing path on other architectures until their numerical
+    baselines are independently validated.
+    """
+    return x.device.type == "cuda" and torch.cuda.get_device_capability(x.device)[0] == 9
+
+
 @triton.jit
 def _indexed_scale_shift_kernel(
     output_ptr,
@@ -102,6 +114,7 @@ def _rms_norm_indexed_scale_shift_kernel(
     stride_scale_row,
     stride_indices,
     row_offset,
+    use_bf16_affine: tl.constexpr,
     block_n: tl.constexpr,
 ):
     row = tl.program_id(0) + row_offset
@@ -112,13 +125,20 @@ def _rms_norm_indexed_scale_shift_kernel(
     x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
     weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
     variance = tl.sum(x * x, axis=0) / hidden_size
-    normalized = x * tl.rsqrt(variance + eps) * weight
-
-    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
-    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    if use_bf16_affine:
+        # Match the unfused RMSNorm output consumed by AdaLN on Hopper.
+        normalized = (x * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
+        shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
+        scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
+        output = normalized * (scale + 1.0) + shift
+    else:
+        normalized = x * tl.rsqrt(variance + eps) * weight
+        shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+        output = normalized * (1.0 + scale) + shift
     tl.store(
         output_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        output,
         mask=mask,
     )
 
@@ -143,6 +163,7 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
     stride_scale_row,
     stride_indices,
     row_offset,
+    use_bf16_affine: tl.constexpr,
     block_n: tl.constexpr,
 ):
     row = tl.program_id(0) + row_offset
@@ -159,12 +180,20 @@ def _indexed_gate_rms_norm_scale_shift_kernel(
 
     weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
     variance = tl.sum(updated * updated, axis=0) / hidden_size
-    normalized = updated * tl.rsqrt(variance + eps) * weight
-    shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
-    scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+    if use_bf16_affine:
+        # Match the unfused RMSNorm output consumed by AdaLN on Hopper.
+        normalized = (updated * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
+        shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
+        scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
+        output = normalized * (scale + 1.0) + shift
+    else:
+        normalized = updated * tl.rsqrt(variance + eps) * weight
+        shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+        scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+        output = normalized * (1.0 + scale) + shift
     tl.store(
         modulated_out_ptr + row * hidden_size + columns,
-        normalized * (1.0 + scale) + shift,
+        output,
         mask=mask,
     )
 
@@ -271,6 +300,7 @@ def rms_norm_indexed_scale_shift(
             shift.stride(0),
             scale.stride(0),
             indices.stride(0),
+            use_bf16_affine=_use_hopper_bf16_affine_semantics(x),
             block_n=triton.next_power_of_2(hidden_size),
             num_warps=8,
         )
@@ -324,6 +354,7 @@ def indexed_gate_rms_norm_scale_shift(
             shift.stride(0),
             scale.stride(0),
             indices.stride(0),
+            use_bf16_affine=_use_hopper_bf16_affine_semantics(residual),
             block_n=triton.next_power_of_2(hidden_size),
             num_warps=8,
         )

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -1953,22 +1953,7 @@ class MiniMaxH3Pipeline(
 
     def _encode_local_media(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaConditioning:
         group, rank, world_size = _dit_rank_world()
-        # Keep image/keyframe encoding on the output rank.  The legacy
-        # single-stage path encoded images locally and only used the tiled VAE
-        # collective for reference videos.  Treating every visual input as a
-        # distributed video after the unified encoder path was introduced
-        # changes the H100 keyframe latent and therefore the I2VA trajectory.
-        distributed_vae = self.video_vae.is_distributed_enabled()
-        distributed_video = False
-        if distributed_vae:
-            has_video = torch.tensor(
-                [int(media is not None and bool(media.videos))],
-                dtype=torch.long,
-                device=self.device,
-            )
-            if world_size > 1:
-                dist.broadcast(has_video, src=0, group=group)
-            distributed_video = bool(has_video.item())
+        distributed_video = self.video_vae.is_distributed_enabled()
         if distributed_video:
             media = self._distribute_media_inputs(media)
         conditioning = None
@@ -1981,7 +1966,6 @@ class MiniMaxH3Pipeline(
                     video_vae=self.video_vae,
                     audio_vae=self.audio_vae,
                     emit_conditioning=rank == 0,
-                    encode_images=rank == 0,
                     component_scope=self._component_on_device,
                 )
         except ValueError as exc:

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -602,13 +602,8 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         if getattr(self, "decode_only", False):
             raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode images")
         previous_parallel = self.model.parallel_tiling
-        parallel_size = int(getattr(self, "parallel_size", 1))
-        if parallel_size <= 1:
+        if int(getattr(self, "parallel_size", 1)) <= 1:
             self.model.parallel_tiling = False
-        # Keyframe conditioning follows the official single-image contract.
-        # Do not let it enter the tiled VAE collective even when reference
-        # videos use the same VAE instance with patch parallelism enabled.
-        tiling_context = self._rank_local_tiling() if parallel_size > 1 else nullcontext()
         parameter = next(self.parameters())
         previous_dtype = parameter.dtype
         if previous_dtype != torch.float32:
@@ -616,7 +611,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         devices = [parameter.device] if parameter.device.type != "cpu" else []
         try:
             with (
-                tiling_context,
+                self._encoder_tiling_context(image.height, image.width),
                 torch.random.fork_rng(
                     devices=devices,
                     device_type=parameter.device.type,

vllm_omni/model_executor/models/minimax_h3/encoder_processing.py

@@ -481,7 +481,6 @@ def encode_media(
     video_vae: Any,
     audio_vae: Any,
     emit_conditioning: bool,
-    encode_images: bool = True,
     component_scope: Callable[[Any], AbstractContextManager[Any]] = nullcontext,
 ) -> MiniMaxH3EncoderMediaConditioning | None:
     """Run codecs on participating ranks; only the output rank runs audio.
@@ -496,11 +495,10 @@ def encode_media(
         if video_vae is None:
             raise RuntimeError("MiniMax H3 video VAE is not resident on this rank")
         with component_scope(video_vae):
-            if encode_images:
-                for value in media.images:
-                    image = _image_from_tensor(value)
-                    visual_rows.append(video_vae.encode_image(image))
-                    visual_shapes.append((1, image.height // 16, image.width // 16))
+            for value in media.images:
+                image = _image_from_tensor(value)
+                visual_rows.append(video_vae.encode_image(image))
+                visual_shapes.append((1, image.height // 16, image.width // 16))
             for value in media.videos:
                 frames = np.asarray(value.detach().cpu().to(torch.uint8).contiguous().numpy())
                 rows, shape = video_vae.encode_video(frames)