Skip to content

fix: dispose KV cache after Chatterbox generation (#1734) - #1737

Open
m96-chan wants to merge 1 commit into
huggingface:mainfrom
m96-chan:bugfix/dispose-kvcache-1734
Open

fix: dispose KV cache after Chatterbox generation (#1734)#1737
m96-chan wants to merge 1 commit into
huggingface:mainfrom
m96-chan:bugfix/dispose-kvcache-1734

Conversation

@m96-chan

@m96-chan m96-chan commented Aug 2, 2026

Copy link
Copy Markdown

Summary

Fixes #1734.

ChatterboxModel.generate leaked a full KV cache of GPU buffers on every call, so GPU memory grew without bound in any application that synthesizes repeatedly.

Root Cause

PreTrainedModel.generate only disposes the cache when it is not handing it back to the caller:

const keepCacheAlive = 'past_key_values' in kwargs || generation_config.return_dict_in_generate;
if (!keepCacheAlive) {
    await past_key_values.dispose();
}

ChatterboxModel.generate always opts into the dict form, because it needs audio_tokens, speaker_embeddings and speaker_features alongside sequences.
keepCacheAlive is therefore permanently true for this architecture, so the disposal in the base class never runs. The wrapper then destructured only the four values it needed and let past_key_values go out of scope without disposing it,
leaving its GPU buffers allocated — up to max_new_tokens worth of keys and values across every layer, per call. The caller of ChatterboxModel.generate receives only a waveform, so nothing downstream can free it either.

Fix

Destructure past_key_values in ChatterboxModel.generate and dispose it once the
conditional_decoder run is finished:

await past_key_values?.dispose();

The dispose happens after conditional_decoder has run, so the cache stays alive for as long as generation actually needs it. DynamicCache.dispose() already exists and does the right thing; this simply calls it on the one path that never did.
The optional chaining keeps the change safe for configurations where no cache is returned.

Impact

Measured in Chrome on WebGPU (RTX 5090, Dawn/Vulkan), onnx-community/chatterbox-ONNX, 20 generations at max_new_tokens: 256, GPU process sampled once a second:

growth over 20 generations
before:
+1682 MiB — twenty clean steps of ~120 MiB
after:
+36 MiB — flat at 2422 MiB from the third generation onward

The remaining residue is one-off rather than per-call: once it settles it does not move again for the rest of the run. Generation speed is unchanged (5.3 s per call before, 5.3–5.7 s after), so nothing is traded away for it.

Scope is limited to ChatterboxModel.generate; no public API or behavior change — callers already only received the waveform.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Chatterbox KV cache leak

1 participant