diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py index 8bd78a5a0..c1ec12208 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py @@ -140,8 +140,8 @@ def WNConv1d(*args, **kwargs): return weight_norm(nn.Conv1d(*args, **kwargs)) -@torch.jit.script def snake(x, alpha): + # profiling JIT changes rounding after the first call when it fuses this graph shape = x.shape x = x.reshape(shape[0], shape[1], -1) x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index c83674a83..33db17fed 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -40,7 +40,7 @@ logger = init_logger(__name__) # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows # publish. SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion" -SGL_TEST_FILES_CI_DATA_REVISION = "11783b3fbd8ebb1e3509cc4590c1fff476e65511" +SGL_TEST_FILES_CI_DATA_REVISION = "285ffa7fa7a8afcfb90e2344e3f9d6ca12a5e2ea" # The NPU pin is kept as a separate branch so ascend GT can be bumped independently # when it's regenerated on its own cadence. diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py index 473bf04df..a6567a0e9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax-H3 released VAE decode contract.""" +import subprocess +import sys +import textwrap from unittest import mock import pytest @@ -117,3 +120,35 @@ def test_audio_vae_attention_defaults_to_local_sdpa_and_allows_fa(): } assert recording_fa.input_dtype == torch.bfloat16 assert output.dtype == torch.float32 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_audio_snake_first_call_matches_repeated_calls(): + # a fresh process prevents earlier tests from warming a profiling JIT graph + subprocess.run( + [ + sys.executable, + "-c", + textwrap.dedent( + """ + import torch + from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_audio_vae.audio_vae import Snake1d + + torch.manual_seed(42) + activation = Snake1d(64).cuda().eval() + with torch.inference_mode(): + activation.alpha.uniform_(0.1, 2.0) + x = torch.randn(2, 64, 4096, device="cuda") + original_x = x.clone() + original_alpha = activation.alpha.clone() + first = activation(x) + for _ in range(4): + torch.testing.assert_close(activation(x), first, rtol=0, atol=0) + torch.testing.assert_close(x, original_x, rtol=0, atol=0) + torch.testing.assert_close(activation.alpha, original_alpha, rtol=0, atol=0) + """ + ), + ], + check=True, + timeout=120, + )