[Test] Add Stable Audio offline e2e TeaCache Test

[zhangj1an]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

This PR adds the Stable-Audio-Open offline e2e TeaCache test as part of the diffusion-feature L4 test coverage effort tracked in #1832.

For Stable-Audio-Open, TeaCache (#1314) is currently the only supported diffusion acceleration feature in vllm-omni. Other diffusion features listed in #1217, such as Tensor Parallel (#1406), Cache-DiT (#1341), CFG parallelism, VAE parallelism, and HSDP, are not yet supported for this model, so they are intentionally out of scope for this PR. Extended coverage will be covered in future PRs.

Key changes

• for Stable Audio test, split into L2 and L4 level. for L4 level, combine quantisation and cache into the same test. right now cache_backend="tea_cache"
• update Stable-Audio-Open documentation / support matrix to mark TeaCache as supported
• update Stable Audio configs to expose --cache-backend tea_cache and --tea-cache-rel-l1-thresh, following [skip ci][Doc] Refine the Diffusion Features User Guide #1928

Test plan

Run:

pytest -s -v tests/e2e/offline_inference/test_stable_audio.py tests/e2e/offline_inference/test_stable_audio_expansion.py

Test result

2 passed, 16 warnings in 32.91s.

============================= test session starts ==============================
platform linux -- Python 3.11.10, pytest-9.0.2, pluggy-1.6.0 -- /root/vllm-omni/.venv/bin/python
cachedir: .pytest_cache
rootdir: /root/vllm-omni
configfile: pyproject.toml
plugins: anyio-4.13.0, mock-3.15.1, cov-7.1.0, asyncio-1.3.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collecting ... collected 2 items

tests/e2e/offline_inference/test_stable_audio.py::test_stable_audio[linyueqian/stable_audio_random] PASSED [ 50%]
tests/e2e/offline_inference/test_stable_audio_expansion.py::test_stable_audio_quantization_and_teacache[linyueqian/stable_audio_random] PASSED [100%]

=============================== warnings summary ===============================
<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

.venv/lib/python3.11/site-packages/torch/jit/_script.py:362: 14 warnings
  /root/vllm-omni/.venv/lib/python3.11/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================= 2 passed, 16 warnings in 32.91s ========================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

cc: @fhfuih @akshatvishu @linyueqian

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[fhfuih]

Thanks for the contribution. The code itself generally looks good to me. But according to our convention, the expanded coverage test (L4) should be placed in another file test_stable_audio_expansion.

(Also, could you help rename this file to test_stable_audio to match the new L2 test naming convention as well? Thanks)

And the L4 test should have the label advanced_model instead of core_model. You can refer to .buildkite/test-ready and test-merge YAML files to see which and how tests are selected. You can also copy the pytest command with --collect-only to test yourself without actually running.

For the run-validate helper function, could you please also remove this function and reuse assert_audio_valid function in tests/conftest.py? This function can be reused in both two files. This function is currently not used anywhere else (because there is no other audio generation models), so you can debug or change its signature per your need.

Finally, it is possible to combine quantization and cache in one test case in its L4 test? This way the test coverage is more complete.

Thanks for your attention. I presume the maintenance of this model has been stale for quite a while now, and now is a good change to consolidate the new test guidance 😁

[zhangj1an]

Hey @fhfuih, thanks for the comprehensive feedback!

I split the original test file into test_stable_audio and test_stable_audio_expansion, also applied all your suggestions. I re-ran the test and they still work, as shown in updated test result in PR description.

[fhfuih]

Thanks for the contribution. Looks good to me now. A few comments for others who stumble on this PR:

• Since #1794 is still not available, stable audio test uses offline mode.
• The offline test execution helper omni_runner as a counterpart to omni_server does not currently support diffusion yet. @yenuo26 and I are aware of this matter and planning to ship update with some other refactoring. It is currently good for now that this tests creates it own helper

[fhfuih]

Pending a nightly-test tag to run CI. @wtomin @congw729 PTAL

[zhangj1an]

Hey @yenuo26, the CI has passed, please help me add the ready label. Thank you! 🙌

[lishunyang12]

looks good overall, just a minor nit

[akshatvishu]

Hey @zhangj1an ! cache-dit (#1341) got merged today, so feel free to expand the tests for it. If you’re busy with other work, I’m happy to add them myself.

[lishunyang12]

LGTM. Fix conflict please.

[fhfuih]

Hey @zhangj1an ! cache-dit (#1341) got merged today, so feel free to expand the tests for it. If you’re busy with other work, I’m happy to add them myself.

Hey @akshatvishu, congrats on stable audio cache-dit (#1341) getting merged! 🚀 I’ll add the unit test today. 👍

Update: I think it may not be necessary to add a cache-dit unit test for Stable Audio, since it’s not a high-priority model and this PR already covers FP8 + TeaCache. What do you think? @fhfuih

Otherwise, the new test can be added in tests/e2e/offline_inference/test_stable_audio_expansion.py as below.

Agree. No need to cover the complete feature set for this model. Just one or two most common one will do. I think this PR is all good for now. And we can ignore the AMD CI.

And it is OK to not always keep the branch updated (if there is no conflict). Otherwise it keeps aborting CI test and restarting a new one. The repetitive CI request also wastes the GPU resource.

Let's wait for the current version to pass the CI (the CUDA one)

[fhfuih]

Seems the primary CI has passed: https://buildkite.com/vllm/vllm-omni/builds/6480/steps/canvas

But the doc CI fails due to occasional timeout. I think it is OK to remove the nightly tag and put a ready tag, trigger a quick rerun and merge. @wtomin @Gaohan123 PTAL Thanks