[Test] [AMD] diffusion: Make multimodal_gen unit tests ROCm-aware

python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py

@@ -2,13 +2,15 @@
 """CPU-safe FLUX.3 tests for the rowwise-FP8 and ModelOpt NVFP4 DiT paths."""
 
 import contextlib
+import importlib.util
 import json
 import tempfile
 import unittest
 from pathlib import Path
 from types import SimpleNamespace
 from unittest import mock
 
+import pytest
 import torch
 from safetensors.torch import save_file
 
@@ -38,6 +40,14 @@
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
 
+# Building the DiT constructs attention layers, which read the SP group size.
+pytestmark = pytest.mark.usefixtures("single_process_model_parallel")
+
+_requires_flashinfer = unittest.skipUnless(
+    importlib.util.find_spec("flashinfer"),
+    "the PrimsTS weight pack ships with flashinfer (CUDA-only)",
+)
+
 _LEAVES = ("to_qkv", "mlp_in", "attn_out", "mlp_out")
 _ALL_ROLES = {
     "to_qkv": Flux3PrimsTsRole.QKV_QKNORM_ROPE,
@@ -207,6 +217,7 @@ def _roles(block):
 
 
 class TestFlux3Fp8Construction(unittest.TestCase):
+    @_requires_flashinfer
     def test_process_weights_prepares_prims_ts_weights(self):
         from sglang.kernels.ops.diffusion import pack_flux3_mlp_weight
 
@@ -258,6 +269,7 @@ def test_tp_local_shard_shapes_gate_prims_ts(self):
         self.assertEqual(block.to_qkv.output_size_per_partition, 768)
         self.assertEqual(_roles(block), {**_ALL_ROLES, "to_qkv": None})
 
+    @_requires_flashinfer
     def test_swiglu_layout_rejects_generic_apply(self):
         layer = _tiny_native_model(_flux3_fp8_config()).single_blocks[0].mlp_in
         layer.quant_method.process_weights_after_loading(layer)
@@ -270,6 +282,7 @@ def test_metadata_declared_bf16_path_stays_unquantized(self):
         self.assertIsInstance(block.to_qkv.quant_method, Flux3Fp8LinearMethod)
         self.assertIsInstance(block.mlp_in.quant_method, UnquantizedLinearMethod)
 
+    @_requires_flashinfer
     def test_fused_mlp_weight_pack_reverses_each_gate_value_pair(self):
         from sglang.kernels.ops.diffusion import pack_flux3_mlp_weight
 

python/sglang/multimodal_gen/test/unit/test_flux3_prims_ts_config.py

@@ -4,7 +4,11 @@
 import unittest
 from unittest import mock
 
-from flashinfer.prims_ts import gemm as prims_ts_gemm
+import pytest
+
+prims_ts_gemm = pytest.importorskip(
+    "flashinfer.prims_ts.gemm", reason="PrimsTS ships with flashinfer (CUDA-only)"
+)
 
 from sglang.kernels.ops.diffusion.gemm import flux3_prims_ts_config as config_module
 

python/sglang/multimodal_gen/test/unit/test_kandinsky6_layerwise_offload.py

@@ -36,7 +36,10 @@
     Kandinsky6AudioVAE,
 )
 from sglang.multimodal_gen.runtime.models.vaes.kandinsky6_sr_vae import Kandinsky6SRVAE
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
 from sglang.multimodal_gen.runtime.utils import precision
 from sglang.multimodal_gen.test.unit.kandinsky6_sr_tiny_components import (
@@ -219,6 +222,11 @@ def test_layerwise_groups_preserve_repeated_forwards(component, resident_layers)
     assert set(groups) <= set(model.layer_names)
     assert model.layerwise_offload_dit_group_enabled == component.endswith("dit")
     state = {name: value.cpu().clone() for name, value in model.state_dict().items()}
+    tolerance = dict(rtol=0, atol=0)
+    if component == "audio-vae" and current_platform.is_rocm():
+        # On MI300 the offloaded fp32 audio VAE differs from the resident run by
+        # ~1e-6 (CUDA is bit-exact); the weights are still compared exactly below.
+        tolerance = dict(rtol=1e-4, atol=1e-5)
     with (
         torch.inference_mode(),
         set_forward_context(current_timestep=0, attn_metadata=None),
@@ -241,12 +249,12 @@ def test_layerwise_groups_preserve_repeated_forwards(component, resident_layers)
         } == set(groups)
         for _ in range(2):
             model.prepare_for_next_req()
-            torch.testing.assert_close(forward(), expected, rtol=0, atol=0)
+            torch.testing.assert_close(forward(), expected, **tolerance)
             for manager in model.layerwise_offload_managers:
                 assert manager._last_forwarded_layer == manager.num_layers - 1
                 manager.release_all()
         model.disable_offload()
-        torch.testing.assert_close(forward(), expected, rtol=0, atol=0)
+        torch.testing.assert_close(forward(), expected, **tolerance)
     for name, value in model.state_dict().items():
         torch.testing.assert_close(value.cpu(), state[name], rtol=0, atol=0)
 

python/sglang/multimodal_gen/test/unit/test_moss_image_parallel.py

@@ -128,6 +128,11 @@ def _worker(mode):
     dist.destroy_process_group()
 
 
+@pytest.mark.skipif(
+    torch.version.hip is not None,
+    reason="ROCm builds query GPU 0 while importing sglang, which fails with the "
+    "GPUs hidden for these CPU workers",
+)
 @pytest.mark.parametrize("mode", ["sp", "cfg"])
 def test_two_rank_moss_image_matches_serial(mode):
     result = subprocess.run(

python/sglang/multimodal_gen/test/unit/test_qwen3_5_encoder.py

@@ -101,10 +101,15 @@ def test_native_encoder_registry_and_checkpoint(small_text_encoder):
 
 @pytest.mark.parametrize("backend", [None, "torch_sdpa", "fa"])
 def test_text_attention_default_and_explicit_selection(small_text_encoder, backend):
-    from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+    from sglang.multimodal_gen.runtime.platforms import (
+        AttentionBackendEnum,
+        current_platform,
+    )
 
     if backend == "fa" and not torch.cuda.is_available():
         pytest.skip("FA selection requires CUDA")
+    if backend == "fa" and current_platform.is_rocm():
+        pytest.skip("the platform currently resolves FA to SDPA on ROCm")
     get_global_server_args().attention_backend = backend
     config = Qwen3_5EncoderConfig()
     config.update_model_arch(small_text_encoder.config.to_dict())

python/sglang/multimodal_gen/test/unit/test_residual_gate_add2.py

@@ -25,6 +25,11 @@ def test_kernel_is_bit_identical_to_addcmul(self):
                     # Raises instead of falling back when the predicate rejects.
                     got = residual_gate_add2_cuda(residual, a, b, gate)
                     want = torch.addcmul(residual, gate, a + b)
+                    if torch.version.hip:
+                        # On MI300 the fp16 scalar path is not bit-identical to
+                        # addcmul; hold it to the dtype's rounding tolerance.
+                        torch.testing.assert_close(got, want)
+                        continue
                     self.assertTrue(torch.equal(got, want), (dtype, shape, gate.shape))
                 expanded_gate = full_gate[:1].expand_as(residual)
                 self.assertFalse(

python/sglang/multimodal_gen/test/unit/test_rowwise_fp8.py

@@ -174,7 +174,9 @@ def test_per_tensor_serialized_fp8_keeps_a_scalar_weight_scale(self):
 
 
 def _fp8_capable():
-    if not torch.cuda.is_available():
+    # ROCm only takes FP8 weights through the loader, which converts them to
+    # the platform's FP8 layout; this test feeds raw OCP e4m3fn weights.
+    if not torch.cuda.is_available() or torch.version.hip:
         return False
     major, minor = torch.cuda.get_device_capability()
     return major >= 9 or (major == 8 and minor >= 9)