[WIP] ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact) - #172
Draft
zhudianGG wants to merge 1 commit into
Draft
[WIP] ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact)#172zhudianGG wants to merge 1 commit into
zhudianGG wants to merge 1 commit into
Conversation
Two perf-only refactors to MingAudioEncoder, both verified bit-identical to the per-clip path in bf16 (the encoder runtime dtype): - Batch conv1/conv2 across clips (one kernel launch instead of one per clip). Right-pad each mel to the max length and zero conv1's padded tail before conv2 so each clip's real region is unaffected by its neighbours' padding; slice each clip's output back off afterwards. - Compute max_seqlen once from the Python length list in forward() and thread it into every block's attention, instead of a per-block (cu_seqlens...).max().item() that forced a GPU->CPU sync in each of the 32 layers. Direct callers that omit it fall back to the old computation. Adds one regression test to the existing encoder test file (no new file) asserting the batched stem is torch.equal to a per-clip reference in bf16.
zhudianGG
marked this pull request as draft
July 16, 2026 23:20
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.
Follow-up perf optimization stacked on #115 (
noah_ming_mstar_port) — targets that branch sinceaudio_encoder.pydoesn't exist onmainyet.Two perf-only refactors to
MingAudioEncoder, both verified bit-identical to the current per-clip path in bf16 (the encoder's runtime dtype):1. Hoist
max_seqlenout of the per-block attentionThe packed-attention flash path computed
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()inside every block — a GPU→CPU sync in each of the 32 layers per audio forward.max_seqlenis constant across blocks, andforward()already has the per-clip lengths as a plain Python list, so it's now computed once (no GPU touch) and threaded down. Direct callers that omit it fall back to the old computation.2. Batch the conv stem across clips
conv1/conv2ran in a per-clip Python loop. Now the clips are right-padded to the max length and run in one batched conv (one kernel launch instead of B).conv1has a bias, so its output over the zero-padded tail isgelu(bias) != 0; that tail is explicitly zeroed beforeconv2(stride 2) so no clip's last frame reads a neighbour's padding — then each clip's real region is sliced back off. This is what keeps it exact.Skipped on purpose: the projector's per-clip loop (
submodules.py) uses a downsampling Conv1d and would need the same masking plus stride-aware output-length handling for a much smaller gain — not worth the extra code.Verification
torch.equalto a per-clip reference in bf16; fp32 differs only ~3e-7 relative (conv-kernel algorithm selection, sub-ULP — well below bf16 resolution).#1only changes wheremax_seqlenis computed, not its value.Tests
One regression method added to the existing
test_ming_flash_omni_encoders.py(no new test file) asserting the batched stem matches per-clip encoding in bf16.Checklist
ruff check .passes