From b6930f60f637c88f66d9d6de26872f38c7c4a93c Mon Sep 17 00:00:00 2001 From: Dongji Gao Date: Mon, 26 Oct 2026 10:11:29 -0700 Subject: [PATCH 1/2] [Model] Nemotron-3.5-Transcribe: pack the tower's valid frames for MMEncoderAttention 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 Co-authored-by: Cursor Agent --- .../processing/test_nemotron3_5_transcribe.py | 55 +++++- .../models/nemotron3_5_transcribe_audio.py | 160 ++++++++++++++---- 2 files changed, 182 insertions(+), 33 deletions(-) diff --git a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py index cba7d95ac931..22c07c5308c7 100644 --- a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py +++ b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py @@ -240,9 +240,15 @@ def tiny_tower(dist_init): from vllm.model_executor.models.nemotron3_5_transcribe_audio import ( Nemotron3_5TranscribeAudioTower, ) + from vllm.platforms import current_platform torch.manual_seed(0) - tower = Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG).eval() + # vLLM registers the encoder attention ops for the platform's device only. + tower = ( + Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG) + .eval() + .to(current_platform.device_type) + ) with torch.no_grad(): # vLLM's linear layers allocate their weights uninitialized. for module in tower.modules(): @@ -458,3 +464,50 @@ def test_audio_tower_encodes_each_clip_as_if_alone(tiny_tower): ] for got, expected in zip(together, alone): torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4) + + +def test_encoder_packs_each_rows_valid_frames(tiny_tower): + encoder = tiny_tower.encoder.asr_encoder + packed = encoder.pack(torch.tensor([2, 0, 3]), 4, torch.device("cpu")) + assert packed.rows.tolist() == [0, 1, 8, 9, 10] + assert packed.positions.tolist() == [0, 1, 0, 1, 2] + assert packed.cu_seqlens.tolist() == [0, 2, 5] + + +def test_encoder_runs_only_the_valid_frames(tiny_tower): + encoder = tiny_tower.encoder.asr_encoder + device = encoder.final_norm.weight.device + torch.manual_seed(2) + x = torch.randn(3, 6, encoder.d_model, device=device) + lengths = torch.tensor([6, 0, 3]) + with torch.inference_mode(): + states, _ = encoder( + x, lengths.to(device), bypass_pre_encode=True, lengths_cpu=lengths + ) + alone = [ + encoder( + x[row : row + 1, :length], + lengths[row : row + 1].to(device), + bypass_pre_encode=True, + )[0][0] + for row, length in ((0, 6), (2, 3)) + ] + assert not states[1].any() + assert not states[2, :, 3:].any() + torch.testing.assert_close(states[0], alone[0], rtol=1e-4, atol=1e-4) + torch.testing.assert_close(states[2, :, :3], alone[1], rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("name", ["asr_encoder", "diarization_model.encoder"]) +@pytest.mark.parametrize("num_frames", [0, 6]) +def test_encoder_returns_zeros_when_no_row_has_frames(tiny_tower, name, num_frames): + encoder = tiny_tower.encoder.get_submodule(name) + device = encoder.final_norm.weight.device + x = torch.randn(3, num_frames, encoder.d_model, device=device) + lengths = torch.zeros(3, dtype=torch.int64) + with torch.inference_mode(): + states, _ = encoder( + x, lengths.to(device), bypass_pre_encode=True, lengths_cpu=lengths + ) + assert states.shape == (3, encoder.d_model, num_frames) + assert not states.any() diff --git a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py index 47552bca867f..32e15c1e3aa3 100644 --- a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py +++ b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py @@ -21,10 +21,14 @@ from dataclasses import dataclass from typing import Any, cast +import numpy as np import torch import torch.nn.functional as F from torch import nn +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, +) from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.utils.torch_utils import async_tensor_h2d @@ -261,7 +265,26 @@ def forward( x = x.reshape( batch_size, (num_frames + pad) // self.factor, channels * self.factor ) - return self.proj(x), (lengths + self.factor - 1) // self.factor + return self.proj(x), self.output_lengths(lengths) + + def output_lengths(self, lengths: torch.Tensor) -> torch.Tensor: + return (lengths + self.factor - 1) // self.factor + + +@dataclass +class PackedFrames: + """The valid frames of a padded ``(B, T)`` batch, packed row after row. + + ``rows`` indexes the flattened batch and ``positions`` holds each frame's + position in its row. ``cu_seqlens``, ``max_seqlen`` and ``sequence_lengths`` + are in the format of the attention backend. + """ + + rows: torch.Tensor + positions: torch.Tensor + cu_seqlens: torch.Tensor + max_seqlen: torch.Tensor + sequence_lengths: torch.Tensor | None class MultiHeadAttention(nn.Module): @@ -295,25 +318,37 @@ def __init__( if self.qk_norm: self.q_norm = nn.LayerNorm(self.head_dim) self.k_norm = nn.LayerNorm(self.head_dim) + self.attn = MMEncoderAttention( + num_heads=self.n_heads, + head_size=self.head_dim, + prefix=f"{prefix}.attn", + ) def forward( self, x: torch.Tensor, - attn_mask: torch.Tensor, + packed: PackedFrames, cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - batch_size, num_frames, _ = x.shape - qkv = self.w_qkv(x).view(batch_size, num_frames, 3, self.n_heads, self.head_dim) - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + """Attend within each row of the packed ``(N, d_model)`` frames.""" + num_frames = x.shape[0] + qkv = self.w_qkv(x).view(num_frames, 3, self.n_heads, self.head_dim) + q, k, v = qkv.unbind(1) if self.qk_norm: q = self.q_norm(q).to(v.dtype) k = self.k_norm(k).to(v.dtype) q = self.rope.apply(q, cos.to(q.dtype), sin.to(q.dtype)) k = self.rope.apply(k, cos.to(k.dtype), sin.to(k.dtype)) - out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) - out = out.transpose(1, 2).reshape(batch_size, num_frames, self.d_model) - return self.out_proj(out) + out = self.attn( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + cu_seqlens=packed.cu_seqlens, + max_seqlen=packed.max_seqlen, + sequence_lengths=packed.sequence_lengths, + ) + return self.out_proj(out.reshape(num_frames, self.d_model)) class TransformerBlock(nn.Module): @@ -354,11 +389,11 @@ def __init__( def forward( self, x: torch.Tensor, - attn_mask: torch.Tensor, + packed: PackedFrames, cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), attn_mask, cos, sin) + x = x + self.attn(self.norm1(x), packed, cos, sin) return x + self.ffn.net(self.norm2(x)) @@ -367,8 +402,10 @@ class TransformerEncoder(nn.Module): ``forward`` takes ``(B, C, T)`` mel features, or ``(B, T, d_model)`` pre-encoded embeddings with ``bypass_pre_encode=True``, and returns - ``(B, d_model, T')`` states with their lengths. Only keys are masked, so - padded frames never influence valid ones. + ``(B, d_model, T')`` states with their lengths. Only each row's valid frames + run through the layers, packed for variable-length attention; the padded + positions of the output are zero. ``lengths_cpu``, the lengths on the host, + lets the packing skip a device-to-host copy. """ def __init__( @@ -401,28 +438,63 @@ def __init__( ] ) self.final_norm = nn.LayerNorm(self.d_model) + self.attn_backend = self.layers[0].attn.attn.attn_backend def forward( - self, x: torch.Tensor, lengths: torch.Tensor, bypass_pre_encode: bool = False + self, + x: torch.Tensor, + lengths: torch.Tensor, + bypass_pre_encode: bool = False, + lengths_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if not bypass_pre_encode: x, lengths = self.pre_encode(x, lengths) + if lengths_cpu is not None: + lengths_cpu = self.pre_encode.output_lengths(lengths_cpu) lengths = lengths.to(torch.int64) - if self.xscale: - x = x * self.xscale - x = self.embed_norm(x) - num_frames = x.shape[1] - positions = torch.arange(num_frames, device=x.device).unsqueeze(0) - # A row with no valid frame (a short clip in a later window) attends to - # frame 0 to stay finite; its outputs are discarded. - key_mask = (positions < lengths.unsqueeze(1)) | ( - (positions == 0) & (lengths.unsqueeze(1) == 0) - ) - attn_mask = key_mask[:, None, None, :] - cos, sin = self.rope.cos_sin(num_frames, x.dtype, x.device) - for layer in self.layers: - x = layer(x, attn_mask, cos, sin) - return self.final_norm(x).transpose(1, 2), lengths + if lengths_cpu is None: + lengths_cpu = lengths.cpu() + batch_size, num_frames, width = x.shape + out = x.new_zeros(batch_size * num_frames, width) + packed = self.pack(lengths_cpu, num_frames, x.device) + if packed.rows.numel(): + h = x.reshape(-1, width).index_select(0, packed.rows) + if self.xscale: + h = h * self.xscale + h = self.embed_norm(h) + cos, sin = self.rope.cos_sin(num_frames, x.dtype, x.device) + cos = cos.index_select(0, packed.positions).unsqueeze(1) + sin = sin.index_select(0, packed.positions).unsqueeze(1) + for layer in self.layers: + h = layer(h, packed, cos, sin) + out.index_copy_(0, packed.rows, self.final_norm(h)) + return out.view(batch_size, num_frames, width).transpose(1, 2), lengths + + def pack( + self, lengths_cpu: torch.Tensor, num_frames: int, device: torch.device + ) -> PackedFrames: + """Index each row's valid frames; rows without any are left out.""" + lengths_cpu = lengths_cpu.to(torch.int64).clamp(min=0, max=num_frames) + valid = torch.arange(num_frames) < lengths_cpu.unsqueeze(1) + rows = valid.flatten().nonzero().squeeze(1) + positions = torch.arange(num_frames).expand_as(valid)[valid] + seq_lens = lengths_cpu[lengths_cpu > 0].numpy() + cu_seqlens = np.concatenate(([0], np.cumsum(seq_lens))).astype(np.int32) + backend = self.attn_backend + return PackedFrames( + rows=async_tensor_h2d(rows, device=device), + positions=async_tensor_h2d(positions, device=device), + cu_seqlens=MMEncoderAttention.maybe_recompute_cu_seqlens( + backend, cu_seqlens, self.d_model, 1, device + ), + max_seqlen=torch.tensor( + MMEncoderAttention.compute_max_seqlen(backend, cu_seqlens), + dtype=torch.int32, + ), + sequence_lengths=MMEncoderAttention.maybe_compute_seq_lens( + backend, cu_seqlens, device + ), + ) def downsample_preds(preds: torch.Tensor, factor: int) -> torch.Tensor: @@ -718,6 +790,7 @@ def step( state: SpeakerCacheState, left_offset: int, right_offset: int, + lengths_cpu: torch.Tensor | None = None, ) -> tuple[SpeakerCacheState, torch.Tensor]: """Process one window of ``(B, C, T)`` mel features. @@ -730,7 +803,14 @@ def step( embs = torch.cat([state.spkcache, state.fifo, chunk_embs], dim=1) context_len = state.spkcache.shape[1] + state.fifo.shape[1] emb_lengths = context_len + chunk_lengths - states, state_lengths = self.encoder(embs, emb_lengths, bypass_pre_encode=True) + emb_lengths_cpu = ( + None + if lengths_cpu is None + else context_len + self.encoder.pre_encode.output_lengths(lengths_cpu) + ) + states, state_lengths = self.encoder( + embs, emb_lengths, bypass_pre_encode=True, lengths_cpu=emb_lengths_cpu + ) hidden = modules.encoder_proj(states.transpose(1, 2)) high_res_preds = modules.speaker_probs(hidden, state_lengths) @@ -818,7 +898,11 @@ def dtype(self) -> torch.dtype: return self.asr_norm.weight.dtype def forward( - self, mel: torch.Tensor, mel_lengths: torch.Tensor, max_feat_len: int + self, + mel: torch.Tensor, + mel_lengths: torch.Tensor, + max_feat_len: int, + mel_lengths_cpu: torch.Tensor | None = None, ) -> torch.Tensor: total_feat_len = min(mel.shape[-1], max_feat_len) num_windows = max(1, math.ceil(total_feat_len / self.window_feat_len)) @@ -837,9 +921,18 @@ def forward( window_lengths = (mel_lengths - context_start).clamp( min=0, max=context_end - context_start ) + window_lengths_cpu = ( + None + if mel_lengths_cpu is None + else (mel_lengths_cpu - context_start).clamp( + min=0, max=context_end - context_start + ) + ) encoded, _ = self.asr_encoder( - asr_signal[:, :, context_start:context_end], window_lengths + asr_signal[:, :, context_start:context_end], + window_lengths, + lengths_cpu=window_lengths_cpu, ) left_drop = left_offset // self.subsampling_factor core_len = self._output_boundary(end) - self._output_boundary(start) @@ -852,6 +945,7 @@ def forward( state, left_offset, right_offset, + lengths_cpu=window_lengths_cpu, ) speaker_windows.append(align_frames(window_preds, core_len)) @@ -961,7 +1055,9 @@ def forward( mel_lengths_cpu = self.preprocessor.featurizer.get_seq_len( num_samples_cpu ).clamp(max=mel.shape[-1]) - encoded = self.encoder(mel, mel_lengths, int(mel_lengths_cpu.max())) + encoded = self.encoder( + mel, mel_lengths, int(mel_lengths_cpu.max()), mel_lengths_cpu + ) encoded_lengths = ( mel_lengths_cpu + self.subsampling_factor - 1 ) // self.subsampling_factor From f50ff608c403d3fbddf00da5567c2abfa0c7b016 Mon Sep 17 00:00:00 2001 From: Dongji Gao Date: Mon, 26 Oct 2026 12:58:19 -0700 Subject: [PATCH 2/2] [Model] Nemotron-3.5-Transcribe: refuse FP8 encoder attention 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 Co-authored-by: Cursor Agent --- docs/models/supported_models.md | 7 +++++++ .../processing/test_nemotron3_5_transcribe.py | 18 ++++++++++++++++++ .../models/nemotron3_5_transcribe_audio.py | 9 +++++++++ 3 files changed, 34 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9d2210c69c3e..38f088780c4b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -776,6 +776,13 @@ Speech2Text models trained specifically for Automatic Speech Recognition. changes it), and `/v1/audio/transcriptions` splits longer audio into pieces of at most one hour, shorter when `--max-model-len` is below 65,536. + Audio tower: the FP8 and NVFP4 checkpoints quantize the linear layers of its ASR encoder to FP8 through + their `quantization_config`. FP8 encoder attention (`--mm-encoder-attn-dtype fp8`) is not supported and + raises an error when the model is built. The tower attends through vLLM's multimodal encoder attention + backend (FlashAttention on CUDA GPUs that support it; `--mm-encoder-attn-backend` selects another). In + bfloat16, its output can depend on that backend and on the other clips in the batch: small rounding + differences can flip the thresholded speaker features on some frames. + !!! note `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed. diff --git a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py index 22c07c5308c7..d527d724a004 100644 --- a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py +++ b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py @@ -511,3 +511,21 @@ def test_encoder_returns_zeros_when_no_row_has_frames(tiny_tower, name, num_fram ) assert states.shape == (3, encoder.d_model, num_frames) assert not states.any() + + +@pytest.mark.usefixtures("dist_init") +def test_fp8_encoder_attention_is_refused(tmp_path): + from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config + from vllm.model_executor.models.nemotron3_5_transcribe_audio import ( + Nemotron3_5TranscribeAudioTower, + ) + + (tmp_path / "config.json").write_text(json.dumps(CHECKPOINT_CONFIG)) + model_config = ModelConfig(str(tmp_path), mm_encoder_attn_dtype="fp8") + # vLLM's encoder attention refuses FP8 with another message, except on the + # FlashInfer cuDNN and ROCm AITER backends; the tower refuses it on all. + with ( + set_current_vllm_config(VllmConfig(model_config=model_config)), + pytest.raises(ValueError, match=r"Nemotron-3\.5-Transcribe does not support"), + ): + Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG) diff --git a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py index 32e15c1e3aa3..65b54f6d52a0 100644 --- a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py +++ b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py @@ -34,6 +34,7 @@ from vllm.utils.torch_utils import async_tensor_h2d from .utils import maybe_prefix +from .vision import get_multimodal_config # Checkpoint tensors that the inference path does not use. The speaker # encoder's own preprocessor is unused because both encoders read the mel @@ -984,6 +985,14 @@ def __init__( prefix: str = "", ) -> None: super().__init__() + mm_config = get_multimodal_config() + if mm_config is not None and mm_config.mm_encoder_attn_dtype == "fp8": + raise ValueError( + "Nemotron-3.5-Transcribe does not support " + "mm_encoder_attn_dtype='fp8': the audio tower's attention runs in " + "the model dtype. The FP8 checkpoints quantize the tower's linear " + "layers without this setting." + ) self.preprocessor = AudioPreprocessor(audio_config["preprocessor"]) self.encoder = ParallelExpertEncoder( audio_config, quant_config, maybe_prefix(prefix, "encoder")