Tuning CosyVoice3 in Rust: Where the Speed Actually Was

I spent two rounds of work optimising a stage that accounts for 4.4% of runtime, because my own profiler told me it was 17%.

That is the most useful thing I can tell you about this project, so it goes first. The rest — RTF 2.96 down to 0.697 on the same 132-word passage, a second implementation of a model I had already made fast once in PyTorch, this time in Rust on candle — is downstream of getting the instrument right.

RTF is synthesis time over audio duration. Under 1.0 is faster than real time.

Fix the instrument first

Metal dispatch is asynchronous. Encoding GPU work returns before the work runs, so a timer that reads the clock after encoding — but before anything needs the result — charges that stage’s time to whichever later stage happens to block on the queue.

Mine did that. Adding a device.synchronize() before recording each stage:

Amdahl’s law is not the lesson. The lesson is that I chose targets using a number my own instrumentation produced, it was wrong by a factor of four, and I never questioned it because it was my number.

Memory went worse. Peak RSS sat at exactly 3.59 GB across a sixfold range of input length. phys_footprint_peak returned exactly 13.00 GB for three configurations that should have differed enormously. Both had saturated — pinned allocations in one case, a high-water mark in the other — and neither could see the variable I cared about. I built a fix on top of the second reading before admitting it told me nothing.

Arithmetic found it. Each autoregressive lane carries its own KV cache: 2 * n_kv * capacity * head_dim * 4 bytes * 24 layers, about 101 MB. A sixteen-minute section has 118 segments. One lane per segment asks for 11.9 GB of cache alone, and on a 16 GB machine that means swap — plainly visible in the timings while invisible in the memory metrics. The AR stage went RTF 0.19 to 0.70; the engine 0.70 to 1.49.

When a metric returns the same value for inputs that must differ, it is not telling you the system is stable. It is not telling you anything.

Where the time is

With the timers honest:

The flow decoder should dominate. Ten Euler steps, doubled for guidance, across twenty-two DiT blocks is 440 block passes over the full mel length per utterance. The vocoder is one feed-forward pass.

Stock upstream reaches RTF 4.370 here, and that number needs its caveat in the open rather than in a footnote: CosyVoiceModel.__init__ hardcodes cuda if available else cpu. Upstream has no MPS path at all. CPU is not a badly chosen baseline, it is the only thing stock upstream can do on this machine. The MPS-adapted service from the first round reaches about 0.76, and this port is now slightly past it.

What the library was already willing to do

Before any custom kernel, a series of changes that amount to using candle properly. All measured with interleaved runs against per-block fixture gates:

Attention taught me something I should have known already. Assembled, it cost 8.2× a projection while doing five times less arithmetic — 2.6 GMAC against 13.4 per block. The cost is traffic. The assembled form materialises a [2, 16, 798, 798] scores tensor, 81.5 MB, then touches it four times: write, scale, softmax, read. Call it 490 MB of memory traffic to do 2.6 GMAC. The fused kernel never materialises it at all.

When an operation’s cost is wildly out of proportion to its arithmetic, stop optimising the arithmetic and count the bytes.

A 1.25× from renaming things

In this model only head 0 gets rotary embedding. Attention is independent per head, so the rotated head and the other fifteen can be two sdpa calls over views instead of one call over a rebuilt tensor — saving a 6.5 MB reconstruction of q and another of k. Measured 1.25× on a whole DiT block, within 0.25 ms of a variant with the rotary deleted entirely, which is the ceiling for that change.

The direct implementation returns wrong answers. candle 0.10.2’s Metal sdpa mishandles a head axis narrowed to a non-zero offset: narrow(1, 0, 1) agrees with the naive form to 6.7e-7 relative, while narrow(1, 1, 15) comes back at relative 1.24. Not precision. Wrong.

So avoid the bug rather than working around it. Permute the head blocks of to_q, to_k, to_v and to_out at load time so the rotated head sits last. The fifteen unrotated heads become narrow(1, 0, 15) — offset zero, which the kernel handles — and the single rotated head, which does carry an offset, is one sixteenth of the tensor, so making it contiguous costs almost nothing. Relabelling is exact as long as every projection agrees; to_out consumes heads along its input columns, so it takes the same permutation. Block validates at 8.56e-7.

I benchmarked that variant before I checked it was correct. A 1.25× on a function returning garbage nearly became an optimization, and the only thing that caught it was a fixture gate written months earlier for unrelated reasons.

The refutation that was itself wrong

An earlier investigation had ruled out implementing the vocoder’s convolutions as a GEMM and blamed materialisation traffic. Wrong conclusion, and wrong for a specific reason: the experiment measured two things at once.

Split the route in half and the picture inverts. The GEMM runs at 2.4 TFLOP/s — 8.4× faster than the direct convolution. The im2col gather feeding it runs at 4.9 GB/s. The gather was the entire problem. Reordering its rows from tap-interleaved to tap-major, so each thread reads contiguously, made it 3.01× faster for the price of one weight permutation at load. The upsampling decoder went 454.0 ms to 349.7 ms, bit-identical output.

Now the part that matters more than the win. CosyVoice’s end-to-end RTF did not move: 1.612 to 1.602, inside run-to-run noise. The vocoder is 16% of this engine and the isolated decoder is not even all of that. The same change took Audio8 — different engine, codec is 32% of its runtime — from 0.664 to 0.499.

A real 1.30× on a stage the total cannot see is still real. It is worth nothing here and a great deal one crate over. Stage benchmarks tell you whether a change works; only the end-to-end number tells you whether to keep it. I have watched more than one team ship a quarter of work that made a benchmark faster and a product identical.

Scheduling did the heavy lifting

flow.synthesize prepends the reference voice’s 294 speech tokens and 588 mel frames on every call. Decode segment by segment and you pay that every time. On the test passage: 7 × 588 = 4116 prompt frames against 2634 generated ones. Sixty-one percent of the flow stage was one prompt, decoded seven times.

Concatenate the tokens, decode once: 1.71× on the flow, plus 1.72× on the vocoder from making one call instead of seven. No kernel. No arithmetic changed. Deleting work that did not need doing.

It cost something. Fusing removes the inter-segment silence, and WER went 0.000 to 0.023 — three errors in 133 words, sentence boundaries running together. You do not have to give the speed back, because the cut points are exact: two mel frames per speech token, 480 samples per frame, so a segment of n tokens is exactly n * 960 samples. Cut the fused waveform there, re-insert the gaps, WER 0.015. The PyTorch reference scores 0.008 on the same text.

Then I pushed the same lever too far and it turned around on me.

Fusing everything looks strictly better, and my engine preferred it whenever it fitted. It looked good because I measured it on a short passage. DiT attention is quadratic, so per-block cost behaves like a*n + b*n². Fitting coefficients from probes at 798 and 3192 frames and projecting across a seventeen-minute section — a model, not a measurement:

A 3.44× penalty, hidden for weeks behind a fixed frame cap in an unrelated asset that had been quietly forcing the fallback. The engine now targets 2000 generated frames per call, in the middle of a broad flat minimum.

The AR stage gets batched for a different reason — a decode step’s cost is reading twenty-four layers of weights and a host round-trip for sampling, neither of which scales with the sequence — and capped at eight lanes because of the 101 MB each. Right-aligning the prompts keeps it exact to f32, agreeing to 2.0e-6 relative rather than bit-identically. The measured throughput curve is flat by four lanes, so the cap costs nothing and prevents 11.9 GB.

Where quantization stops paying

Only the LLM’s projections are quantized, to q8_0. Relative weight error is 0.55% there, against 4.4% one step down, and the voice is measurably unchanged.

The flow decoder and vocoder are deliberately left alone. They run over full sequences, where candle takes an ordinary GEMM rather than the dedicated matrix–vector kernel, so quantizing them adds a dequantize per call and buys nothing.

One wrinkle if you combine tricks, invisible from either measurement alone: candle’s quantized matvec kernel only fires when the second-to-last dimension is 1. Batch more than one sequence and you lose it to the GEMM path. Batching and quantization are both wins, and they partially cancel.

The list I stopped believing in

The head-transpose kernel is my favourite of these. It did precisely what it promised — 4 to 7× on the operation — and made the thing it was written for very slightly slower, because the transpose was never the cost. It stays in the tree, documented as measured-but-unused, so nobody writes it again.

Keep this list. A tuning pass without a refutation column has not finished measuring.

Where it landed

RTF 0.697 on the 132-word passage, 0.72–0.74 across a seventeen-minute section where longer sequences cost more attention. The largest single un-won item is the DiT’s grouped position-embedding convolution: 24.6 ms for 3.2 GMAC, about 0.13 TFLOP/s, roughly 8% of the flow.

Ranking what produced the 4.2×:

Scheduling first — one flow call per group instead of per segment, batched AR decoding. Largest contributions, no new kernels.

Using the library properly second — fused ops, views instead of copies, flat matmuls. Most of the remainder, nearly free.

Hand-written Metal last. Exact, real, and mostly invisible in this engine’s total while being decisive in another one.

None of the big wins made an operation faster. They deleted operations: a prompt decoded seven times, a scores tensor written and read four times, weights re-read once per sequence instead of once per batch. The kernels were the interesting engineering and the smallest column in the ledger.

And I would have found all of it sooner by auditing the profiler before trusting it.