Skip to content

[WIP] ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact) - #172

Draft
zhudianGG wants to merge 1 commit into
noah_ming_mstar_portfrom
ming-audio-encoder-perf
Draft

[WIP] ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact)#172
zhudianGG wants to merge 1 commit into
noah_ming_mstar_portfrom
ming-audio-encoder-perf

Conversation

@zhudianGG

Copy link
Copy Markdown
Collaborator

Follow-up perf optimization stacked on #115 (noah_ming_mstar_port) — targets that branch since audio_encoder.py doesn't exist on main yet.

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_seqlen out of the per-block attention

The 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_seqlen is constant across blocks, and forward() 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/conv2 ran 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). conv1 has a bias, so its output over the zero-padded tail is gelu(bias) != 0; that tail is explicitly zeroed before conv2 (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

  • Batched conv stem is torch.equal to a per-clip reference in bf16; fp32 differs only ~3e-7 relative (conv-kernel algorithm selection, sub-ULP — well below bf16 resolution).
  • #1 only changes where max_seqlen is 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
  • Bit-exactness verified (bf16)

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 zhudianGG changed the title ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact) [WIP] ming audio encoder: batch conv stem + hoist max_seqlen (perf, bit-exact) Jul 16, 2026
@zhudianGG
zhudianGG marked this pull request as draft July 16, 2026 23:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant