Phonon: run the Mimi transposed convolutions as convolutions on Vulkan, MLX and CUDA (#25020)

examples/models/phonon/README.md

@@ -57,9 +57,10 @@ 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`. The export also
-rewrites the Mimi decoder's transposed convolutions as convolutions, which
-XNNPACK runs faster.
+as the ExecuTorch custom op `llama::sdpa_with_kv_cache`.
+
+For XNNPACK, Vulkan, MLX and CUDA, the export rewrites the Mimi decoder's
+transposed convolutions as convolutions, which those backends run 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

examples/models/phonon/export_phonon.py

@@ -81,14 +81,12 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool):
         if backend != "webgpu":
             raise ValueError("--webgpu-compat-ops only applies to --backend webgpu")
         return [st.use_webgpu_compat_ops]
+    mimi_convs = [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d]
     return {
-        "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d],
-        "mlx": [st.decompose_elu],
-        "vulkan": [
-            st.decompose_elu,
-            st.depthwise_conv_transpose_as_mul,
-            st.conv_transpose1d_as_2d,
-        ],
+        "xnnpack": mimi_convs,
+        "cuda": mimi_convs,
+        "mlx": [st.decompose_elu, *mimi_convs],
+        "vulkan": [st.decompose_elu, *mimi_convs],
     }.get(backend, [])
 
 

examples/models/phonon/source_transformations.py

@@ -13,16 +13,19 @@
 
   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_2d  Older WebGPU runtimes have no 1-D transposed
+                          convolution (use_webgpu_compat_ops).
   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.
+                          Runs the Mimi decoder's transposed convolutions
+                          (kernel = 2 * stride) as a convolution and an
+                          interleave, which XNNPACK, Vulkan, MLX and CUDA run
+                          faster than their transposed convolution.
   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. XNNPACK also runs it
-                          faster this way.
+                          multiply and an overlap-add. XNNPACK, MLX and CUDA
+                          use it for the upsampler too, which
+                          conv_transpose1d_as_conv1d skips.
   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
@@ -171,7 +174,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         y = x.unsqueeze(3) * self.weight  # [b, c, n, 2s]
         head = y[..., :s].reshape(b, c, n * s)
         tail = y[..., s:].reshape(b, c, n * s)
-        zeros = torch.zeros(b, c, s, dtype=x.dtype)
+        zeros = x.new_zeros(b, c, s)
         return torch.cat([head, zeros], 2) + torch.cat([zeros, tail], 2)
 
 
@@ -262,6 +265,8 @@ def decompose_dynamic_layer_norm(programs: dict[str, nn.Module]) -> int:
     nn.ConvTranspose1d,
     nn.Embedding,
     ConvTranspose1dAs2d,
+    ConvTranspose1dAsConv1d,
+    DepthwiseConvTranspose1dAsMul,
 )
 
 

examples/models/phonon/tests/test_phonon.py

@@ -32,6 +32,7 @@
     Phonon,
     PhononConfig,
     program_metadata,
+    StreamingConvTranspose1d,
 )
 from executorch.examples.models.phonon.run_phonon import (
     EagerPrograms,
@@ -43,6 +44,7 @@
 )
 from executorch.examples.models.phonon.source_transformations import (
     compute_weights_in,
+    ComputeIn,
     ConvTranspose1dAsConv1d,
     use_webgpu_compat_ops,
 )
@@ -213,6 +215,8 @@ def test_mimi_matches_reference(self):
     def test_weight_names_are_unique_across_programs(self):
         # The CUDA backend shares a weight between methods when the names match.
         programs = build_programs(self.model, "sdpa")
+        for transform in source_transforms_for("cuda", False):
+            transform(programs)
         compute_weights_in(programs, torch.bfloat16)
         owner = {}
         for name, program in programs.items():
@@ -258,6 +262,13 @@ def test_xnnpack_swaps_preserve_outputs(self):
         # The depthwise upsampler, 2 decoder transposed convs.
         self._check_swaps("xnnpack", [1, 2])
 
+    def test_cuda_swaps_preserve_outputs(self):
+        self._check_swaps("cuda", [1, 2])
+
+    def test_mlx_swaps_preserve_outputs(self):
+        # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs.
+        self._check_swaps("mlx", [7, 1, 2])
+
     def test_conv_transpose_as_conv_matches(self):
         torch.manual_seed(0)
         for stride, steps, bias in (
@@ -275,13 +286,23 @@ def test_conv_transpose_as_conv_matches(self):
                 self.assertLess(max_diff(got, convtr(x)), 1e-5, stride)
 
     def test_compute_weights_in_keeps_float32_io(self):
+        # --dtype bf16 is for CUDA and MLX, which apply the Mimi swaps first.
         inputs = example_inputs(TINY)
         programs = build_programs(self.model)
         with torch.no_grad():
             before = programs["mimi_decode"](*inputs["mimi_decode"])
+        for transform in source_transforms_for("cuda", False):
+            transform(programs)
         self.assertGreater(compute_weights_in(programs, torch.bfloat16), 0)
+        mimi = programs["mimi_decode"]
+        convtrs = [
+            m.convtr for m in mimi.modules() if isinstance(m, StreamingConvTranspose1d)
+        ]
+        self.assertEqual(len(convtrs), 3)
+        for convtr in convtrs:
+            self.assertIsInstance(convtr, ComputeIn)
         with torch.no_grad():
-            after = programs["mimi_decode"](*inputs["mimi_decode"])
+            after = mimi(*inputs["mimi_decode"])
         for a, b in zip(before, after):
             self.assertEqual(b.dtype, torch.float32)
             self.assertLess(max_diff(a, b), 0.05 * (a.abs().max().item() + 1e-6))