fix(asr): free cache-aware slots on delete_state and session reset

[dundysm]

What does this PR do ?

Frees cache-aware feature-bufferer and encoder-cache slots when a stream ends via delete_state() or when a session is reset/closed/opened, so callers that never send is_last=True (notably NeMoStreamingPipelineAdapter) no longer exhaust num_slots.

Collection: ASR

cc @naymaraq @nithinraok @artbataev @lilithgrigoryan

Changelog

• BatchedCacheFeatureBufferer.free_stream(stream_id): idempotent return of the stream's slot (no-op if already freed via is_last in update()).
• BatchedCacheFeatureBufferer.reset(): clear maps, refill available_slots, zero feature buffer and audio bufferers.
• CacheAwareContextManager.free_stream(stream_id): membership-guarded free so double free cannot corrupt free_slots (unlike raw _reset_slots).
• CacheAwareRNNTPipeline and CacheAwareCTCPipeline: override delete_state to free both slot pools, and reset_session to restore all slots (covers open_session / close_session via the base path).
• CPU unit tests with num_slots=2, no checkpoint / HF download.

The problem

On main, both slot pools free a stream only when a request carries is_last=True:

• BatchedCacheFeatureBufferer.update appends to slots_to_free only for frame.is_last
• CacheAwareContextManager.reset_slots frees only streams whose eos flag is true (eos_flags come from request.is_last)

BasePipeline.delete_state and session open/close only clear _state_pool. The simulstream adapter always sends is_last=False and ends with delete_state(), so after num_slots streams the next allocate raises RuntimeError: No free slots available. close_session() did not restore the pools either.

Approach (design-preserving)

Keep the intentional is_last free path for graceful EOU / decoding / bias release. Extend abrupt teardown:

1. delete_state(stream_id) frees both slots if still mapped (idempotent after a prior is_last free).
2. reset_session (and therefore open_session / close_session) rebuilds both pools.

Do not flip the adapter to is_last=True as the sole fix: that would change EOU / keep_all_outputs / hypothesis-reset semantics and still leave aborted streams / session boundaries leaking.

Why not only adapter is_last=True

Reporter already flagged that option and has not measured output impact. Slot free via delete_state matches how the adapter documents end-of-stream today. Session reset remains necessary either way.

Tests

tests/collections/asr/inference/test_cache_aware_slot_release.py (CPU, no HF):

• bufferer / context-manager free_stream returns capacity and is idempotent
• reset restores a fully leaked pool
• is_last path via reset_slots(..., eos=True) still works; following free_stream does not over-fill the queue
• many short streams ending only via delete_state recover free counts (issue repro shape with num_slots=2), for RNNT and CTC
• delete_state after a simulated is_last free is a no-op on the queues
• close_session / open_session restore leaked slots
• control: BasePipeline.delete_state alone still leaks (documents the bug)

Local run:

13 passed

Related, deliberately not in this PR

• fix(asr): release a stream's biasing model when its state is deleted #16308 (biasing release on delete_state): sibling leak; leave to that PR.
• Cache-aware streaming skips a final chunk shorter than sampling_frames, so a stream's last word depends on its batch #16310 (sampling_frames tail): different file and semantics.
• Speed up the cache-aware streaming RNNT pipeline #16231 (cache-aware speedup): same files, still is_last-only free; expect rebase if it lands first. This PR adds helpers beside the hot update path and does not rewrite it.
• Buffered BatchedAudioBufferer per-stream dict growth: soft leak, out of scope.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc) — no

PR Type:

• New Feature
• Bugfix
• Documentation

Additional Information

• Fixes Cache-aware streaming pipeline runs out of slots when streams end with delete_state() instead of an is_last request #16309
• DCO: Signed-off-by trailer on the commit

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dundysm]

pinging @naymaraq @nithinraok @artbataev @lilithgrigoryan for cache-aware streaming / asr review. happy to rebase if #16231 lands first (same files, different concern).

[kzos]

Thanks for picking this up, @dundysm. I ran this branch (ff02c54) and main (00278b0) on CPU in float32 with the in-tree cache_aware_rnnt.yaml (nemotron-speech-streaming-en-0.6b, 256 slots) and cache_aware_ctc.yaml (stt_en_fastconformer_hybrid_large_streaming_multi, 1024 slots), otherwise unchanged apart from ITN/NMT off. Slot rows: one 1.12 s noise chunk per stream, ended with delete_state().

	main	this PR
RNN-T free slots after 1 / 128 / 256 / 300 streams	255 / 128 / 0 / stream 257 raises No free slots available	256 throughout
CTC free slots after 1 / 300 / 1024 / 1100 streams	1023 / 724 / 0 / stream 1025 raises	1024 throughout
RNN-T, 100 streams left open, then close_session()	156 free	256 free
pipeline.run() (batch_size=4), 12 LibriSpeech test-other utterances, segment and word level		identical to main for both models: text, segments, word times; confidences compared for CTC only (RNN-T gave 0.0 for every word with the in-tree config)
id reused after delete_state() without is_last, 6 utterance pairs: text differs from a fresh stream	RNN-T 1/6, CTC 5/6	0/6, 0/6
simulstream adapter, 12 utterances, a new instance per utterance vs one instance: transcripts differ	RNN-T 4/12, CTC 4/12	0/12, 0/12

Your test file: 13 passed here; copied onto main, 10 of the 13 fail.

The last two rows are extra motivation: on main the id keeps both slots after delete_state(), encoder cache not cleared, and the next stream on that id continues from them (the adapter starts every instance at stream id 0 on one class-level pipeline). I have a small CPU reuse test that passes here and would catch a free_stream() that skipped the zeroing (all 13 tests here still pass with that change); happy to share it.

#16308 adds CacheAwareRNNTPipeline.delete_state at the same place, so the two conflict. On this branch alone, #16308's adapter run (a biasing request per stream) no longer stops at stream 257, so the biasing-model leak it addresses keeps growing (300 active after 300 streams). With the resolution below (and no-op slot pools in #16308's test fixture) the merge keeps 256 free slots and 0 active models after every stream, passes both test files (17) and tests/collections/asr/inference (184 passed, 7 skipped), and close_session() still releases all biasing models. If this lands first, I'll rebase #16308 with it; otherwise the same resolution applies here, and I'm happy to help.

def delete_state(self, stream_id: int) -> None:
    state = self.get_state(stream_id)
    if (
        state is not None
        and self.decoding_computer is not None
        and self.decoding_computer.per_stream_biasing_enabled
    ):
        release_auto_managed_stream_biasing(state, self.decoding_computer.biasing_multi_model)
    self._free_cache_aware_slots(stream_id)
    super().delete_state(stream_id)

The buffered pipelines' per-stream bufferers (RNN-T and CTC, left out of scope here, and SALM's audio_bufferer) are handled in my separate PR #16322.

A question: open_session() and close_session(), both called by pipeline.run(), now build a new encoder cache through get_initial_cache_state() while the old one is still referenced: +1.6 GiB (RNN-T) / +2.6 GiB (CTC) peak RSS and roughly 0.2–0.5 s per call on CPU, where main allocates nothing. The initial cache is all zeros, so would zeroing it in place be enough? Locally that took roughly 0.1–0.16 s with no extra peak (not measured on GPU).

[kzos]

Two inline notes on the lines they refer to; the measurements are in my comment above.

[kzos]

CacheAwareContextManager.reset() builds a new cache through get_initial_cache_state(num_slots) while the old one is still referenced, and pipeline.run() reaches this twice (through open_session() and close_session()). On CPU with nemotron-speech-streaming-en-0.6b and 256 slots I measured +1.6 GiB peak RSS and roughly 0.2–0.3 s per call here, where main allocates nothing (details in my comment above; not measured on GPU). The initial cache is all zeros for this model, and zeroing the existing tensors in place took roughly 0.1 s with no extra peak. Would in-place zeroing work for you here? The CTC pipeline's reset_session() makes the same call (+2.6 GiB and roughly 0.3–0.5 s with 1024 slots).

[kzos]

#16308 (mine) adds CacheAwareRNNTPipeline.delete_state() at this same spot, to release the stream's per-stream biasing model, so the two conflict whichever lands second. The combined method is in my comment above: release the biasing model while the state still exists, then _free_cache_aware_slots(), then super().delete_state(). If this lands first, I'll rebase #16308 onto it.

[dundysm]

thanks @kzos for the thorough cpu verification and for confirming the fix on both rnnt and ctc, including the stream-id reuse / simulstream correctness rows. that extra motivation is helpful.

agree on both follow-ups:

1. in-place zero of the encoder cache on session reset (instead of realloc via get_initial_cache_state while the old tensors are still referenced). that matches what _reset_slots and the bufferer reset already do, and avoids the peak rss / time you measured on cpu for open_session/close_session inside pipeline.run(). i have not re-measured on gpu, so i will not claim gpu numbers; happy to take your cpu figures as the bar for the change.

2. coordinate delete_state with fix(asr): release a stream's biasing model when its state is deleted #16308: fold your combined order into CacheAwareRNNTPipeline.delete_state here (release_auto_managed_stream_biasing while state exists, then _free_cache_aware_slots, then super). ctc stays slot-only. if this lands first you can rebase fix(asr): release a stream's biasing model when its state is deleted #16308 to keep the buffered path + tests; happy to credit you on the rnnt hunk.

i pushed a follow-up commit for (1) and (2), kept the existing 13 cpu unit tests green, and added a small check that reset() zeros in place rather than reallocating. your offered cpu reuse test that catches a free_stream skipping zeroing is welcome if you want to share it; free_stream here already goes through _reset_slots so it should already zero.

buffered per-stream bufferers / salm left to #16322 as you noted.