Repository navigation
perf(lfm2_audio): reuse the turn before's prefill and encoder output - #845
Merged
Merged
Conversation
The chunked prefill that later S2S turns run builds its 256-step block graphs with the backbone's weights, which with lfm2_audio.cpu_repack may sit in CPU extra buffers, but validated them with validate_backend_graph_supported, where the backbone's other graphs and those of the encoder, depthformer and detokenizer use validate_backend_graph_with_cpu_extra_buffers. The block graphs use those weights only through LinearModule's matmuls, as the one-shot prefill does, so the extended checks pass and nothing changes; a later change that viewed such a weight, or fed it a strided input, now fails when the graph is built rather than reading misplaced rows. With lfm2_audio.cpu_repack on, main's default, every later turn of a ten-turn English conversation passed the checks with the Q4_0 package on an Intel Core i9-13900HK laptop CPU and on a Xeon 8358, and with the Q4_0 and Q8_0 packages on the CPU of an M3 Ultra, types that ggml's repacked kernels take there. Together with the next commit, which changes no output either, that conversation streamed with Q4_0 and its questions asked alone give main's bytes on all three, and cpu_extra_buffers_test and the lfm2_audio unit tests pass.
A later turn's prompt is the turn before's with its reply and the new question added, but every chunked prefill ran from position 0, so a turn's prefill grew with the conversation. In a ten-turn English conversation, streamed with the Q4_0 package on an Intel Core i9-13900HK laptop CPU at 8 threads, the tenth turn's prefill took 10.4 s of its 11.1 s to first audio. A block of the chunked prefill computes the same bits from the same span, on the same keys, values and conv state, whatever the cache's length. So the backbone now keeps a host copy of what each full block of the last chunked prefill wrote: its rows of each attention layer's keys and values, 24 KiB per step for the published checkpoints, the conv state after it, 160 KiB, and the block's span. The next chunked prefill compares its blocks' spans with the kept ones from the start, ids, positions and frame rows by value and audio rows by their bytes, restores the blocks that match into the cache it cleared, as running them would have left it, and runs the rest as before. Its last block always runs, as its logits are the prefill's. A shorter last block is never kept, nor is anything a one-shot prefill or a decode step wrote, and a one-shot prefill leaves the kept blocks as they are. The decode cache now allocates the host scratch it stages a whole state in only when a one-shot prefill imports one, so a cache that only chunked prefills have written no longer holds 24 KiB per cache step of zeros there. The bytes it writes are the same. In that conversation, in one run of each, the tenth turn now restores 2560 of its prompt steps, prefills in 1.5 s and reaches its first audio in 2.2 s, and the ten turns' prefills take 16.3 s instead of 56.4 s. The conversation's peak RSS went from 1445 MiB to 1386 MiB. With Q4_0, the tenth turn reaches its first audio in 1775 ms instead of 8225 ms on a Xeon 8358 at 8 threads, in 625 ms instead of 1993 ms on an M3 Ultra's CPU at 16 threads, in 350 ms instead of 751 ms on its Metal and in 218 ms instead of 474 ms on an A10 (CUDA). Peak RSS went from 1605 MiB to 1549 MiB on the Xeon and, with the F16 package offline, from about 750 MiB to 683 MiB on the A10; on the M3 Ultra, and on the A10 with Q4_0, it moved by under 10 MB. Each turn's reply artifact, text, audio and stream events are main's bytes, and so are those of the ten questions asked alone: offline and streamed on the M3 Ultra's CPU with Q4_0 and Q8_0, on its Metal with Q4_0 and on the A10 with Q4_0 and F16, and streamed with Q4_0 on the i9-13900HK and the Xeon. Each A10 tree carried the fix for ggml-cuda's block_reduce race, which can otherwise change a run's bytes. The timing log gets lfm2_audio.prefill.resumed_steps, and the runtime resumed_prefill_steps(). lfm2_audio_backbone_test checks prompts of 300 to 1000 steps on one runtime against a fresh one, bit for bit, in the prefill's logits and six steps after it, at 1, 4 and 8 threads and with caches of other lengths, and how much a bit flipped in an audio row, a changed frame code, a one-shot request and a text-only chunked request in between let the next prefill resume.
Each later turn also ran every earlier question through the encoder again: the tenth turn of the ten-turn English conversation, streamed with the Q4_0 package on an Intel Core i9-13900HK laptop CPU at 8 threads, spent 0.45 s encoding its ten questions. The S2S session now keeps the questions of its last request with history or return_codes, as the mono 16 kHz samples the encoder took, with their encoder output. An earlier question whose samples are those of a kept one, compared by size and bytes, takes that output instead of the features and the encoder, which on one session give the same bits for the same samples whatever they encoded before. The new question always goes through the encoder. A request without history or return_codes keeps nothing and leaves the kept questions as they are, so the session keeps none of that request's audio. The conversation limit is checked at the same points with the same counts, so a request fails with the same error as before. The trace log gets lfm2_audio.session.reused_questions, and the session resumed_prefill_steps() and reused_questions(). A request replaces the kept questions only once it has passed the conversation limit, so one that fails leaves them as they were. The questions it takes from them it shares with them, rather than holding a second copy of their samples and output. The tenth turn of that conversation now encodes one question, in 47 ms, and reaches its first audio in 1.7 s, against 2.2 s with the commit before and 11.1 s on main; the ten turns' encoder calls take 0.5 s instead of 2.8 s. With the Q8_0 package its first audio comes after 4.6 s against 24.3 s on main, and with Q4_0 and lfm2_audio.cpu_repack off after 3.6 s against 26.3 s. With Q4_0, it comes after 1400 ms, against 1775 ms and 8225 ms, on a Xeon 8358 at 8 threads, after 401 ms against 625 ms and 1993 ms on an M3 Ultra's CPU at 16 threads, after 125 ms against 350 ms and 751 ms on its Metal, and after 118 ms against 218 ms and 474 ms on an A10 (CUDA). The conversation's peak RSS on the i9-13900HK was 1390 MiB, against 1386 MiB with the commit before and 1445 MiB on main, and on the Xeon 1554 MiB, against 1549 MiB and 1605 MiB. The ten turns and the ten questions asked alone give main's bytes: offline and streamed on the i9-13900HK with Q4_0 and Q8_0, and streamed with Q4_0 and cpu_repack off; offline and streamed on the M3 Ultra's CPU with Q4_0 and Q8_0, on its Metal with Q4_0 and on the A10 with Q4_0 and F16, each A10 tree with ggml-cuda's block_reduce race fix; and streamed on the Xeon with Q4_0. On the Xeon the ten turns streamed with Q8_0 give main's bytes too. The tenth turn sent with the same history to a new process gives the conversation's tenth turn, on the A10 and on the Xeon. lfm2_audio_encoder_test now compares bit for bit where it checks that a runtime gives a fresh one's output after chunks of other lengths, one past 30 s whose buffer is freed among them, and the same output for the same input after other input. lfm2_audio_chat_test runs a four-turn conversation in one session, offline and streamed, whose prompts pass 256 and 512 steps: each turn's text, audio and reply artifact are a fresh session's, with 0, 0, 256 and 512 prefill steps resumed and 0 to 3 questions reused. The fourth turn then runs again after each of these: another conversation's first two turns, after which it resumes nothing; that conversation's first turn alone, with return_codes, after which it encodes its questions again but still resumes; itself over the limit, which fails with a fresh session's error after each question and on the whole prompt; and a request without history or return_codes. The session's comment no longer says it keeps nothing between requests.
The Conversations section said the S2S session keeps nothing between requests and that each turn prefills the whole conversation again. It now says what a session keeps for the next turn of one conversation: the questions of its last request with history or return_codes, with their encoder output, about 160 KiB per second of question, and the full prefill blocks of its last turn with history, 6 to 8 MiB per 256 steps, up to about 350 MiB in all at the conversation limit. It says that a turn gives the same bytes with or without them, which requests replace which of them and what the next turn then runs again, that a turn that leaves out the oldest turns runs its prefill in full, and that a question sent back as other samples is encoded again. The validation section covers the new tests, and the memory section points to the Conversations section.
Owner
|
@ykhrustalev Thansk, PR merged! |
This was referenced Oct 9, 2026
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
2 participants
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.
Follow-up to #838: a later S2S turn restores the prefill blocks it shares with the turn before and reuses the earlier questions' encoder output, instead of running the whole conversation again. Output does not change.
Problem
validate_backend_graph_with_cpu_extra_buffers, as the backbone's prefill and decode graphs do, instead ofvalidate_backend_graph_supported. No failure from the omission was foundWhat this PR changes
BlockGraphvalidates withvalidate_backend_graph_with_cpu_extra_buffers. It uses the weights only throughLinearModule's matmuls, as the one-shot prefill does, so the checks pass and nothing changesreturn_codes, as mono 16 kHz samples, with their encoder output. An earlier question with the same samples (size and bytes) takes the kept output; the new question is always encoded. A request without history orreturn_codeskeeps nothing and leaves what is kept. A request over the conversation limit fails as before and leaves the kept questions as they werelfm2_audio.prefill.resumed_steps(timing) andlfm2_audio.session.reused_questions(trace). The backbone runtime getsresumed_prefill_steps(), the S2S sessionresumed_prefill_steps()andreused_questions()lfm2_audio_backbone_testchecks restored prefills against a fresh runtime, bit for bit, for prompts of 300 to 1000 steps at 1, 4 and 8 threads and with caches of other lengths, and what stops a resume.lfm2_audio_chat_testruns a four-turn conversation in one session, offline and streamed, against fresh sessions, then its fourth turn again after another conversation, over the limit and after a request without history orreturn_codes.lfm2_audio_encoder_testnow compares bit for bit, as the reuse needs the same bits for the same samplesWhat changes for users
return_codesrun as beforereturn_codes, it encodes its earlier questions again but still restores its prefill blocks. A turn that leaves out the oldest turns runs its prefill in full, though its questions still take the kept encoder output. A question sent back as other samples than its turn had is encoded again, and the prefill runs again from the block it starts in; a float WAV holds the samples exactlyTesting
Three trees built from main at c0d84e1, #838's merge: main, main with the first two commits, and this PR. This PR is the same four patches rebased onto main at adedaf2, which since c0d84e1 adds only ef860c9 (the encoder's CPU repack check moves to an extended test) and a whistle_asr change. On each machine, each tree ran the same ten-turn English conversation in one process, with fixed seeds and
return_codes, sending its own replies back as history, and the ten questions asked alone. Every output file was compared byte for byte: the reply artifact's payload, meta and kind, the text, the audio and the stream events.lfm2_audio.cpu_repackwas on, main's default, unless noted.cpu_repack=false: this PR gives main's bytes in every configuration, and the first two commits do with Q4_0 streamed; 900 of 900 files match over the 14 comparisons. Main and this PR, run twice with Q4_0 streamed, repeat themselvesblock_reducerace fix, which is not merged yet. Without it ggml-cuda's softmax can change a run's bytes, so a difference between trees could come from the race rather than from this PRcpu_extra_buffers_testpass on main and this PR. On this PR rebased onto adedaf2, built with the extended tests, 12 of 12 pass: those,lfm2_audio_encoder_cpu_repack_testandmodel_spec_system_testtest_lfm2_audio_asr,test_lfm2_audio_ttsandtest_lfm2_audio_s2s(CPU) included, and so doescpu_extra_buffers_test; main gives every test the same status.lfm2_audio_chat_testnow takes 3.2 s there instead of 0.6 s, andlfm2_audio_backbone_test1.6 s instead of 1.1 slfm2_audio_cuda_testandcpu_extra_buffers_testpass in all three trees. On the Xeon, a CPU build of this PR without 🚨 fix(ggml-cuda): backport the block_reduce shared-memory race fix (llama.cpp #26385) #837 passes the lfm2_audio unit tests andcpu_extra_buffers_test. In the M3 Ultra's and the A10's lfm2_audio-only builds, 13 tests of other families do not run because their executables do not link, the same on mainTurn 10 of the conversation, streamed, main / first two commits / this PR. Main and the first two commits encode ten questions, this PR one. Replies, and so prompt lengths, differ between machines and packages, so the steps restored differ too. The first two commits were not run with Q8_0 or with
cpu_repack=falseon the i9-13900HK, nor with Q8_0 on the Xeon.cpu_repack=falsecpu_repack=false), 40.0 / 11.5 / 11.8 s on the Xeon with Q4_0, 8.8 / 2.5 / 2.5 s on the M3 Ultra's CPU with Q4_0 (offline), 2.4 / 0.7 / 0.7 s on Metal and 1.30 / 0.46 / 0.46 s on the A10 with Q4_0. The last two commits do not change the prefill, so the differences between the first two commits and this PR there are run-to-run noiseBuild and run