From a41a877c5e5d831671e4d50455925a2dbeb37548 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Fri, 13 Nov 2026 23:49:36 -0800 Subject: [PATCH] Phonon: run the Mimi decoder's transposed convolutions as convolutions on XNNPACK (#25004) Summary: XNNPACK's deconvolution is the most expensive op in Phonon's Mimi decoder: about 45% of each frame on x86 and 38% on arm64. Every transposed convolution in Mimi has kernel = 2 * stride, so each one is exactly a kernel-2 convolution with stride * out_channels outputs, one per output phase, whose phases are then interleaved into the output steps. The XNNPACK export now applies that rewrite (conv_transpose1d_as_conv1d) to the three decoder transposed convolutions, and the existing depthwise_conv_transpose_as_mul to the upsampler. The output is the same up to float rounding: the exported 31M program writes audio within 1 LSB (16-bit PCM) of the program without the rewrite, with the same length and frames. Running the transposed convolutions as 2-D convolutions with the sequence on the width axis (conv_transpose1d_as_2d) does not help: Mimi stays at 8.2 ms per frame, so the gain comes from the convolution kernel rather than the tiling axis. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: shoumikhin, JakeStevens Differential Revision: D131605676 --- examples/models/phonon/README.md | 4 +- examples/models/phonon/export_phonon.py | 1 + .../models/phonon/source_transformations.py | 57 ++++++++++++++++++- examples/models/phonon/tests/test_phonon.py | 33 +++++++++-- 4 files changed, 88 insertions(+), 7 deletions(-) diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index be74338d754..7f471832a30 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -57,7 +57,9 @@ cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --s ``` XNNPACK delegates the linear and convolution ops; the flow LM's attention runs -as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. The export also +rewrites the Mimi decoder's transposed convolutions as convolutions, which +XNNPACK runs faster. `--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow LM's transformer and flow net with the scheme of the backend; the codec stays in diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index 7784fc65d33..2bae25d8f47 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -82,6 +82,7 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") return [st.use_webgpu_compat_ops] return { + "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d], "mlx": [st.decompose_elu], "vulkan": [ st.decompose_elu, diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py index 69a2c228a54..4c6a1dcdb9b 100644 --- a/examples/models/phonon/source_transformations.py +++ b/examples/models/phonon/source_transformations.py @@ -4,7 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Function-preserving module swaps for backends that lack an op. +"""Function-preserving module swaps for backends that lack an op or run it +slowly. The model in model.py uses plain PyTorch ops. Each swap below computes the same function from more widely supported ops; export_phonon.py selects them @@ -13,10 +14,15 @@ decompose_elu MLX and Vulkan have no aten.elu kernel; without it the Mimi decoder splits at every ELU. conv_transpose1d_as_2d Vulkan runs transposed convolution on 4-D inputs only. + conv_transpose1d_as_conv1d + XNNPACK's deconvolution is slow; this runs the Mimi + decoder's transposed convolutions (kernel = + 2 * stride) as a convolution and an interleave. depthwise_conv_transpose_as_mul Vulkan computes the Mimi upsampler (depthwise, kernel = 2 * stride) incorrectly; this runs it as a - multiply and an overlap-add. + multiply and an overlap-add. XNNPACK also runs it + faster this way. use_webgpu_compat_ops WebGPU runtimes built before these runtime fixes (off by default; --webgpu-compat-ops): * native_layer_norm resize hook (the flow LM's @@ -93,6 +99,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y.squeeze(2) +class ConvTranspose1dAsConv1d(nn.Module): + """ConvTranspose1d with kernel_size == 2 * stride as a kernel-2 Conv1d + with stride * out_channels outputs (one per output phase), interleaved + into the output steps.""" + + def __init__(self, convtr: nn.ConvTranspose1d): + super().__init__() + (self.stride,) = convtr.stride + s = self.stride + w = convtr.weight.detach() # [cin, cout, 2s] + cin, cout, _ = w.shape + # Phase j of output step p reads input p - 1 through tap j + s and + # input p through tap j. + taps = torch.stack([w[..., s:], w[..., :s]], -1) # [cin, cout, s, 2] + self.weight = nn.Parameter( + taps.permute(1, 2, 0, 3).reshape(cout * s, cin, 2).contiguous(), + requires_grad=False, + ) + self.bias = convtr.bias + self.out_channels = cout + + def forward(self, x: torch.Tensor) -> torch.Tensor: + z = F.conv1d(x, self.weight, padding=1) # [b, cout * s, n + 1] + b, _, n = z.shape + s = self.stride + y = z.view(b, self.out_channels, s, n).transpose(2, 3) + y = y.reshape(b, self.out_channels, n * s) + return y if self.bias is None else y + self.bias.view(1, -1, 1) + + class ComputeIn(nn.Module): """Runs `module` with its weights in `dtype` on a float32 input and returns float32. Other attributes (bias, kernel_size, ...) are read from @@ -171,6 +207,23 @@ def conv_transpose1d_as_2d(programs: dict[str, nn.Module]) -> int: return count +def conv_transpose1d_as_conv1d(programs: dict[str, nn.Module]) -> int: + count = 0 + for program in programs.values(): + for module in program.modules(): + if not isinstance(module, StreamingConvTranspose1d): + continue + convtr = module.convtr + if ( + isinstance(convtr, nn.ConvTranspose1d) + and convtr.groups == 1 + and convtr.kernel_size[0] == 2 * convtr.stride[0] + ): + module.convtr = ConvTranspose1dAsConv1d(convtr) + count += 1 + return count + + def depthwise_conv_transpose_as_mul(programs: dict[str, nn.Module]) -> int: count = 0 for program in programs.values(): diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index ce952bfec70..63b79dee939 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -43,6 +43,7 @@ ) from executorch.examples.models.phonon.source_transformations import ( compute_weights_in, + ConvTranspose1dAsConv1d, use_webgpu_compat_ops, ) from executorch.examples.models.phonon.tests.reference import ( @@ -236,19 +237,43 @@ def test_webgpu_compat_ops_preserve_outputs(self): for a, b in zip(outputs, programs[name](*inputs[name])): self.assertLess(max_diff(a, b), 1e-4) - def test_vulkan_swaps_preserve_outputs(self): + def _check_swaps(self, backend: str, counts: list[int]): inputs = example_inputs(TINY) programs = build_programs(self.model) with torch.no_grad(): before = programs["mimi_decode"](*inputs["mimi_decode"]) - # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. - counts = [t(programs) for t in source_transforms_for("vulkan", False)] - self.assertEqual(counts, [7, 1, 2]) + self.assertEqual( + [t(programs) for t in source_transforms_for(backend, False)], counts + ) with torch.no_grad(): after = programs["mimi_decode"](*inputs["mimi_decode"]) for a, b in zip(before, after): self.assertLess(max_diff(a, b), 1e-4) + def test_vulkan_swaps_preserve_outputs(self): + # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("vulkan", [7, 1, 2]) + + def test_xnnpack_swaps_preserve_outputs(self): + # The depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("xnnpack", [1, 2]) + + def test_conv_transpose_as_conv_matches(self): + torch.manual_seed(0) + for stride, steps, bias in ( + (1, 1, True), + (4, 16, False), + (5, 1, True), + (6, 16, True), + ): + convtr = torch.nn.ConvTranspose1d( + 6, 3, 2 * stride, stride=stride, bias=bias + ) + x = torch.randn(1, 6, steps) + with torch.no_grad(): + got = ConvTranspose1dAsConv1d(convtr)(x) + self.assertLess(max_diff(got, convtr(x)), 1e-5, stride) + def test_compute_weights_in_keeps_float32_io(self): inputs = example_inputs(TINY) programs = build_programs(self.model)