A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 08:45 UTC
Remove `head_mask` from Ultravox and Swin (#30764)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Dec 16, 2025 at 16:02 UTC
0b0acc758ed3f0eecd8d95b3e232f8dd91bb8473
2 files changed
+16
-17
vllm/model_executor/models/swin.py
+1
-15
@@ -102,7 +102,6 @@ class SwinSelfAttention(nn.Module):
102
self,
103
hidden_states: torch.Tensor,
104
attention_mask: torch.FloatTensor | None = None,
105
- head_mask: torch.FloatTensor | None = None,
105
output_attentions: bool | None = False,
106
) -> tuple[torch.Tensor, ...]:
107
batch_size, dim, num_channels = hidden_states.shape
@@ -201,12 +200,9 @@ class SwinAttention(nn.Module):
200
self,
201
hidden_states: torch.Tensor,
202
attention_mask: torch.FloatTensor | None = None,
204
- head_mask: torch.FloatTensor | None = None,
203
output_attentions: bool | None = False,
204
) -> tuple[torch.Tensor]:
207
- self_outputs = self.self(
208
- hidden_states, attention_mask, head_mask, output_attentions
209
- )
205
+ self_outputs = self.self(hidden_states, attention_mask, output_attentions)
206
attention_output = self.output(self_outputs[0], hidden_states)
207
outputs = (attention_output,) + self_outputs[1:]
208
return outputs
@@ -339,18 +335,14 @@ class SwinStage(nn.Module):
335
self,
336
hidden_states: torch.Tensor,
337
input_dimensions: tuple[int, int],
342
- head_mask: torch.FloatTensor | None = None,
338
output_attentions: bool | None = False,
339
always_partition: bool | None = False,
340
) -> tuple[torch.Tensor]:
341
height, width = input_dimensions
342
for i, layer_module in enumerate(self.blocks):
348
- layer_head_mask = head_mask[i] if head_mask is not None else None
349
-
343
layer_outputs = layer_module(
344
hidden_states,
345
input_dimensions,
353
- layer_head_mask,
346
output_attentions,
347
always_partition,
348
)
@@ -425,17 +417,13 @@ class SwinEncoder(nn.Module):
417
self,
418
hidden_states: torch.Tensor,
419
input_dimensions: tuple[int, int],
428
- head_mask: torch.FloatTensor | None = None,
420
output_attentions: bool | None = False,
421
always_partition: bool | None = False,
422
) -> tuple[torch.Tensor]:
423
for i, layer_module in enumerate(self.layers):
433
- layer_head_mask = head_mask[i] if head_mask is not None else None
434
-
424
layer_outputs = layer_module(
425
hidden_states,
426
input_dimensions,
438
- layer_head_mask,
427
output_attentions,
428
always_partition,
429
)
@@ -473,7 +461,6 @@ class SwinModel(nn.Module):
461
def forward(
462
self,
463
pixel_values: torch.FloatTensor | None = None,
476
- head_mask: torch.FloatTensor | None = None,
464
output_attentions: bool | None = None,
465
) -> tuple[torch.Tensor]:
466
embedding_output, input_dimensions = self.embeddings(pixel_values)
@@ -481,7 +468,6 @@ class SwinModel(nn.Module):
468
encoder_outputs = self.encoder(
469
embedding_output,
470
input_dimensions,
484
- head_mask=head_mask,
471
output_attentions=output_attentions,
472
)
473
vllm/model_executor/models/ultravox.py
+15
-2
@@ -5,6 +5,7 @@
5
"""PyTorch Ultravox model."""
6
7
import copy
8
+import inspect
9
from collections.abc import Iterable, Mapping, Sequence
10
from types import SimpleNamespace
11
from typing import Annotated, Any, Literal, TypeAlias
@@ -380,11 +381,17 @@ class UltravoxTransformerProjector(nn.Module, ModuleUtilsMixin):
381
)
382
hidden_states = hidden_states + positions
383
384
+ # Backward compatibility for Transformers v4 where layer_head_mask
385
+ # was a required argument for WhisperEncoderLayer.forward
386
+ kwargs = {}
387
+ if "layer_head_mask" in inspect.signature(self.layers[0].forward).parameters:
388
+ kwargs["layer_head_mask"] = None
389
+
390
for layer in self.layers:
391
layer_outputs = layer(
392
hidden_states,
393
attention_mask=extended_attention_mask,
387
- layer_head_mask=None,
394
+ **kwargs,
395
)
396
hidden_states = layer_outputs[0]
397
@@ -479,11 +486,17 @@ class ModifiedWhisperEncoder(WhisperEncoder):
486
487
attention_mask = self.get_attention_mask_by_audio_len(audio_lens, hidden_states)
488
489
+ # Backward compatibility for Transformers v4 where layer_head_mask
490
+ # was a required argument for WhisperEncoderLayer.forward
491
+ kwargs = {}
492
+ if "layer_head_mask" in inspect.signature(self.layers[0].forward).parameters:
493
+ kwargs["layer_head_mask"] = None
494
+
495
for encoder_layer in self.layers:
496
layer_outputs = encoder_layer(
497
hidden_states,
498
attention_mask,
486
- layer_head_mask=None,
499
+ **kwargs,
500
)
501
502
hidden_states = layer_outputs[0]