From d0e2d1deef628dacc07c763402c0dafa250f457d Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:31:23 -0700 Subject: [PATCH 1/3] [diffusion] fix: make five multimodal_gen unit tests ROCm-aware - test_flux3_prims_ts_config: importorskip flashinfer, which ROCm does not ship - test_rowwise_fp8: skip the raw OCP e4m3fn apply test on ROCm - test_qwen3_5_encoder: skip the FA selection case on ROCm, where FA resolves to SDPA - test_residual_gate_add2: keep bit-exactness on CUDA, use fp16/bf16 tolerance on ROCm - test_kandinsky6_layerwise_offload: small tolerance for the fp32 audio VAE forward on ROCm; weights are still compared exactly --- .../test/unit/test_flux3_prims_ts_config.py | 6 +++++- .../test/unit/test_kandinsky6_layerwise_offload.py | 14 +++++++++++--- .../test/unit/test_qwen3_5_encoder.py | 7 ++++++- .../test/unit/test_residual_gate_add2.py | 5 +++++ .../multimodal_gen/test/unit/test_rowwise_fp8.py | 4 +++- 5 files changed, 30 insertions(+), 6 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_flux3_prims_ts_config.py b/python/sglang/multimodal_gen/test/unit/test_flux3_prims_ts_config.py index 9197b514ff01..70a2ad7274e1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_flux3_prims_ts_config.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_kandinsky6_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_kandinsky6_layerwise_offload.py index 5e27f81017fc..5fbe69d5327d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_kandinsky6_layerwise_offload.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen3_5_encoder.py b/python/sglang/multimodal_gen/test/unit/test_qwen3_5_encoder.py index 2f2895538e32..ffb874ca8ee8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen3_5_encoder.py +++ b/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()) diff --git a/python/sglang/multimodal_gen/test/unit/test_residual_gate_add2.py b/python/sglang/multimodal_gen/test/unit/test_residual_gate_add2.py index 042444601563..3db5d97052b4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_residual_gate_add2.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/unit/test_rowwise_fp8.py b/python/sglang/multimodal_gen/test/unit/test_rowwise_fp8.py index 0622c93df62a..a19e55426ca0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_rowwise_fp8.py +++ b/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) From db8dc3480dc0efba67fedfd3e4c5af658e8dade0 Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:38:40 -0700 Subject: [PATCH 2/3] [diffusion] fix: two more multimodal_gen unit tests on ROCm - test_flux3_image_dit_quant: use the single_process_model_parallel fixture (building the DiT reads the SP group size) and skip the two cases that import the flashinfer-backed PrimsTS weight pack when flashinfer is absent - test_moss_image_parallel: skip on ROCm, where importing sglang queries GPU 0 and fails with the GPUs hidden for the CPU workers --- .../test/unit/test_flux3_image_dit_quant.py | 12 ++++++++++++ .../test/unit/test_moss_image_parallel.py | 5 +++++ 2 files changed, 17 insertions(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py b/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py index 6c91d6a83965..16fc428de7d7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py @@ -2,6 +2,7 @@ """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 @@ -9,6 +10,7 @@ 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 @@ -270,6 +281,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 diff --git a/python/sglang/multimodal_gen/test/unit/test_moss_image_parallel.py b/python/sglang/multimodal_gen/test/unit/test_moss_image_parallel.py index 2b4d52660fb6..67d3dd1f14dd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_moss_image_parallel.py +++ b/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( From cfd9bf0a62dcba9621038ce0b55d24fe6f115974 Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:50:06 -0700 Subject: [PATCH 3/3] [diffusion] fix: skip the FLUX.3 SwiGLU layout test without flashinfer Its process_weights_after_loading packs the PrimsTS SwiGLU weight, which imports flashinfer. --- .../multimodal_gen/test/unit/test_flux3_image_dit_quant.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py b/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py index 16fc428de7d7..95e669cbea84 100644 --- a/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_flux3_image_dit_quant.py @@ -269,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)