[Kernel] Fuse Ming streaming ISTFT overlap-add and normalization (#7338)

Signed-off-by: Yash Jain <yashjainjain1704@gmail.com>

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()

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]

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.

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