[Model] Nemotron-3.5-Transcribe: pack the audio tower's valid frames for MMEncoderAttention

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.
 

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,68 @@ 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()
+
+
+@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)

vllm/model_executor/models/nemotron3_5_transcribe_audio.py

@@ -21,15 +21,20 @@
 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
 
 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
@@ -261,7 +266,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 +319,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 +390,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 +403,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 +439,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 +791,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 +804,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 +899,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 +922,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 +946,7 @@ def forward(
                 state,
                 left_offset,
                 right_offset,
+                lengths_cpu=window_lengths_cpu,
             )
             speaker_windows.append(align_frames(window_preds, core_len))
 
@@ -890,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")
@@ -961,7 +1064,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