Breeze encoder chunked vram - #431
Conversation
The encoder graph was built at the exact reference-audio length, so conv activations grew linearly (~45 MiB/s of reference) and every new length triggered a full graph rebuild; a 60 s reference cost ~2.5 GB extra over a 6 s one. Split the encoder into two graphs. The conv stack now runs on fixed 5 s chunks (120000 samples) preceded by a 9600-sample left overlap that covers the stack's exact 5240-sample receptive field; chunk lengths are multiples of the 960x transformer stride, so no per-stage right padding occurs and the discarded overlap frames absorb the zero left pads that represent audio start in the first chunk. Stitched outputs are bit-identical to a single-pass encode of the same input (verified over 68 frames x 16 codebooks). The transformer, downsample, and projections run once over the full frame sequence at frame scale, where even minute-long references cost only tens of MiB. Measured on a 2080 Ti (Vulkan, native q8_0 GGUF, peak minus idle baseline): the VRAM slope over reference length drops from ~45 MiB/s to ~11 MiB/s (remaining slope is the frame-scale transformer graph and the longer AR prefill from reference codes), and a 60 s reference peaks ~1.4 GB lower. Encode time for 60 s improves from 3561 ms to 2197 ms.
The transformer graph was rebuilt at the exact frame count for every distinct reference length. Round the capacity up to 125-frame (5 s) buckets so lengths within a bucket share one graph. Unused bucket frames are replicate-padded to match the downsample conv's Replicate right pad; causal attention keeps padding frames invisible to real frames. Verified bit-identical reference codes vs exact-length graphs at 6 s and 15 s; odd lengths show sub-1% last-frame diffs from flash-attention tiling, the same accepted noise class as the pre-existing length sensitivity. Single-run peak VRAM is unchanged.
Vulkan previously paid a cast round trip (f32->bf16->f32, two kernels, a bf16 intermediate tensor) at every activation-rounding point of the breeze decoder. Add a round_bf16 compute shader (f32/f16/bf16 in, always f32 out, round-to-nearest-even via the same fp32_to_bf16 bit trick the cpy shaders use), register pipelines indexed by source type, handle the widened f32 dst in the unary pipeline selection and op-support checks, and enable fused_round for Vulkan in the breeze activation-cast policy. Verified bit-identical breeze reference codes vs the cast round trip at 6 s and 15 s references. Peak VRAM on a 2080 Ti drops ~250 MiB at a 60 s reference (5491 -> 5239 MiB); no measurable change at 6 s.
|
@IIIIIllllIIIIIlllll On my machine, this PR causes regressions: (1) the CUDA mispronunciation issue seems to be back, and (2) Vulkan produces noise. Can you check? |
No problem. |
The breeze activation-rounding policy admits row-strided views into ggml_round_bf16 (ggml_is_contiguous_rows gate in qwen_decoder). The Vulkan port dispatched every input to the flat shader, which indexes the source as a contiguous array, so row-strided views read garbage and clone output degenerated into noise. Route non-contiguous inputs to a new round_bf16_strided shader built on generic_unary_head (same pattern as sigmoid_strided), keeping the flat fast path for contiguous inputs.
|
@0xShug0 Both issues investigated with your exact commands. Findings: (2) Vulkan noise — real bug, root-caused and fixed in 3215194. (1) CUDA mispronunciation — not a regression from this PR, as far as I can tell. On my 2080 Ti with your exact command (q8_0 native load, 60 s ref, seed 42) the coin flip lands the other way: the dev build says "audio CPE" and the PR build says "audio.cpp" correctly. Encoder reference codes at 60 s differ dev-vs-PR in 5.0% of code elements — below the old implementation's own length sensitivity (7.1% for the same audio encoded at different reference lengths, no PR code involved). The bf16 activation rounding makes the AR trajectory chaotic, so any last-ulp encoder difference can flip marginal tokens like "audio.cpp" in either direction per machine. Happy to dig further if you see a systematic (multi-seed) degradation rather than single-sample flips. Could you re-test Vulkan on 3215194? |
|
Follow-up with one more data point: I ran the official PyTorch Breeze-TTS-2 implementation (eager, same inputs, seed 42) as ground truth on both scenarios:
So on the CUDA scenario the official reference itself produces the "mispronounced" variant, and each build flips this marginal token depending on machine and last-ulp noise. There is no systematic direction to the divergence — it is the known bf16-trajectory sensitivity, not a regression introduced here. |
|
@IIIIIllllIIIIIlllll Looks good to me and ready to merge. I will just consider pronunciation instability as the model issue.
|
2 participants
Make clone reference VRAM (near-)constant instead of O(reference length)
Changes
speech_encoder.cpp): fixed 5 schunks + 9600-sample overlap covering the exact 5240-sample receptive
field; chunk boundaries align the 960x transformer stride, so stitched
outputs are bit-identical to a single-pass encode. The transformer still
runs once at frame scale.
replicate-padded): one graph per bucket, no rebuild churn in
long-lived sessions.
round_bf16.comp): singlekernel replacing the f32→bf16→f32 cast round trip, numerically identical.
Results (2080 Ti, q8_0, clone, peak VRAM)
VRAM slope vs reference length: ~45 MiB/s → ~10 MiB/s.
Testing
implementation.
retention.
AI usage: Kimi K3
If you have time, could you please take a look and see if this improvement direction is feasible? Thanks~