From 5f50fc3daf4e65ce777c2ecacd974f11d29d6bf4 Mon Sep 17 00:00:00 2001 From: Yash Jain Date: Mon, 14 Sep 2026 23:25:25 +0100 Subject: [PATCH] [Kernel] Fuse Ming streaming ISTFT overlap-add and normalization (#7338) Signed-off-by: Yash Jain --- benchmarks/kernels/benchmark_ming_istft.py | 218 ++++++++++++++++++ .../models/ming_tts/test_fused_istft.py | 147 ++++++++++++ .../models/common/ming/audio_dsp.py | 22 +- .../models/common/ming/fused_istft.py | 177 ++++++++++++++ 4 files changed, 563 insertions(+), 1 deletion(-) create mode 100644 benchmarks/kernels/benchmark_ming_istft.py create mode 100644 tests/model_executor/models/ming_tts/test_fused_istft.py create mode 100644 vllm_omni/model_executor/models/common/ming/fused_istft.py diff --git a/benchmarks/kernels/benchmark_ming_istft.py b/benchmarks/kernels/benchmark_ming_istft.py new file mode 100644 index 000000000..b67c9bbfc --- /dev/null +++ b/benchmarks/kernels/benchmark_ming_istft.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Measure Ming ISTFT and optionally its pretrained streaming AudioVAE decoder. + +Example (no model weights needed for the operator benchmark): + python benchmarks/kernels/benchmark_ming_istft.py --compile --iterations 100 + +Add --decoder to measure the published decoder with fixed synthetic latents. +This is a decoder-stage benchmark, not text-to-speech serving latency. +""" + +import argparse +import json +import statistics +import time +from collections.abc import Callable +from functools import partial + +import torch + +from vllm_omni.model_executor.models.common.ming.audio_dsp import ISTFT +from vllm_omni.model_executor.models.common.ming.fused_istft import fused_istft + + +def measure(operation: Callable, warmup: int, iterations: int) -> dict: + for _ in range(warmup): + operation() + torch.accelerator.synchronize() + samples = [] + for _ in range(iterations): + start = time.perf_counter() + operation() + torch.accelerator.synchronize() + samples.append((time.perf_counter() - start) * 1000) + torch.accelerator.reset_peak_memory_stats() + before = torch.accelerator.memory_allocated() + operation() + torch.accelerator.synchronize() + return { + "p50_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "stdev_ms": statistics.pstdev(samples), + "extra_peak_mib": (torch.accelerator.max_memory_allocated() - before) / 1024**2, + } + + +def native_post_fft(module, frames, audio, window): + """The native middle-chunk post-FFT sequence, isolated from cuFFT.""" + count = frames.shape[-1] + size = (count - 1) * module.hop_length + module.win_length + + def fold(value): + return torch.nn.functional.fold( + value, output_size=(1, size), kernel_size=(1, module.win_length), stride=(1, module.hop_length) + )[:, 0, 0, :] + + samples = fold(frames * module.window[None, :, None]) + samples, audio = module._buffer_process(samples, audio, 0, streaming=True) + envelope = fold(module.window.square().expand(1, count, -1).transpose(1, 2)) + envelope, window = module._buffer_process(envelope, window, 0, streaming=True) + if not (envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + return samples / envelope, audio, window + + +def assert_outputs(actual, reference): + maximum = 0.0 + for value, expected in zip(actual, reference, strict=True): + if expected is None: + assert value is None + else: + torch.testing.assert_close(value, expected, atol=5e-7, rtol=1e-5) + maximum = max(maximum, (value - expected).abs().max().item()) + return maximum + + +def benchmark_istft(args): + for dtype in (torch.float32, torch.bfloat16): + module = ISTFT(3528, 882, 3528).to(device="cuda", dtype=dtype) + for count in (16, 100, 400): + spec = torch.randn(1, 1765, count, device="cuda", dtype=torch.complex64) + _, audio, window = module.forward_native(spec, streaming=True) + frames = torch.fft.irfft(spec, 3528, dim=1) + for scope in ("post_fft", "istft"): + if scope == "post_fft": + native = partial(native_post_fft, module, frames, audio, window) + fused = partial(fused_istft, frames, module.window, 882, audio, window, True, False) + else: + native = partial(module.forward_native, spec, audio, window, streaming=True) + fused = partial(module, spec, audio, window, streaming=True) + error = assert_outputs(fused(), native()) + results = {"native": measure(native, args.warmup, args.iterations)} + if args.compile: + compiled = torch.compile(native) + compiled_output, reference = compiled(), native() + # Inductor can change BF16 envelope rounding. Report that + # difference without weakening the fused-kernel parity gate. + compiled_matches = all( + torch.allclose(value, expected, atol=5e-7, rtol=1e-5) + for value, expected in zip(compiled_output, reference, strict=True) + ) + results["compiled"] = measure(compiled, args.warmup, args.iterations) + results["compiled"]["matches_native_tolerance"] = compiled_matches + results["compiled"]["waveform_max_abs_error"] = ( + (compiled_output[0] - reference[0]).abs().max().item() + ) + results["fused"] = measure(fused, args.warmup, args.iterations) + print( + json.dumps( + { + "scope": scope, + "window_dtype": str(dtype), + "frames": count, + "max_abs_error": error, + "results": results, + } + ), + flush=True, + ) + + +def benchmark_decoder(args): + from huggingface_hub import hf_hub_download + from safetensors import safe_open + + from vllm_omni.model_executor.models.common.ming.audio_vae import Decoder + + model = "inclusionAI/Ming-omni-tts-0.5B" + revision = "9154772e7fbc585907b6237e3190790676f28975" + config_path = hf_hub_download(model, "config.json", revision=revision) + with open(config_path) as source: + config = json.load(source)["audio_tokenizer_config"] + kwargs = config["dec_kwargs"] + backbone = kwargs["backbone"] + backbone["_attn_implementation"] = "sdpa" + decoder = Decoder(backbone, kwargs["output_dim"], kwargs["latent_dim"], config["patch_size"]) + checkpoint = hf_hub_download(model, "model.safetensors", revision=revision) + prefix = "audio.decoder." + with safe_open(checkpoint, framework="pt") as weights: + state = {key.removeprefix(prefix): weights.get_tensor(key) for key in weights.keys() if key.startswith(prefix)} + decoder.load_state_dict(state, strict=True) + decoder = decoder.eval().to(device="cuda", dtype=torch.bfloat16) + chunks = [ + torch.randn(1, size, kwargs["latent_dim"], device="cuda", dtype=torch.bfloat16) for size in (4, 25, 25, 5) + ] + istft = decoder.head.istft + fused_forward = istft.forward + + def run(): + cache = state = None + outputs = [] + for index, chunk in enumerate(chunks): + output, state, cache = decoder.low_level_reconstruct( + chunk, past_key_values=cache, use_cache=True, stream_state=state, last_chunk=index == len(chunks) - 1 + ) + outputs.append(output) + return torch.cat(outputs, dim=-1) + + istft.forward = istft.forward_native + expected = run() + istft.forward = fused_forward + actual = run() + torch.testing.assert_close(actual, expected, atol=5e-7, rtol=1e-5) + # Alternate the order to expose drift when the decoder-level gain is small. + results = {"native": [], "fused": []} + for order in (("native", "fused"), ("fused", "native"), ("native", "fused")): + for name in order: + istft.forward = istft.forward_native if name == "native" else fused_forward + results[name].append(measure(run, args.warmup, args.decoder_iterations)) + print( + json.dumps( + { + "scope": "pretrained_decoder", + "model": model, + "revision": revision, + "backend": "sdpa", + "dtype": "bfloat16", + "latent_chunks": [4, 25, 25, 5], + "max_abs_error": (actual - expected).abs().max().item(), + "results": results, + } + ), + flush=True, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--compile", action="store_true") + parser.add_argument("--decoder", action="store_true") + parser.add_argument("--decoder-iterations", type=int, default=30) + args = parser.parse_args() + if args.warmup < 0 or min(args.iterations, args.decoder_iterations) < 1: + parser.error("warmup must be nonnegative and iteration counts must be positive") + torch.manual_seed(0) + print( + json.dumps( + { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "warmup": args.warmup, + "iterations": args.iterations, + "decoder_iterations": args.decoder_iterations, + } + ) + ) + with torch.inference_mode(): + benchmark_istft(args) + if args.decoder: + benchmark_decoder(args) + + +if __name__ == "__main__": + main() diff --git a/tests/model_executor/models/ming_tts/test_fused_istft.py b/tests/model_executor/models/ming_tts/test_fused_istft.py new file mode 100644 index 000000000..fa7d072d7 --- /dev/null +++ b/tests/model_executor/models/ming_tts/test_fused_istft.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.model_executor.models.common.ming.audio_dsp import ISTFT + +pytestmark = [pytest.mark.core_model, pytest.mark.tts] +cuda = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None or not HAS_TRITON, + reason="NVIDIA CUDA and Triton required", +) + + +def _assert_outputs_close(actual, expected): + for value, reference in zip(actual, expected, strict=True): + if reference is None: + assert value is None + else: + assert value.dtype == reference.dtype + torch.testing.assert_close(value, reference, atol=5e-7, rtol=1e-5) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("hop", [4, 320, 882]) +@pytest.mark.parametrize("window_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("batch,frames", [(1, 2), (1, 25), (3, 7), (2, 100)]) +@pytest.mark.parametrize("mode", ["offline", "first", "middle", "last", "first_last"]) +@torch.inference_mode() +def test_fused_istft_matches_native(hop, window_dtype, batch, frames, mode, monkeypatch): + torch.manual_seed(0) + module = ISTFT(4 * hop, hop, 4 * hop).to(device="cuda", dtype=window_dtype) + # Slice a larger spectrum so irfft must handle non-contiguous inputs. + spec = torch.randn(batch, 2 * hop + 1, 2 * frames, device="cuda", dtype=torch.complex64)[..., ::2] + kwargs = {"streaming": mode != "offline", "last_chunk": mode in ("last", "first_last")} + if mode in ("middle", "last"): + _, audio, window = module.forward_native(spec, streaming=True) + kwargs.update(audio_buffer=audio, window_buffer=window) + saved_audio, saved_window = audio.clone(), window.clone() + expected = module.forward_native(spec, **kwargs) + + def no_fallback(*args, **kwargs): + pytest.fail("supported input fell back to native ISTFT") + + monkeypatch.setattr(module, "forward_native", no_fallback) + _assert_outputs_close(module(spec, **kwargs), expected) + if mode in ("middle", "last"): + torch.testing.assert_close(audio, saved_audio, rtol=0, atol=0) + torch.testing.assert_close(window, saved_window, rtol=0, atol=0) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("hop", [320, 882]) +@pytest.mark.parametrize("window_dtype", [torch.float32, torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_fused_istft_streaming_state(hop, window_dtype): + torch.manual_seed(1) + module = ISTFT(4 * hop, hop, 4 * hop).to(device="cuda", dtype=window_dtype) + spec = torch.randn(2, 2 * hop + 1, 40, device="cuda", dtype=torch.complex64) + # Reusing the same module for a new request must not reuse the old tails. + for _ in range(2): + actual_state = reference_state = (None, None) + actual_chunks, reference_chunks = [], [] + for start, end in [(0, 4), (4, 29), (29, 31), (31, 40)]: + chunk = spec[..., start:end] + actual = module(chunk, *actual_state, streaming=True, last_chunk=end == 40) + expected = module.forward_native(chunk, *reference_state, streaming=True, last_chunk=end == 40) + _assert_outputs_close(actual, expected) + actual_state, reference_state = actual[1:], expected[1:] + actual_chunks.append(actual[0]) + reference_chunks.append(expected[0]) + output = torch.cat(actual_chunks, dim=-1) + torch.testing.assert_close(output, torch.cat(reference_chunks, dim=-1), atol=5e-7, rtol=1e-5) + if window_dtype == torch.float32: + # Low-precision window envelopes already round differently between + # streaming and offline native execution; preserve each contract. + torch.testing.assert_close(output, module.forward_native(spec)[0], atol=5e-7, rtol=1e-5) + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize("invalid", [False, True]) +@torch.inference_mode() +def test_fused_istft_silence_and_invalid_envelope(invalid): + module = ISTFT(1280, 320, 1280).cuda() + spec = torch.zeros(2, 641, 25, device="cuda", dtype=torch.complex64) + if invalid: + module.window.zero_() + for forward in (module.forward_native, module.forward): + with pytest.raises(RuntimeError, match="window envelope underflowed"): + forward(spec) + else: + _assert_outputs_close(module(spec), module.forward_native(spec)) + assert torch.count_nonzero(module(spec)[0]) == 0 + + +@pytest.mark.parametrize( + "device", + [pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=[pytest.mark.cuda, cuda])], +) +@pytest.mark.parametrize("padding", ["same", "center"]) +def test_istft_autograd_fallback(padding, device, monkeypatch): + module = ISTFT(32, 8, 32, padding=padding).to(device) + spec = torch.randn(2, 17, 7, device=device, dtype=torch.complex64, requires_grad=True) + native = module.forward_native + calls = [] + + def tracked(*args, **kwargs): + calls.append(True) + return native(*args, **kwargs) + + monkeypatch.setattr(module, "forward_native", tracked) + result = module(spec) + assert calls == [True] + reference = native(spec) + if padding == "same": + _assert_outputs_close(result, reference) + result = result[0] + else: + torch.testing.assert_close(result, reference) + result.square().sum().backward() + assert spec.grad is not None and torch.isfinite(spec.grad).all() + + +@cuda +@pytest.mark.cuda +@pytest.mark.parametrize( + "n_fft,hop,frames,dtype", [(32, 8, 1, torch.complex64), (30, 10, 7, torch.complex64), (32, 8, 7, torch.complex128)] +) +@torch.inference_mode() +def test_istft_unsupported_cuda_falls_back(n_fft, hop, frames, dtype, monkeypatch): + module = ISTFT(n_fft, hop, n_fft).cuda() + spec = torch.randn(1, n_fft // 2 + 1, frames, device="cuda", dtype=dtype) + native = module.forward_native + calls = [] + + def tracked(*args, **kwargs): + calls.append(True) + return native(*args, **kwargs) + + monkeypatch.setattr(module, "forward_native", tracked) + _assert_outputs_close(module(spec), native(spec)) + assert calls == [True] diff --git a/vllm_omni/model_executor/models/common/ming/audio_dsp.py b/vllm_omni/model_executor/models/common/ming/audio_dsp.py index be8cd489c..d32415352 100644 --- a/vllm_omni/model_executor/models/common/ming/audio_dsp.py +++ b/vllm_omni/model_executor/models/common/ming/audio_dsp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py import torch @@ -66,6 +66,26 @@ class ISTFT(nn.Module): window_buffer: torch.Tensor | None = None, streaming: bool = False, last_chunk: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if self.padding == "same" and spec.is_cuda and not torch.is_grad_enabled(): + from .fused_istft import fused_istft, fused_istft_supported + + if fused_istft_supported( + spec, self.window, self.n_fft, self.hop_length, audio_buffer, window_buffer, streaming + ): + frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + return fused_istft( + frames, self.window, self.hop_length, audio_buffer, window_buffer, streaming, last_chunk + ) + return self.forward_native(spec, audio_buffer, window_buffer, streaming, last_chunk) + + def forward_native( + self, + spec: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: """ Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. diff --git a/vllm_omni/model_executor/models/common/ming/fused_istft.py b/vllm_omni/model_executor/models/common/ming/fused_istft.py new file mode 100644 index 000000000..95da3ff70 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/fused_istft.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Post-FFT ISTFT fusion for Ming's four-way overlapping audio frames. + +Each lane owns one output sample and gathers its contributing frames. Audio +and window tails remain unnormalized so the next streaming chunk can add them. +The low-precision window rounding follows PyTorch's square, fold, and add +boundaries; the inverse FFT and the envelope-underflow check remain unchanged. +""" + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _istft_kernel( + frames, + window, + previous_audio, + previous_window, + output, + envelope_output, + next_audio, + next_window, + frame_stride_b, + frame_stride_w, + frame_stride_t, + audio_stride_b, + audio_stride_t, + window_stride_t, + num_frames, + raw_size, + output_start, + output_size, + hop_size: tl.constexpr, + streaming: tl.constexpr, + has_buffer: tl.constexpr, + block_size: tl.constexpr, + ): + batch = tl.program_id(0) + sample = tl.program_id(1) * block_size + tl.arange(0, block_size) + valid = sample < raw_size + last_frame = sample // hop_size + audio = tl.full((block_size,), 0, tl.float32) + envelope = tl.full((block_size,), 0, tl.float32) + window_dtype = window.dtype.element_ty + + # Match CUDA col2im's ascending frame order. Disable FMA at launch to + # preserve the separately rounded window multiplication. + for overlap in tl.static_range(3, -1, -1): + frame = last_frame - overlap + window_index = sample % hop_size + overlap * hop_size + mask = valid & (frame >= 0) & (frame < num_frames) + weight = tl.load(window + window_index, mask=mask, other=0).to(tl.float32) + value = tl.load( + frames + batch * frame_stride_b + window_index * frame_stride_w + frame * frame_stride_t, + mask=mask, + other=0, + ) + audio += value * weight + envelope += (weight * weight).to(window_dtype).to(tl.float32) + + # F.fold accumulates in FP32, then stores in the window's dtype. + envelope = envelope.to(window_dtype).to(tl.float32) + tail_size: tl.constexpr = 3 * hop_size + if has_buffer: + carry_mask = valid & (sample < tail_size) + carry_audio = tl.load( + previous_audio + batch * audio_stride_b + sample * audio_stride_t, mask=carry_mask, other=0 + ) + carry_window = tl.load(previous_window + sample * window_stride_t, mask=carry_mask, other=0) + audio += carry_audio + envelope = (envelope + carry_window.to(tl.float32)).to(window_dtype).to(tl.float32) + + output_index = sample - output_start + emit = valid & (output_index >= 0) & (output_index < output_size) + tl.store(output + batch * output_size + output_index, audio / envelope, mask=emit) + if batch == 0: + tl.store(envelope_output + output_index, envelope, mask=emit) + + if streaming: + tail_index = sample - (raw_size - tail_size) + tail_mask = valid & (tail_index >= 0) + tl.store(next_audio + batch * tail_size + tail_index, audio, mask=tail_mask) + if batch == 0: + tl.store(next_window + tail_index, envelope, mask=tail_mask) + + +def fused_istft_supported( + spec: torch.Tensor, + window: torch.Tensor, + n_fft: int, + hop: int, + audio_buffer: torch.Tensor | None, + window_buffer: torch.Tensor | None, + streaming: bool, +) -> bool: + if not ( + HAS_TRITON + and spec.is_cuda + and torch.version.hip is None + and spec.dtype == torch.complex64 + and spec.ndim == 3 + and spec.shape[0] > 0 + and spec.shape[1] == n_fft // 2 + 1 + and spec.shape[2] >= 2 + and hop > 0 + and n_fft == 4 * hop + and window.shape == (n_fft,) + and window.is_contiguous() + and window.device == spec.device + and window.dtype in (torch.float32, torch.bfloat16, torch.float16) + ): + return False + if not streaming or (audio_buffer is None and window_buffer is None): + return True + return ( + audio_buffer is not None + and window_buffer is not None + and audio_buffer.shape == (spec.shape[0], 3 * hop) + and window_buffer.shape == (1, 3 * hop) + and audio_buffer.device == window_buffer.device == spec.device + and audio_buffer.dtype == torch.float32 + and window_buffer.dtype == window.dtype + ) + + +def fused_istft( + frames: torch.Tensor, + window: torch.Tensor, + hop: int, + audio_buffer: torch.Tensor | None, + window_buffer: torch.Tensor | None, + streaming: bool, + last_chunk: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Fuse supported post-FFT inputs; dispatch is checked by ``ISTFT``.""" + batch, win_length, num_frames = frames.shape + tail_size = win_length - hop + pad = tail_size // 2 + raw_size = (num_frames - 1) * hop + win_length + has_buffer = streaming and audio_buffer is not None + start = 0 if has_buffer else pad + end = raw_size - (tail_size if streaming and not last_chunk else pad) + output_size = end - start + output = frames.new_empty((batch, output_size)) + envelope = window.new_empty((output_size,)) + next_audio = frames.new_empty((batch, tail_size)) if streaming else audio_buffer + next_window = window.new_empty((1, tail_size)) if streaming else window_buffer + + _istft_kernel[(batch, triton.cdiv(raw_size, 256))]( + frames, + window, + audio_buffer, + window_buffer, + output, + envelope, + next_audio, + next_window, + *frames.stride(), + audio_buffer.stride(0) if streaming and audio_buffer is not None else 0, + audio_buffer.stride(1) if streaming and audio_buffer is not None else 0, + window_buffer.stride(1) if streaming and window_buffer is not None else 0, + num_frames, + raw_size, + start, + output_size, + hop_size=hop, + streaming=streaming, + has_buffer=has_buffer, + block_size=256, + enable_fp_fusion=False, + ) + if not (envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + return output, next_audio, next_window