fix: dispose KV cache after Chatterbox generation (#1734) - #1737
Open
m96-chan wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #1734.
ChatterboxModel.generateleaked 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.generateonly disposes the cache when it is not handing it back to the caller: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.