TTSKit: decode-path performance improvements for Qwen3 streaming

[mjfrey]

Four small optimizations to the Qwen3 TTS decode hot path, found while profiling on-device streaming synthesis in a production app (A19 iPhone, 12Hz 0.6B models):

• Stateful decoders no longer copy K/V updates into the host-side cache. For stateful models the MLState buffers already received this step's key/value updates during prediction, but the decode path still materialized the update tensors into the host-side cache arrays on every step. Only the cache position and attention masks need to advance. Adds KVCache.update() for that, with a unit test covering the position/mask bookkeeping.
• K/V tensor materialization is concurrent. The key and value MLTensor outputs are independent, but the async update paths materialized them sequentially — two round trips per decode step. async let runs both toMLMultiArray() calls concurrently in update(keyTensor:valueTensor:), updateStateCache, and updateWithHiddenContext.
• BinaryDistinctString.init(NSString) copies UTF-16 code units directly. Converting through Array(String) walks Swift extended grapheme clusters and re-expands them to UTF-16 — needlessly expensive while loading large tokenizer vocabularies, and it can lose the exact NSString representation the type exists to preserve. getCharacters(_:range:) copies the code units as-is.
• Top-k sampling selects the top-k logits before softmax instead of computing probabilities over the entire vocabulary each step. Equivalent after renormalization, and avoids materializing full-vocabulary probabilities on the host.

Measurements

A/B on an Apple silicon Mac (CLI, v1.0.0 vs v1.0.0 + these changes only, same seed/text/models, W8A16 0.6B, temperature 0.7 / top-k 20, 3 runs each):

Metric	Before	After	Δ
Tokenizer load (Qwen3 151k vocab)	0.95 s	0.64 s	−32%
MultiCodeDecoder sampling	0.91 ms/step	0.81 ms/step	−11%
Total process CPU (16.5 s generation)	~9.9 s	~9.6 s	−2.5%
Throughput / RTF / peak RSS	—	—	unchanged

End-to-end throughput is unchanged on the Mac because generation there is ANE-bound; the host-side savings are aimed at phone-class devices, where the decode loop shares CPU and memory bandwidth with the host app. The original motivation was an Instruments capture on an A19 iPhone during streaming synthesis, which attributed ~16% of captured cycles to the sampling path that these changes shrink.

All changes are behavior-preserving. swift build and the TTSKit unit tests (86) pass; the changes have been validated in continuous day-to-day on-device streaming use.

[ZachNagengast]

Thanks for the submission @mjfrey! Were you able to quantify the optimizations in terms of speed up or memory usage?

[Copilot]

Pull request overview

This PR optimizes several hot-path operations in Qwen3 streaming TTS decoding/sampling to reduce per-step overhead (cache updates, tensor materialization, tokenizer string handling, and sampling).

Changes:

• Avoids unnecessary host-side work for stateful decode steps by only advancing KV cache position/masks.
• Materializes independent key/value tensors concurrently via async let to reduce serialized latency.
• Improves sampling performance by selecting top‑k logits before softmax and computing the (renormalized) distribution on the reduced set.
• Speeds up tokenizer vocabulary loading by copying UTF‑16 code units directly from NSString.
• Adds a unit test covering stateful KV cache position/mask advancement.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
Tests/TTSKitTests/TTSKitUnitTests.swift	Adds a unit test for stateful KVCache.update() mask/position advancement.
Sources/TTSKit/Utilities/Sampling.swift	Adds a top‑k-before-softmax sampling fast path for multi-head tensor sampling.
Sources/TTSKit/Utilities/KVCache.swift	Runs key/value tensor materialization concurrently in several update paths.
Sources/TTSKit/Qwen3TTS/Qwen3CodeDecoder.swift	Skips tensor materialization for stateful cache updates (advances only position/masks).
Sources/ArgmaxCore/External/Hub/BinaryDistinct.swift	Optimizes BinaryDistinctString.init(NSString) by directly copying UTF‑16 code units.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mjfrey]

Thanks for the submission @mjfrey! Were you able to quantify the optimizations in terms of speed up or memory usage?

@ZachNagengast Yes — A/B on an Apple silicon Mac (CLI, v1.0.0 vs v1.0.0 + these changes only, same seed/text/models, W8A16 0.6B, temperature 0.7 / top-k 20, 3 runs each):

Metric	Before	After	Δ
Tokenizer load (151k vocab)	0.95 s	0.64 s	−32% (consistent across runs)
MultiCodeDecoder sampling	0.91 ms/step	0.81 ms/step	−11%
Total process CPU (16.5 s generation)	~9.9 s	~9.6 s	−2.5%
Throughput / RTF / peak memory	19.8 steps/s, 0.64, ~470 MB	same	unchanged

Throughput and memory are unchanged on the Mac because generation there is ANE-bound — these are host-CPU savings, aimed at phone-class devices where the decode loop shares CPU and memory bandwidth with the host app. The original motivation was an Instruments capture on an A19 iPhone during streaming synthesis, which attributed ~16% of captured cycles to the sampling path.