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%. Everything here is downstream of repairing that instrument: RTF 2.96 down to 0.697 on the same 132-word passage, reimplementing in Rust on candle a model I had already made fast once in PyTorch. RTF is synthesis time over audio duration, so under 1.0 beats real time.

Fix the instrument first

Metal dispatch is asynchronous, so a timer reading the clock after encoding charges that stage’s time to whichever later stage blocks on the queue. Adding a device.synchronize() before each stage moved two numbers a long way:

I chose targets from 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 and phys_footprint_peak returned exactly 13.00 GB for three configurations that should have differed enormously, both saturated, and I built a fix on top of the second reading before admitting it told me nothing. A metric returning the same value for inputs that must differ is not reporting stability; it is reporting nothing.

Arithmetic found what the metrics could not. Each autoregressive lane carries its own KV cache, at 2 * n_kv * capacity * head_dim * 4 bytes * 24 layers or roughly 101 MB, and a sixteen-minute section has 118 segments, so one lane per segment asks for 11.9 GB. On a 16 GB machine that means swap, plainly visible in the timings and invisible in every memory metric I had: the AR stage went from RTF 0.19 to 0.70 and the engine from 0.70 to 1.49.

With honest timers, the distribution is unsurprising. Ten Euler steps doubled for guidance across twenty-two DiT blocks is 440 block passes per utterance, against one feed-forward pass for the vocoder.

Stock upstream reaches 4.370 here, and the caveat belongs in the open: CosyVoiceModel.__init__ hardcodes cuda if available else cpu, so CPU is the only thing it can do on this machine. The MPS-adapted service from my first round reaches about 0.76, and this port is slightly past it.

Using the library properly

Before any custom kernel, changes that amount to reading candle’s API carefully, measured with interleaved runs against per-block fixture gates:

Assembled attention cost 8.2x a projection while doing five times less arithmetic, 2.6 GMAC against 13.4 per block, because the cost is traffic: it materialises a [2, 16, 798, 798] scores tensor of 81.5 MB and touches it four times to write, scale, softmax and read, roughly 490 MB for 2.6 GMAC, while the fused kernel never materialises it at all. When cost is out of proportion to arithmetic, stop optimising the arithmetic and count the bytes.

Two kernel stories

Only head 0 receives rotary embedding, and attention is independent per head, so the rotated head and the other fifteen can be two sdpa calls over views rather than one call over a rebuilt tensor, saving a 6.5 MB reconstruction of q and another of k. That measured 1.25x on a DiT block, within 0.25 ms of deleting the rotary entirely, which is the ceiling.

The direct implementation returns wrong answers, because 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. So avoid the bug rather than working around it, by permuting the head blocks of to_q, to_k, to_v and to_out at load so the rotated head sits last. The fifteen unrotated heads become narrow(1, 0, 15) at offset zero, which the kernel handles, and the one head carrying an offset is a sixteenth of the tensor, so making it contiguous costs almost nothing. The block validates at 8.56e-7.

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

The second story is a refutation that was itself wrong. An earlier investigation had ruled out implementing the vocoder’s convolutions as a GEMM and blamed materialisation traffic, but that experiment measured two things at once. Split the route in half and the picture inverts: the GEMM runs at 2.4 TFLOP/s, 8.4x faster than the direct convolution, while the im2col gather feeding it runs at 4.9 GB/s. Reordering the gather’s rows from tap-interleaved to tap-major made it 3.01x faster for one weight permutation at load, taking the upsampling decoder from 454.0 ms to 349.7 ms, bit-identical.

Then the part that matters more than the win. End-to-end RTF did not move, 1.612 to 1.602, inside noise, because the vocoder is 16% of this engine. The same change took Audio8, whose codec is 32% of runtime, from 0.664 to 0.499. Stage benchmarks tell you whether a change works; only the end-to-end number tells you whether to keep it.

Scheduling did the heavy lifting

flow.synthesize prepends the reference voice’s 294 speech tokens and 588 mel frames on every call, so on the test passage 7 x 588 = 4116 prompt frames against 2634 generated ones: sixty-one percent of the flow stage was one prompt decoded seven times. Concatenating the tokens and decoding once gave 1.71x on the flow plus 1.72x on the vocoder from one call instead of seven. No kernel, no arithmetic changed, just deleting work.

Fusing removes the inter-segment silence, and WER went from 0.000 to 0.023, three errors in 133 words with sentence boundaries running together. The speed need not be given back, because the cut points are exact: two mel frames per speech token at 480 samples per frame makes a segment of n tokens exactly n * 960 samples, so cutting there and re-inserting the gaps returns WER to 0.015 against the PyTorch reference’s 0.008.

Then I pushed the same lever too far. Fusing everything looks strictly better and my engine preferred it whenever it fitted, which looked good only because I measured it on a short passage. DiT attention is quadratic, so per-block cost behaves like a*n + b*n², and fitting coefficients from probes at 798 and 3192 frames to project across a seventeen-minute section gives a model rather than a measurement:

A 3.44x 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 a broad flat minimum.

The AR stage batches for a different reason, since a decode step’s cost is reading twenty-four layers of weights plus a host round-trip for sampling, neither of which scales with sequence length. It caps at eight lanes because of the 101 MB each; right-aligning the prompts keeps it exact to 2.0e-6 relative, and the throughput curve is flat by four lanes, so the cap costs nothing and prevents 11.9 GB.

Only the LLM’s projections are quantized, to q8_0, where relative weight error is 0.55% against 4.4% one step down. The flow decoder and vocoder run over full sequences, where candle takes an ordinary GEMM rather than its matrix-vector kernel, so quantizing them adds a dequantize per call and buys nothing. One wrinkle is invisible from either measurement alone: that matvec kernel only fires when the second-to-last dimension is 1, so batching loses 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, because it did precisely what it promised and made the thing it was written for slightly slower, since the transpose was never the cost. It stays in the tree documented as measured-but-unused so nobody writes it again. A tuning pass without a refutation column has not finished measuring.

Where it landed

RTF 0.697 on the 132-word passage and 0.72 to 0.74 across a seventeen-minute section. The largest un-won item is the DiT’s grouped position-embedding convolution at 24.6 ms for 3.2 GMAC, about 0.13 TFLOP/s and roughly 8% of the flow.

Ranking the 4.2x: scheduling first, one flow call per group and batched AR decoding, the largest contributions and no new kernels; using the library properly second, most of the remainder nearly free; hand-written Metal last, real and mostly invisible in this engine while decisive in another. 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.

CosyVoice3 is no longer what I narrate with. Qwen3-TTS replaced it at twice the size and RTF 0.253, where the same lesson about scheduling arrived again and I still failed to apply it first.