From 2d35fa129da1c63a3455f08ca8c2d9a34726f8c2 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Sat, 14 Nov 2026 10:37:44 -0800 Subject: [PATCH] Phonon: keep the Mimi decoder state on the GPU for CUDA (#25021) Summary: mimi_decode takes its streaming state (15 tensors, 2.1 MB, mostly the decoder transformer's keys and values) as inputs and returns it as outputs, so on CUDA every 80 ms frame copied it to the GPU and back, with a synchronization each way. That made mimi_decode take 4 ms in the runtime against 1.3 ms for the compiled graph. StatefulMimiDecode wraps MimiDecodeStep and keeps the state in mutable buffers of the program (MimiState): forward(latent, reset) -> pcm, where reset = 1 starts an utterance from the zero state. The CUDA export uses it (mimi_state_in_program in export_phonon.py); the other backends keep the state as inputs and outputs. The program records the choice in a new mimi_state_in_program constant and phonon.json, the C++ and Python runners pass (latent, reset) when it is set, and format_version is now 2, so the runners ask to re-export older programs. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: shoumikhin, Gasoonjia Differential Revision: D131675810 --- examples/models/phonon/README.md | 3 + examples/models/phonon/export_phonon.py | 16 +++-- examples/models/phonon/model.py | 68 +++++++++++++++++---- examples/models/phonon/run_phonon.py | 10 ++- examples/models/phonon/runner.cpp | 17 +++++- examples/models/phonon/runner.h | 3 + examples/models/phonon/tests/test_phonon.py | 58 ++++++++++++++++++ 7 files changed, 155 insertions(+), 20 deletions(-) diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index 35d4406422b..c7e1964177e 100644 --- a/examples/models/phonon/README.md +++ b/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 diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index e66c2849ce2..a45265d3d6f 100644 --- a/examples/models/phonon/export_phonon.py +++ b/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, } diff --git a/examples/models/phonon/model.py b/examples/models/phonon/model.py index dbc58b668f4..66cc7e61c8e 100644 --- a/examples/models/phonon/model.py +++ b/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), }, } diff --git a/examples/models/phonon/run_phonon.py b/examples/models/phonon/run_phonon.py index b35e00cbcd2..a8339ca13ef 100644 --- a/examples/models/phonon/run_phonon.py +++ b/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: diff --git a/examples/models/phonon/runner.cpp b/examples/models/phonon/runner.cpp index def47451416..2de2e1dbe54 100644 --- a/examples/models/phonon/runner.cpp +++ b/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 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(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 tensors; - tensors.reserve(state.size() + 1); + tensors.reserve(state.size() + 2); tensors.push_back(from_blob( step.latent.data(), {1, 1, static_cast(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; diff --git a/examples/models/phonon/runner.h b/examples/models/phonon/runner.h index 646696b91be..679762541f7 100644 --- a/examples/models/phonon/runner.h +++ b/examples/models/phonon/runner.h @@ -105,7 +105,10 @@ class Runner { max_prefill_; double frame_rate_, eos_threshold_, temperature_; std::vector 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> mimi_state_sizes_; + bool mimi_state_in_program_ = false; }; } // namespace phonon diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index a551b871ef6..9bc3d0de6fe 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/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__":