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