Add tensor parallel support to stable audio open 1.0

[akshatvishu]

Part of #1217

Purpose

Add tensor parallel support to stable audio open

Test Plan

All experiments were conducted on Kaggle using 2× Tesla T4 GPUs.
Due to the hardware constraints, the maximum tensor parallelism (tp_size) tested was 2.

Environment

GPUs: 2 × Tesla T4

Precision: float16 & float32

Maximum tensor parallel size (TP): 2

Global Test Configuration:

STEPS = 100
AUDIO_LENGTH = 10.0  # seconds

Model Initialization:

parallel_config = DiffusionParallelConfig(
    tensor_parallel_size=tp_size
)

# Initialize Omni model
omni = Omni(
    model=MODEL_PATH,
    dtype="float16",
    parallel_config=parallel_config,
)

generator = torch.Generator(
    device=current_omni_platform.device_type
).manual_seed(SEED)

params = OmniDiffusionSamplingParams(
    num_inference_steps=STEPS,
    guidance_scale=7.0,
    generator=generator,
    extra_args={
        "audio_start_in_s": 0.0,
        "audio_end_in_s": AUDIO_LENGTH,
    },
omni.generate({"prompt": "A danceable electronic track in the genre of dance.", "negative_prompt": "Low quality, noisy"}, params)
)

Test Result

For float16 :

Name	tp_size	Time(seconds)	Speedup	File
hf_fp16	-	28.67	-	dance_track.mp3
baseline_tp1_fp16	1	26.40	-	TP_1_fp16.mp3
tp2_fp16	2	20.01	1.32x	TP_2_fp16.mp3

For float32 :

Name	tp_size	Time(sec)	Speedup	File
hf_fp32	-	194.15	-	dance_track_fp32.mp3
baseline_tp_1_fp32	1	145.00	-	TP_1_fp32.mp3
tp_2_fp32	2	78.01	1.86x	TP_2_fp32.mp3

• files are in .mp3 as github don't allows .wav
• hf_ prefix entries means that we run stable audio hugging face diffuser's pipeline with the same parameters.
• baseline_ prefix entries are the ones from against which the speedup is calculated.

For extended testing and results, please refer to this kaggle-notebook.

---

Notes

• At kaggle-notebook, The default final_sigmas_type="zero" in Diffusers produces a harmless torchsde boundary warnings when using CosineDPMSolverMultistepScheduler. This is expected upstream behavior and not a vLLM Omni issue.

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)

[lishunyang12]

Thanks for the contribution. The TP approach works, but broadcasting after every transformer block effectively defeats the purpose of tensor parallelism. I suspect there is a subtle numerical divergence elsewhere. A few high-level concerns and questions inline.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b5dbdf6ca6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[akshatvishu]

Thanks for the review! Just a heads-up: I accidentally performed a force-with-lease on a follow-up after the initial squash . This shifted the hunk headers, so my previous responses to your line-specific comments might point to the wrong code blocks now but I’ve manually verified that all requested changes are addressed.

[hsliuustc0106]

@vllm-omni-reviewer

[hsliuustc0106]

Hi @akshatvishu 👋

This tensor parallel support PR for stable audio hasn't been updated for 16 days. Is this still being worked on?

Thanks!

[akshatvishu]

Hey @hsliuustc0106 !
I think apart from a minor style fix in one of the comment ; this is ready to go from my side! Happy to provide any more test if needed!

Update: Done with the style changes!

[lishunyang12]

Any progress?

[akshatvishu]

Any progress?

Its ready for my side! Any test or benchmark you want me to test/run?

[linyueqian]

🟡 [important] VAE chunk_size = 1 serializes decode for every caller, not just T4.

The earlier code did self.vae.decode(latents_for_vae).sample in a single call. Hardcoding chunk_size = 1 here means that once num_waveforms_per_prompt > 1 or batched inference lands, users on H100 / A100 pay a serial VAE decode that they do not need. The comment says safe default for 47s audio on T4 but the default is now applied to everyone.

Could we either:

1. Expose vae_chunk_size as a sampling param / config option, defaulting to latents_for_vae.shape[0] (whole-batch decode, matches old behavior), so T4 users can opt into vae_chunk_size=1, or
2. Pick the chunk size dynamically from available VRAM or batch size so the common single-waveform path stays a single decode call?

As-is, this is a silent perf regression for non-T4 users.

[akshatvishu]

Took approach 1: vae_chunk_size is now an explicit sampling param with default latents_for_vae.shape[0], so the common path stays a single VAE decode call. vae_chunk_size=1 is still available for low-VRAM use.

I added a unit test for the chunking behavior. Should I keep that, or remove it?

[linyueqian]

@princepride ptal as well thanks

[akshatvishu]

failed tests:

=========================== short test summary info ============================
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_modulated_input_shape
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_run_transformer_blocks_callable
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_postprocess_callable
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_postprocess_output_shape
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_postprocess_return_tuple_when_return_dict_false
FAILED tests/diffusion/cache/test_teacache_extractors.py::TestFluxExtractor::test_without_guidance

error:

E NotImplementedError: Could not run 'vllm::rocm_unquantized_gemm' with arguments from the 'CPU' backend

The CI CPU tests for FLUX extractors are crashing on ROCm builds probably because ReplicatedLinear dispatches to GPU kernels. I think we can fix this by using the same monkeypatch logic used in the adalayernorm tests to force the default CPU-compatible GEMM path during test execution.

Happy to raise a PR to fix this if this is indeed the preferred way forward!

ref:
https://github.com/vllm-project/vllm-omni/blob/main/tests/diffusion/layers/test_adalayernorm.py#L37

cc: @linyueqian