Phonon: keep the Mimi decoder state on the GPU for CUDA (#25021)

examples/models/phonon/README.md

@@ -61,6 +61,9 @@ 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.
+For CUDA, the Mimi decoder also keeps its streaming state on the GPU between
+frames: its `mimi_decode` takes `(latent, reset)` instead of the latent and
+the state (see `phonon.json`).
 
 `--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

@@ -66,6 +66,13 @@
 }
 
 
+def mimi_state_in_program(backend: str) -> bool:
+    """Whether mimi_decode keeps its streaming state in the program
+    (StatefulMimiDecode). On CUDA, state passed as inputs and outputs is copied
+    between host and device on every frame."""
+    return backend == "cuda"
+
+
 def attention_for(backend: str) -> str:
     """KV cache / attention implementation of the flow LM per backend (see
     FlowLMStep): the ExecuTorch custom op where a kernel exists for it, plain
@@ -246,8 +253,9 @@ def lower_phonon(
     qlinear, dtype = quantize_scheme(backend, quantize, dtype)
     packing = "tile_packed_to_4d" if backend == "cuda" and qlinear == "4w" else None
     # Before the module swaps, which change the model's modules in place.
-    metadata = program_metadata(model)
-    programs = build_programs(model, attention_for(backend))
+    stateful_mimi = mimi_state_in_program(backend)
+    metadata = program_metadata(model, stateful_mimi)
+    programs = build_programs(model, attention_for(backend), stateful_mimi)
     for transform in source_transforms_for(backend, webgpu_compat_ops):
         print(f"{transform.__name__}: {transform(programs)}")
     if dtype == "bf16":
@@ -275,7 +283,7 @@ def lower_phonon(
                 skip_incompatible_shapes=True,
             )
 
-    inputs, shapes = example_inputs(cfg), dynamic_shapes(cfg)
+    inputs, shapes = example_inputs(cfg, stateful_mimi), dynamic_shapes(cfg)
     if packing:
         inputs["flow_lm"] = tuple(t.to("cuda") for t in inputs["flow_lm"])
     partitioners, passes, decompositions = partitioners_for(
@@ -356,7 +364,7 @@ def export_phonon(
         "constants": {k: v for k, v in metadata.items() if not torch.is_tensor(v)},
         "constant_tensors": {"bos_embeds": [1, 1, cfg.dim]},
         "dynamic": {"S": [1, cfg.max_prefill]},
-        "methods": io_contract(cfg),
+        "methods": io_contract(cfg, mimi_state_in_program(backend)),
         "mimi_state": [[n, list(s)] for n, s in mimi_state_spec(cfg)],
         "delegation": delegation,
     }

examples/models/phonon/model.py

@@ -810,7 +810,8 @@ class MimiDecodeStep(nn.Module):
     the SEANet decoder. Every piece of streaming state (conv tails, transposed
     conv partial sums, the window's keys and values, which window slots are
     filled, and the RoPE offset) is an explicit input and output, so no backend
-    has to keep mutable state alive between calls.
+    has to keep mutable state alive between calls. StatefulMimiDecode keeps it
+    in the program instead.
     """
 
     def __init__(self, model: Phonon):
@@ -915,7 +916,41 @@ def forward(self, latent: torch.Tensor, *state: torch.Tensor):
         )
 
 
-def example_inputs(cfg: PhononConfig) -> dict[str, Tuple[torch.Tensor, ...]]:
+class MimiState(nn.Module):
+    """Buffers for the state of mimi_state_spec, in its order."""
+
+    def __init__(self, cfg: PhononConfig):
+        super().__init__()
+        for name, shape in mimi_state_spec(cfg):
+            self.register_buffer(name, torch.zeros(shape), persistent=False)
+
+
+class StatefulMimiDecode(nn.Module):
+    """forward(latent f32[1, 1, latent_dim], reset f32[1]) -> pcm f32[1, 1, frame]
+
+    MimiDecodeStep with its streaming state in mutable buffers of the program,
+    so a backend that runs on another device (CUDA) keeps it there instead of
+    copying about 2 MB in and out per frame. reset = 1 starts an utterance from
+    the zero state, reset = 0 continues it.
+    """
+
+    def __init__(self, step: MimiDecodeStep):
+        super().__init__()
+        self.cfg = step.cfg
+        self.step = step
+        self.state = MimiState(step.cfg)
+
+    def forward(self, latent: torch.Tensor, reset: torch.Tensor) -> torch.Tensor:
+        keep = 1.0 - reset
+        pcm, *new_state = self.step(latent, *(s * keep for s in self.state.buffers()))
+        for buf, value in zip(self.state.buffers(), new_state):
+            buf.copy_(value)
+        return pcm
+
+
+def example_inputs(
+    cfg: PhononConfig, mimi_state_in_program: bool = False
+) -> dict[str, Tuple[torch.Tensor, ...]]:
     g = torch.Generator().manual_seed(0)
     seq = 5
     return {
@@ -928,7 +963,7 @@ def example_inputs(cfg: PhononConfig) -> dict[str, Tuple[torch.Tensor, ...]]:
         ),
         "mimi_decode": (
             torch.randn(1, 1, cfg.latent_dim, generator=g),
-            *mimi_zero_state(cfg),
+            *([torch.ones(1)] if mimi_state_in_program else mimi_zero_state(cfg)),
         ),
     }
 
@@ -944,23 +979,28 @@ def dynamic_shapes(cfg: PhononConfig) -> dict[str, Any]:
 
 
 def build_programs(
-    model: Phonon, attention: str = "custom_sdpa"
+    model: Phonon, attention: str = "custom_sdpa", mimi_state_in_program: bool = False
 ) -> dict[str, nn.Module]:
+    mimi = MimiDecodeStep(model)
     return {
         "text_embed": TextEmbedding(model).eval(),
         "voice_embed": VoiceEmbedding(model).eval(),
         "flow_lm": FlowLMStep(model, attention).eval(),
-        "mimi_decode": MimiDecodeStep(model).eval(),
+        "mimi_decode": (
+            StatefulMimiDecode(mimi) if mimi_state_in_program else mimi
+        ).eval(),
     }
 
 
-def program_metadata(model: Phonon) -> dict[str, Any]:
+def program_metadata(
+    model: Phonon, mimi_state_in_program: bool = False
+) -> dict[str, Any]:
     """Constant methods stored in the .pte next to the four programs."""
     cfg = model.cfg
     with torch.no_grad():
         bos = model.flow_lm.bos_embeds().contiguous()
     return {
-        "format_version": 1,
+        "format_version": 2,
         "sample_rate": cfg.sample_rate,
         "frame_samples": cfg.frame_samples,
         "frame_rate": cfg.frame_rate,
@@ -971,13 +1011,17 @@ def program_metadata(model: Phonon) -> dict[str, Any]:
         "eos_threshold": cfg.eos_threshold,
         "temperature": cfg.temperature,
         "num_mimi_state": len(mimi_state_spec(cfg)),
+        "mimi_state_in_program": int(mimi_state_in_program),
         "bos_embeds": bos,
     }
 
 
-def io_contract(cfg: PhononConfig) -> dict[str, Any]:
+def io_contract(
+    cfg: PhononConfig, mimi_state_in_program: bool = False
+) -> dict[str, Any]:
     """Inputs and outputs of every program; "S" is the dynamic sequence length
-    (1 <= S <= max_prefill). Identical for every backend."""
+    (1 <= S <= max_prefill). mimi_decode takes and returns its state unless the
+    program keeps it (StatefulMimiDecode)."""
     d, ld = cfg.dim, cfg.latent_dim
     state = [
         {"name": name, "shape": list(shape), "dtype": "float32"}
@@ -1009,7 +1053,9 @@ def t(name, shape, dtype="float32"):
             ],
         },
         "mimi_decode": {
-            "inputs": [t("latent", [1, 1, ld])] + state,
-            "outputs": [t("pcm", [1, 1, cfg.frame_samples])] + state,
+            "inputs": [t("latent", [1, 1, ld])]
+            + ([t("reset", [1])] if mimi_state_in_program else state),
+            "outputs": [t("pcm", [1, 1, cfg.frame_samples])]
+            + ([] if mimi_state_in_program else state),
         },
     }

examples/models/phonon/run_phonon.py

@@ -172,8 +172,9 @@ def __init__(self, programs, tokenize: Callable[[str], List[int]]):
         self.run = programs
         self.tokenize = tokenize
         c = programs.constant
-        if c("format_version") != 1:
+        if c("format_version") != 2:
             raise ValueError("Unsupported Phonon program; re-export it")
+        self.mimi_state_in_program = bool(c("mimi_state_in_program"))
         self.sample_rate = c("sample_rate")
         self.frame_samples = c("frame_samples")
         self.frame_rate = c("frame_rate")
@@ -213,7 +214,10 @@ def noise_fn(k):
         if voice_len > self.max_prefill:
             raise ValueError("Voice prompt is longer than max_prefill")
         self._flow(voice_embeds, 0, torch.zeros(1, self.latent_dim))
-        mimi_state = self.run.mimi_zero_state()
+        # The inputs of mimi_decode after the latent: the state, or reset = 1
+        # then 0 when the program keeps the state.
+        stateful = self.mimi_state_in_program
+        mimi_state = [torch.ones(1)] if stateful else self.run.mimi_zero_state()
         prepared, _ = prepare_text(text)
         chunks = split_sentences(
             prepared, lambda s: len(self.tokenize(s)), MAX_CHUNK_TOKENS
@@ -237,7 +241,7 @@ def noise_fn(k):
                 pos, step, frames = pos + 1, step + 1, frames + 1
                 out = self.run("mimi_decode", [latent.reshape(1, 1, -1), *mimi_state])
                 pcm.append(out[0].reshape(-1).numpy().copy())
-                mimi_state = out[1:]
+                mimi_state = [torch.zeros(1)] if stateful else out[1:]
                 info["latents"].append(latent.reshape(-1).numpy().copy())
                 info["eos_logits"].append(eos_logit)
                 if eos_logit > self.eos_threshold and countdown is None:

examples/models/phonon/runner.cpp

@@ -28,7 +28,7 @@ using executorch::runtime::EValue;
 
 // The official runtime's sentence chunk size.
 constexpr size_t kMaxChunkTokens = 50;
-constexpr int kFormatVersion = 1;
+constexpr int kFormatVersion = 2;
 
 std::vector<EValue> execute(
     Module& model,
@@ -293,6 +293,7 @@ Runner::Runner(
   frame_rate_ = real("frame_rate");
   eos_threshold_ = real("eos_threshold");
   temperature_ = real("temperature");
+  mimi_state_in_program_ = integer("mimi_state_in_program") != 0;
   bos_embeds_ = copy_floats(execute(model_, "bos_embeds").at(0));
   if (static_cast<int64_t>(bos_embeds_.size()) != dim_) {
     throw std::runtime_error("bos_embeds does not match dim");
@@ -308,6 +309,12 @@ Runner::Runner(
   if (!meta.ok()) {
     throw std::runtime_error("Could not read mimi_decode metadata");
   }
+  if (mimi_state_in_program_) {
+    if (meta->num_inputs() != 2) {
+      throw std::runtime_error("mimi_decode must take (latent, reset)");
+    }
+    return;
+  }
   for (size_t i = 1; i < meta->num_inputs(); ++i) {
     auto info = meta->input_tensor_meta(i);
     if (!info.ok() || info->scalar_type() != ScalarType::Float) {
@@ -434,6 +441,8 @@ Synthesis Runner::synthesize(
     }
     state.emplace_back(numel, 0.0f);
   }
+  // For a program that keeps the Mimi state: 1 on the first frame.
+  float mimi_reset = 1.0f;
 
   Synthesis result;
   const auto count_tokens = [this](const std::string& s) {
@@ -471,9 +480,12 @@ Synthesis Runner::synthesize(
 
       // The tensors must outlive the EValues that point to them.
       std::vector<executorch::extension::TensorPtr> tensors;
-      tensors.reserve(state.size() + 1);
+      tensors.reserve(state.size() + 2);
       tensors.push_back(from_blob(
           step.latent.data(), {1, 1, static_cast<SizesType>(latent_dim_)}));
+      if (mimi_state_in_program_) {
+        tensors.push_back(from_blob(&mimi_reset, {1}));
+      }
       for (size_t i = 0; i < state.size(); ++i) {
         tensors.push_back(from_blob(state[i].data(), mimi_state_sizes_[i]));
       }
@@ -491,6 +503,7 @@ Synthesis Runner::synthesize(
         std::memcpy(
             state[i].data(), float_data(out.at(i + 1)), state[i].size() * 4);
       }
+      mimi_reset = 0.0f;
 
       if (step.eos_logit > eos_threshold_ && countdown < 0) {
         countdown = frames_after_eos;

examples/models/phonon/runner.h

@@ -105,7 +105,10 @@ class Runner {
       max_prefill_;
   double frame_rate_, eos_threshold_, temperature_;
   std::vector<float> bos_embeds_, noise_;
+  // Shapes of the Mimi state passed in and out of mimi_decode; empty when the
+  // program keeps the state (mimi_state_in_program) and takes a reset flag.
   std::vector<std::vector<executorch::aten::SizesType>> mimi_state_sizes_;
+  bool mimi_state_in_program_ = false;
 };
 
 } // namespace phonon

examples/models/phonon/tests/test_phonon.py

@@ -16,6 +16,7 @@
 import tempfile
 import unittest
 
+import numpy as np
 import torch
 
 from executorch.examples.models.phonon.export_phonon import (
@@ -52,6 +53,8 @@
     FlowLMReference,
     MimiReference,
 )
+from executorch.exir import to_edge_transform_and_lower
+from torch.export import export
 
 TINY = PhononConfig(
     dim=64,
@@ -212,6 +215,20 @@ def test_mimi_matches_reference(self):
             self.assertEqual(tuple(pcm.shape), (1, 1, cfg.frame_samples))
             self.assertLess(max_diff(pcm, ref.decode(latent)), 1e-4)
 
+    def test_stateful_mimi_matches_state_io(self):
+        mimi = build_programs(self.model)["mimi_decode"]
+        stateful = build_programs(self.model, mimi_state_in_program=True)["mimi_decode"]
+        g = torch.Generator().manual_seed(3)
+        # Two utterances: the second reset must clear the first one's state.
+        for _ in range(2):
+            state = mimi_zero_state(TINY)
+            for k in range(4):
+                latent = torch.randn(1, 1, TINY.latent_dim, generator=g)
+                with torch.no_grad():
+                    want, *state = mimi(latent, *state)
+                    got = stateful(latent, torch.tensor([1.0 if k == 0 else 0.0]))
+                self.assertLess(max_diff(got, want), 1e-6)
+
     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")
@@ -384,6 +401,32 @@ def test_xnnpack_program_runs_from_file(self):
         for a, b in zip(got, want):
             self.assertLess(max_diff(a, b), 1e-4)
 
+    def test_stateful_mimi_program_keeps_state_between_calls(self):
+        try:
+            from executorch.extension.pybindings.portable_lib import (
+                _load_for_executorch,
+            )
+        except ImportError:
+            self.skipTest("ExecuTorch pybindings are not available")
+        model = random_model().to_half_split_rope()
+        stateful = build_programs(model, mimi_state_in_program=True)["mimi_decode"]
+        inputs = example_inputs(TINY, mimi_state_in_program=True)["mimi_decode"]
+        program = to_edge_transform_and_lower(
+            {"mimi_decode": export(stateful, inputs, strict=False)}
+        ).to_executorch()
+        g = torch.Generator().manual_seed(4)
+        with tempfile.TemporaryDirectory() as d:
+            with open(f"{d}/mimi.pte", "wb") as f:
+                program.write_to_file(f)
+            module = _load_for_executorch(f"{d}/mimi.pte")
+            for reset in (1.0, 0.0, 0.0, 1.0, 0.0):
+                args = (torch.randn(1, 1, TINY.latent_dim, generator=g),)
+                args += (torch.tensor([reset]),)
+                with torch.no_grad():
+                    want = stateful(*args)
+                got = module.run_method("mimi_decode", list(args))[0]
+                self.assertLess(max_diff(got, want), 1e-4)
+
     def test_bf16_only_on_supported_backends(self):
         with self.assertRaises(ValueError):
             lower_phonon(random_model().to_half_split_rope(), "vulkan", dtype="bf16")
@@ -467,10 +510,25 @@ def test_runner_on_eager_programs(self):
         self.assertEqual(pcm.size, info["frames"] * TINY.frame_samples)
         again, _ = runner.synthesize("hi there", speaker, temperature=0.0)
         self.assertTrue((pcm == again).all())
+        stateful = PhononRunner(
+            EagerPrograms(
+                build_programs(model, mimi_state_in_program=True),
+                program_metadata(model, mimi_state_in_program=True),
+            ),
+            lambda text: [1, 2, 3][: len(text.split())],
+        )
+        for _ in range(2):
+            got, _ = stateful.synthesize("hi there", speaker, temperature=0.0)
+            self.assertLess(np.abs(got - pcm).max(), 1e-5)
 
     def test_contract_is_json_serializable(self):
         json.dumps(io_contract(PhononConfig()))
         self.assertEqual(len(io_contract(PhononConfig())["mimi_decode"]["inputs"]), 16)
+        stateful = io_contract(PhononConfig(), mimi_state_in_program=True)
+        self.assertEqual(
+            [x["name"] for x in stateful["mimi_decode"]["inputs"]], ["latent", "reset"]
+        )
+        self.assertEqual(len(stateful["mimi_decode"]["outputs"]), 1)
 
 
 if __name__ == "__main__":