Repository navigation
[Model] Nemotron-3.5-Transcribe: pack the audio tower's valid frames for MMEncoderAttention - #64944
Merged
Isotr0py merged 3 commits intoNov 2, 2026
Conversation
…EncoderAttention The ASR and speaker encoders padded every clip of a batch to the longest one and masked the padded keys in SDPA. They now run only each row's valid frames, packed row after row, and attend through MMEncoderAttention with cu_seqlens, so vLLM picks the variable-length kernel of the platform. The rotary embedding reads each frame's position in its row, rows without valid frames are skipped, and the padded positions of the output are zero. The lengths come from the host, so the packing needs no device-to-host copy. The tower tests run on the platform's device, where vLLM registers the encoder attention ops. Signed-off-by: Dongji Gao <dongjig@nvidia.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
With the tower on MMEncoderAttention, mm_encoder_attn_dtype='fp8' reaches its attention. vLLM's encoder attention refuses it on most backends with its own error. Where it accepts it, as on FlashInfer cuDNN, the layers add FP8 scale buffers that no checkpoint has, so loading fails, and the packed cu_seqlens follow the BF16 layout. The old SDPA path ignored the setting. The tower now refuses it with a clear error on every backend when it is built; the FP8 checkpoints quantize the tower's linear layers without it. The model's note in the supported models list says so, and that the encoder attention backend and the other clips in a batch can change the tower's bfloat16 output. Signed-off-by: Dongji Gao <dongjig@nvidia.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Contributor
|
Documentation preview: https://vllm--64944.org.readthedocs.build/en/64944/ |
mergify
Bot
added
documentation
DongjiGao
marked this pull request as ready for review
November 1, 2026 05:50
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
6 tasks done
Isotr0py
approved these changes
Nov 1, 2026
Member
|
/ci run |
|
✅ @DongjiGao, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #98996 for commit ℹ️ This PR is 3 commits behind upstream |
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.
Overview
Follow-up to #64282, from @Isotr0py's review there ("Can we use MMEncoderAttention with cu_seqlens for FA?"). Both
encoders of the Nemotron-3.5-Transcribe audio tower, the ASR encoder and the Sortformer speaker encoder, now run only
each row's valid frames, packed row after row, and attend through
MMEncoderAttentionwithcu_seqlensinstead ofpadded, masked SDPA. The PR changes the model file, its tests and its docs note only.
Claims
(FlashAttention on the tested GPUs), and
--mm-encoder-attn-backendapplies to the tower.6000, 1.51× and 2.15× on H100. On H100 it is slower on 8 clips and on 4 long clips (0.85× and 0.84×). Offline
end-to-end RTFx on the RTX PRO 6000 quick sets does not change.
speaker attribution the packed tower changes in every run (six runs per arm, with GPUs swapped and request order
shuffled); Earnings22 WER rises by 0.32 points from a few looping recordings whose link to the tower change is not
known.
mm_encoder_attn_dtype='fp8'is refused with a clear error. The FP8 and NVFP4 checkpoints quantize the linearlayers of the tower's ASR encoder without this setting and are not affected.
Validation
Setups:
nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4Bat the released revision525be0e; the quick sets ran earlier withthe pre-release revision
f6e8392of the same repository. The exactness and accuracy check, the batch-dependenceprobe and the tower timing ran this branch; the quick and long sets ran the same model-file changes on pre-merge
builds of [Model] Add Nemotron-3.5-Transcribe #64282 (the long-set build, in both arms, also carries follow-up commits outside the tower).
0.32.1rc1.dev326+gc722f3a53with this PR's tower code (the sameexecutable code; only the license header differs) and torch 2.14.0+cu130. The exactness and accuracy check used
weights byte-identical to the released checkpoint; the tower timing and batch-dependence probes used the
pre-release GA-RC1 BF16 weights (SHA-256
fe3b898b42a65837180c65dedd2e62b92836cee704ae3cbee2f473946fc7d0c9).Tests (this branch on main
e0d00e4b8e; Python 3.12.14, torch 2.14.0+cu132, pytest 9.1.1):The new tests fail without their code: the empty-batch test with the guard removed (vLLM's SDPA wrapper raises
torch.cat(): expected a non-empty list of Tensors), and the FP8 test with the refusal removed (vLLM's generic FP8check raises a different message). The tiny tower is float32 and selects TORCH_SDPA, so the tests do not cover
FlashAttention; the real-tower measurements below do. CI runs the file in the CPU Multimodal Processor step.
Exactness and accuracy (each clip encoded alone; 64 short and 4 long clips; the released checkpoint). With
float32 and vLLM's TORCH_SDPA backend, the packed tower is bit-identical to main's padded tower for 68 of 68 clips at
the ASR states, speaker probabilities and embeddings, on both the RTX PRO 6000 and H100. In BF16, against that
float32 reference, both towers have the same typical error (ASR relative L2 median about 7e-3 on short clips on both
GPUs). On the four long clips, FlashAttention 2 (RTX PRO 6000) matches main's error (788 against 843 flipped speaker
frames of 4,314), while FlashAttention 3 (H100) has somewhat larger worst-case errors (ASR maximum 0.043 against
0.036; 358 against 316 flipped frames).
Tower speed (BF16, tower only, median of 10 alternating runs):
Offline end-to-end RTFx does not change (RTX PRO 6000, quick sets, English / multilingual: 386 / 428 packed, 390 /
426 padded): the tower does not limit serving throughput.
Accuracy, quick sets (offline; Open ASR leaderboard English, 18,341 rows scored as 18,006 entries; multilingual,
15,486 clips):
Accuracy, long sets (multi-speaker SOT, 164 sessions, cpWER; long-form, 451 recordings, WER with the loop
guardrail; padded and packed served side by side, each run twice):
CORAAL's mean moves by +0.04 points, against repeat differences of 1.19 (padded) and 1.23 (packed); Earnings21 moves
by −0.07. AMI-IHM's +1.73 comes from one of its 16 meetings (EN2002c): in both packed runs, about 800 words of one
speaker go to an extra speaker tag. The words are unchanged (WER without speaker tags moves by +0.07), so this is
speaker attribution, not recognition; an earlier H100 comparison with different weights moved the same meeting the
other way. Earnings22's +0.32 comes from nine recordings where generation loops or reaches the token limit, mostly as
deleted text; the other 116 recordings improve slightly. On H100 (one run per arm, an earlier checkpoint and chat
template), AMI-IHM and Earnings22 improved instead, while Mixer6 and three CORAAL locations got worse. A crossover on
SOT (four more runs per arm, with GPUs swapped and request order shuffled) reproduced the EN2002c split in all six
packed runs (about 915 words under the extra tag) and in none of the six padded runs (about 115 words), so it follows
the tower, not the GPU or the batch order. The other 15 AMI-IHM meetings are equal or slightly better with the packed
tower in every pair. Earnings22 was not part of the crossover. TED-LIUM (11 recordings) varies too much between runs
of the same code to compare (up to 5.9 points).
Batch dependence (BF16). A clip's tower output can depend on the other clips in its batch, before and after this
PR: rounding changes with the batch, the speaker encoder's probabilities move, and the thresholded speaker features
flip on some frames. Batched against alone:
On the four long clips on H100, forcing SDPA removed the additional speaker-feature variation of FlashAttention 3,
and the ASR encoder's output stayed identical with FlashAttention 3. Forcing SDPA does not make the tower
batch-invariant: on H100, half of the short clips change with the batch in all three towers (embeddings up to 0.033
relative L2 apart, no speaker flips), and on the RTX PRO 6000 all three towers flip speaker frames on three of the
four long clips. FlashAttention therefore stays the default. Forced SDPA on the RTX PRO 6000 takes 23.4, 63.8, 113.0
and 100.3 ms on the four batches of the speed table.
Details
TransformerEncoder.packindexes each row's valid frames from host lengths and buildscu_seqlenswith
MMEncoderAttention's helpers for the selected backend. Rotary embeddings use each frame's position in itsrow. Rows without valid frames are left out, and a batch without any valid frame skips the layers, because vLLM's
SDPA wrapper cannot split an empty
cu_seqlens. The padded positions of the output are zero. The Sortformer stepadds its speaker-cache and FIFO context to the host lengths as to the device lengths.
TORCH_SDPA, vLLM's SDPA wrapper still copies the cumulative lengths to the host once per layer(
vllm/v1/attention/ops/vit_attn_wrappers.py).mm_encoder_attn_dtype='fp8'only on the FlashInfercuDNN backend and on ROCm AITER. There, the tower's attention layers would add FP8 scale buffers that no checkpoint
has, so loading fails, and the FlashInfer path expects a different
cu_seqlenslayout. The tower therefore refusesthe setting on every backend when it is built. The padded SDPA path ignored it.
docs/models/supported_models.mdstates the refusal and that the encoder attentionbackend and the other clips in a batch can change the tower's BF16 output.
serving throughput does not change.
MMEncoderAttentionhas no per-layer backend option, so SDPA in the speakerencoder alone would need a model-specific override. Measured on RTX PRO 6000 and H100 only, not on ROCm or other
GPUs.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
AI assistance (Cursor agent) was used to prepare this change. Dongji Gao reviewed every changed line and the
validation.