@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f3f89529d
Fork 0

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 9, 18:15 UTC

[Voxtral] Add new streaming arch (#32861)

Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Patrick von Platen committed Jan 23, 2026 at 12:41 UTC 3f3f89529dc3fbaa5bf22c86d9f0833b49dc76ea
9 files changed +768 -314
tests/models/multimodal/generation/test_voxtral_streaming.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import asdict
4 +
5 +import pytest
6 +from mistral_common.audio import Audio
7 +from mistral_common.protocol.instruct.chunk import RawAudio
8 +from mistral_common.protocol.transcription.request import (
9 + StreamingMode,
10 + TranscriptionRequest,
11 +)
12 +from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
13 +
14 +from vllm import LLM, EngineArgs, SamplingParams
15 +from vllm.assets.audio import AudioAsset
16 +
17 +
18 +def _get_engine(path: str) -> LLM:
19 + engine_args = EngineArgs(
20 + model=path,
21 + max_model_len=8192,
22 + max_num_seqs=1,
23 + limit_mm_per_prompt={"audio": 1},
24 + config_format="mistral",
25 + load_format="mistral",
26 + tokenizer_mode="mistral",
27 + enforce_eager=True,
28 + gpu_memory_utilization=0.4,
29 + )
30 + return LLM(**asdict(engine_args))
31 +
32 +
33 +@pytest.mark.skip(reason="Voxtral streaming is not yet public")
34 +def test_voxtral_streaming_forward():
35 + audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
36 +
37 + model_name = "mistralai/Voxtral-Mini-3B-Realtime-2602"
38 + tokenizer = MistralTokenizer.from_hf_hub(model_name)
39 + audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
40 +
41 + def from_file(file_path: str):
42 + audio = Audio.from_file(file_path, strict=False)
43 + req = TranscriptionRequest(
44 + audio=RawAudio.from_audio(audio),
45 + streaming=StreamingMode.OFFLINE,
46 + language=None,
47 + )
48 + tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
49 +
50 + return (tokenized.tokens, tokenized.audios[0].audio_array)
51 +
52 + tokenized_list = [
53 + from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
54 + ]
55 +
56 + inputs = []
57 + sampling_params = []
58 +
59 + for tokens, audio_array in tokenized_list:
60 + num_samples = audio_array.shape[0]
61 + max_tokens = (
62 + audio_config.num_audio_tokens(num_samples)
63 + - audio_config.num_delay_tokens
64 + - 1
65 + )
66 + sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens))
67 +
68 + input_dict = {
69 + "multi_modal_data": {"audio": [(audio_array, None)]},
70 + "prompt_token_ids": tokens,
71 + }
72 + inputs.append(input_dict)
73 +
74 + llm = _get_engine(model_name)
75 + outputs = llm.generate(
76 + inputs,
77 + sampling_params=sampling_params,
78 + )
79 +
80 + texts = [out.outputs[0].text for out in outputs]
81 + expected = [
82 + (
83 + " First words I spoke in the original phonograph. "
84 + "A little piece of practical poetry. Mary had a little lamb,"
85 + " it sleeps with quite a snow, and everywhere that Mary went, "
86 + "the lamb was sure to go."
87 + ),
88 + (
89 + " And the 0-1 pitch on the way to Edgar Martinez. Swung on"
90 + " the line. Down the left field line for OBS. Here comes Joy. "
91 + "Here is Junior to third base. They're going to wave him in. "
92 + "The throw to the plate will be late. The Mariners are going"
93 + " to play. For the American League Championship, "
94 + "I don't believe it. It just continues. My oh, my."
95 + ),
96 + ]
97 + assert texts == expected
vllm/model_executor/models/llama.py
+4 -1
@@ -404,6 +404,7 @@ class LlamaModel(nn.Module):
404 positions: torch.Tensor,
405 intermediate_tensors: IntermediateTensors | None,
406 inputs_embeds: torch.Tensor | None = None,
407 + **extra_layer_kwargs,
408 ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
409 if get_pp_group().is_first_rank:
410 if inputs_embeds is not None:
@@ -422,7 +423,9 @@ class LlamaModel(nn.Module):
423 ):
424 if idx in self.aux_hidden_state_layers:
425 aux_hidden_states.append(hidden_states + residual)
425 - hidden_states, residual = layer(positions, hidden_states, residual)
426 + hidden_states, residual = layer(
427 + positions, hidden_states, residual, **extra_layer_kwargs
428 + )
429
430 if not get_pp_group().is_last_rank:
431 return IntermediateTensors(
vllm/model_executor/models/mistral.py
+108
@@ -10,6 +10,12 @@ from transformers import LlamaConfig
10
11 from vllm.compilation.decorators import support_torch_compile
12 from vllm.config import CacheConfig, VllmConfig
13 +from vllm.model_executor.layers.activation import SiluAndMul
14 +from vllm.model_executor.layers.linear import (
15 + ColumnParallelLinear,
16 + MergedColumnParallelLinear,
17 + RowParallelLinear,
18 +)
19 from vllm.model_executor.layers.quantization import QuantizationConfig
20 from vllm.model_executor.models.llama import (
21 LlamaAttention,
@@ -17,11 +23,57 @@ from vllm.model_executor.models.llama import (
23 LlamaForCausalLM,
24 LlamaModel,
25 )
26 +from vllm.sequence import IntermediateTensors
27 from vllm.v1.attention.backend import AttentionType
28
29 from .utils import AutoWeightsLoader
30
31
32 +class MistralMLP(nn.Module):
33 + def __init__(
34 + self,
35 + hidden_size: int,
36 + intermediate_size: int,
37 + hidden_act: str,
38 + quant_config: QuantizationConfig | None = None,
39 + bias: bool = False,
40 + gate_up_proj_bias: bool | None = None,
41 + prefix: str = "",
42 + reduce_results: bool = True,
43 + disable_tp: bool = False,
44 + ) -> None:
45 + super().__init__()
46 + gate_up_proj_bias = bias if gate_up_proj_bias is None else gate_up_proj_bias
47 + self.gate_up_proj = MergedColumnParallelLinear(
48 + input_size=hidden_size,
49 + output_sizes=[intermediate_size] * 2,
50 + bias=gate_up_proj_bias,
51 + quant_config=quant_config,
52 + disable_tp=disable_tp,
53 + prefix=f"{prefix}.gate_up_proj",
54 + )
55 + self.down_proj = RowParallelLinear(
56 + input_size=intermediate_size,
57 + output_size=hidden_size,
58 + bias=bias,
59 + quant_config=quant_config,
60 + reduce_results=reduce_results,
61 + disable_tp=disable_tp,
62 + prefix=f"{prefix}.down_proj",
63 + )
64 + if hidden_act != "silu":
65 + raise ValueError(
66 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."
67 + )
68 + self.act_fn = SiluAndMul()
69 +
70 + def forward(self, x):
71 + x, _ = self.gate_up_proj(x)
72 + x = self.act_fn(x)
73 + x, _ = self.down_proj(x)
74 + return x
75 +
76 +
77 class MistralAttention(LlamaAttention):
78 def __init__(
79 self,
@@ -114,6 +166,50 @@ class MistralDecoderLayer(LlamaDecoderLayer):
166 self.input_layernorm.quant_scaling_from = self.self_attn.qkv_proj
167 self.post_attention_layernorm.quant_scaling_from = self.mlp.gate_up_proj
168
169 + if getattr(config, "ada_rms_norm_t_cond", False):
170 + self.ada_rms_norm_t_cond = nn.Sequential(
171 + ColumnParallelLinear(
172 + input_size=config.hidden_size,
173 + output_size=config.ada_rms_norm_t_cond_dim,
174 + bias=False,
175 + return_bias=False,
176 + ),
177 + nn.GELU(),
178 + RowParallelLinear(
179 + input_size=config.ada_rms_norm_t_cond_dim,
180 + output_size=config.hidden_size,
181 + bias=False,
182 + return_bias=False,
183 + ),
184 + )
185 + else:
186 + self.ada_rms_norm_t_cond = None
187 +
188 + def forward(
189 + self,
190 + positions: torch.Tensor,
191 + hidden_states: torch.Tensor,
192 + residual: torch.Tensor | None,
193 + t_cond: torch.Tensor | None = None,
194 + ) -> tuple[torch.Tensor, torch.Tensor]:
195 + # Self Attention
196 + if residual is None:
197 + residual = hidden_states
198 + hidden_states = self.input_layernorm(hidden_states)
199 + else:
200 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
201 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
202 +
203 + # Fully Connected
204 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
205 +
206 + if self.ada_rms_norm_t_cond is not None:
207 + assert t_cond is not None
208 + hidden_states = hidden_states * (1 + self.ada_rms_norm_t_cond(t_cond))
209 +
210 + hidden_states = self.mlp(hidden_states)
211 + return hidden_states, residual
212 +
213
214 @support_torch_compile
215 class MistralModel(LlamaModel):
@@ -126,6 +222,18 @@ class MistralModel(LlamaModel):
222 ):
223 super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
224
225 + def forward(
226 + self,
227 + input_ids: torch.Tensor | None,
228 + positions: torch.Tensor,
229 + intermediate_tensors: IntermediateTensors | None,
230 + inputs_embeds: torch.Tensor | None = None,
231 + t_cond: torch.Tensor | None = None,
232 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
233 + return super().forward(
234 + input_ids, positions, intermediate_tensors, inputs_embeds, t_cond=t_cond
235 + )
236 +
237
238 class MistralForCausalLM(LlamaForCausalLM):
239 # Mistral: We don't support LoRA on the embedding layers.
vllm/model_executor/models/voxtral.py
+39 -4
@@ -4,7 +4,7 @@
4 import inspect
5 import math
6 from collections.abc import Iterable, Mapping, Sequence
7 -from functools import cached_property
7 +from functools import cached_property, partial
8 from math import ceil
9 from typing import Literal, cast
10
@@ -33,7 +33,11 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
33 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
34 from vllm.model_executor.models import SupportsPP
35 from vllm.model_executor.models.module_mapping import MultiModelKeys
36 -from vllm.model_executor.models.whisper import WhisperEncoder
36 +from vllm.model_executor.models.whisper import (
37 + WhisperEncoder,
38 + _create_fake_bias_for_k_proj,
39 +)
40 +from vllm.model_executor.models.whisper_causal import WhisperCausalEncoder
41 from vllm.multimodal import MULTIMODAL_REGISTRY
42 from vllm.multimodal.inputs import (
43 MultiModalDataDict,
@@ -543,6 +547,7 @@ class VoxtralForConditionalGeneration(
547 }
548 ).named_parameters()
549 )
550 + weights = _create_fake_bias_for_k_proj(weights, ".wk.weight")
551
552 loaded_weights = set()
553
@@ -730,6 +735,10 @@ class VoxtralEncoderModel(nn.Module):
735 r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)", # noqa: E501
736 r"whisper_encoder.layers.\1.mlp.fc2.\2",
737 ),
738 + (
739 + r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w3\.(weight|bias)",
740 + r"whisper_encoder.layers.\1.mlp.fc3.\2",
741 + ), # noqa: E501
742 (
743 r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)",
744 r"whisper_encoder.layers.\1.final_layer_norm.\2",
@@ -749,10 +758,15 @@ class VoxtralEncoderModel(nn.Module):
758 super().__init__()
759 self.config = cast(WhisperConfig, vllm_config.model_config.hf_config)
760 self.dtype: torch.dtype = vllm_config.model_config.dtype
752 - self.whisper_encoder = WhisperEncoder(
761 + self.is_causal = getattr(self.config, "is_causal", False)
762 + if self.is_causal:
763 + WhisperEncoderCls = WhisperCausalEncoder
764 + else:
765 + WhisperEncoderCls = partial(WhisperEncoder, init_in_fp32=True)
766 +
767 + self.whisper_encoder = WhisperEncoderCls(
768 vllm_config=vllm_config,
769 prefix=maybe_prefix(prefix, "whisper_encoder"),
755 - init_in_fp32=True,
770 )
771 mel_filters = mel_filter_bank(
772 num_frequency_bins=1 + self.config.window_size // 2,
@@ -843,6 +857,22 @@ class VoxtralEncoderModel(nn.Module):
857 ("qkv_proj", "k_proj", "k"),
858 ("qkv_proj", "v_proj", "v"),
859 ]
860 + params_mapping = []
861 +
862 + if self.is_causal:
863 + # For `WhisperCausalEncoder` we need
864 + # some more renaming
865 + stacked_params_mapping.extend(
866 + [
867 + (".mlp.gate_up_proj", ".mlp.fc1", 0),
868 + (".mlp.gate_up_proj", ".mlp.fc3", 1),
869 + ]
870 + )
871 + params_mapping.extend(
872 + [
873 + (".mlp.down_proj", ".mlp.fc2"),
874 + ]
875 + )
876 params_dict = dict(self.named_parameters())
877
878 name, loaded_weight = weight
@@ -860,6 +890,11 @@ class VoxtralEncoderModel(nn.Module):
890 weight_loader(param, loaded_weight, shard_id)
891 break
892 else:
893 + for param_name, weight_name in params_mapping:
894 + if weight_name not in name:
895 + continue
896 + name = name.replace(weight_name, param_name)
897 +
898 param = params_dict[name]
899 weight_loader = getattr(param, "weight_loader", default_weight_loader)
900 weight_loader(param, loaded_weight)
vllm/model_executor/models/voxtral_streaming.py
+21 -4
@@ -112,6 +112,18 @@ class TimeEmbedding(torch.nn.Module):
112 return torch.cat((emb.cos(), emb.sin()), dim=-1) # (B, D) or (B, T, D)
113
114
115 +def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor:
116 + # 1. Multiply by the scaling factor (e.g. 4)
117 + base = input_tensor * scaling
118 +
119 + # 2. Create the offsets, e.g. [0, 1, 2, 3]
120 + offsets = torch.arange(scaling, device=input_tensor.device)
121 +
122 + # 3. Use broadcasting, e.g. (N, 1) + (4,) results in (N, 4)
123 + # Then flatten back to 1D
124 + return (base.unsqueeze(1) + offsets).view(-1)
125 +
126 +
127 @MULTIMODAL_REGISTRY.register_processor(
128 VoxtralStreamingMultiModalProcessor,
129 info=VoxtralProcessingInfo,
@@ -175,8 +187,9 @@ class VoxtralStreamingGeneration(VoxtralForConditionalGeneration):
187 inputs_embeds.shape[0] * pool_size, inputs_embeds.shape[1] // pool_size
188 )
189
178 - audio_hidden_states = self.whisper_encoder.whisper_encoder.forward_layers(
179 - inputs_embeds
190 + whisper_positions = _expand_tensor(positions, pool_size)
191 + audio_hidden_states = self.whisper_encoder.whisper_encoder(
192 + inputs_embeds, whisper_positions
193 )
194
195 num_tokens, audio_hidden_size = audio_hidden_states.shape
@@ -197,10 +210,14 @@ class VoxtralStreamingGeneration(VoxtralForConditionalGeneration):
210 device=inputs_embeds.device,
211 dtype=inputs_embeds.dtype,
212 )
200 - inputs_embeds = inputs_embeds + self.time_embedding(time_tensor)
213 + t_cond = self.time_embedding(time_tensor)
214
215 hidden_states = self.language_model.model(
203 - input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
216 + input_ids,
217 + positions,
218 + intermediate_tensors,
219 + inputs_embeds=inputs_embeds,
220 + t_cond=t_cond,
221 )
222
223 return hidden_states
vllm/model_executor/models/whisper.py
+31 -72
@@ -5,7 +5,6 @@ import enum
5 import math
6 from collections.abc import Iterable, Mapping, Sequence
7 from contextlib import nullcontext
8 -from functools import partial
8 from typing import Annotated, Literal, cast
9
10 import numpy as np
@@ -39,8 +38,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39 from vllm.model_executor.models.whisper_utils import (
40 ISO639_1_SUPPORTED_LANGS,
42 - WhisperAttentionWithBlockPooling,
43 - WhisperCausalConv1d,
41 )
42 from vllm.multimodal import MULTIMODAL_REGISTRY
43 from vllm.multimodal.inputs import (
@@ -78,7 +75,7 @@ logger = init_logger(__name__)
75
76 class WhisperPosEmbedType(enum.Enum):
77 SINUSOIDAL = "sinusoidal"
81 - NOPE = "nope"
78 + ROPE = "rope"
79 LEARNED = "learned"
80
81
@@ -140,7 +137,6 @@ class WhisperAttention(nn.Module):
137 bias: bool = True,
138 attn_type: AttentionType = AttentionType.DECODER,
139 per_layer_sliding_window: int | None = None,
143 - block_pool_size: int = 1,
140 cache_config: CacheConfig | None = None,
141 quant_config: QuantizationConfig | None = None,
142 prefix: str = "",
@@ -199,14 +195,7 @@ class WhisperAttention(nn.Module):
195 attn_type=self.attn_type,
196 )
197 else: # AttentionType.DECODER (regular decoder self-attention)
202 - if block_pool_size > 1:
203 - attn_cls = partial(
204 - WhisperAttentionWithBlockPooling, block_pool_size=block_pool_size
205 - )
206 - else:
207 - attn_cls = Attention
208 -
209 - self.attn = attn_cls(
198 + self.attn = Attention(
199 self.num_heads,
200 self.head_dim,
201 self.scaling,
@@ -351,9 +340,7 @@ class WhisperEncoderLayer(nn.Module):
340 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
341 super().__init__()
342 config = vllm_config.model_config.hf_config
354 - is_causal = getattr(config, "is_causal", False)
343 sliding_window = getattr(config, "sliding_window", None)
356 - block_pool_size = getattr(config, "block_pool_size", 1)
344 cache_config = vllm_config.cache_config
345 quant_config = vllm_config.quant_config
346
@@ -361,8 +348,7 @@ class WhisperEncoderLayer(nn.Module):
348 self.self_attn = WhisperAttention(
349 embed_dim=self.embed_dim,
350 num_heads=config.encoder_attention_heads,
364 - attn_type=AttentionType.DECODER if is_causal else AttentionType.ENCODER,
365 - block_pool_size=block_pool_size,
351 + attn_type=AttentionType.ENCODER,
352 per_layer_sliding_window=sliding_window,
353 cache_config=cache_config,
354 quant_config=quant_config,
@@ -470,13 +456,8 @@ class WhisperEncoder(nn.Module):
456 self.max_source_positions = config.max_source_positions
457 self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
458
473 - self.is_causal = getattr(config, "is_causal", False)
474 - Conv1d = (
475 - WhisperCausalConv1d if self.is_causal else partial(nn.Conv1d, padding=1)
476 - )
477 -
478 - self.conv1 = Conv1d(self.num_mel_bins, embed_dim, kernel_size=3)
479 - self.conv2 = Conv1d(embed_dim, embed_dim, stride=2, kernel_size=3)
459 + self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
460 + self.conv2 = nn.Conv1d(embed_dim, embed_dim, stride=2, kernel_size=3, padding=1)
461
462 self.total_stride = self.conv1.stride[0] * self.conv2.stride[0]
463 self.start_layer, self.end_layer, self.layers = make_layers(
@@ -488,33 +469,29 @@ class WhisperEncoder(nn.Module):
469 )
470 self.layer_norm = nn.LayerNorm(config.d_model)
471
491 - if self.is_causal and self.pos_embed_type != WhisperPosEmbedType.NOPE:
492 - raise ValueError(
493 - "Only NOPE position embeddings are supported "
494 - f"for causal models, but got {self.pos_embed_type}"
495 - )
496 - elif self.pos_embed_type in (
472 + if self.pos_embed_type not in (
473 WhisperPosEmbedType.SINUSOIDAL,
474 WhisperPosEmbedType.LEARNED,
475 ):
500 - maybe_fp32_init_ctx = (
501 - set_default_torch_dtype(torch.float32)
502 - if init_in_fp32
503 - else nullcontext()
476 + raise ValueError(
477 + "Only sinusoidal or learned position embeddings are supported "
478 + f"for non-causal models, but got {self.pos_embed_type}"
479 + )
480 +
481 + maybe_fp32_init_ctx = (
482 + set_default_torch_dtype(torch.float32) if init_in_fp32 else nullcontext()
483 + )
484 +
485 + with (
486 + torch.no_grad(),
487 + maybe_fp32_init_ctx,
488 + ):
489 + self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
490 + self.embed_positions.weight.copy_(
491 + sinusoids(*self.embed_positions.weight.shape)
492 )
493
506 - with (
507 - torch.no_grad(),
508 - maybe_fp32_init_ctx,
509 - ):
510 - self.embed_positions = nn.Embedding(
511 - self.max_source_positions, embed_dim
512 - )
513 - self.embed_positions.weight.copy_(
514 - sinusoids(*self.embed_positions.weight.shape)
515 - )
516 -
517 - def forward_conv(
494 + def forward(
495 self, input_features: torch.Tensor | list[torch.Tensor]
496 ) -> torch.Tensor:
497 hidden_states = []
@@ -523,44 +500,26 @@ class WhisperEncoder(nn.Module):
500 embeds = nn.functional.gelu(self.conv1(features))
501 embeds = nn.functional.gelu(self.conv2(embeds))
502
526 - if self.pos_embed_type in (
527 - WhisperPosEmbedType.SINUSOIDAL,
528 - WhisperPosEmbedType.LEARNED,
529 - ):
530 - embeds = embeds.transpose(-1, -2)
531 - embeds = (
532 - embeds + self.embed_positions.weight[: embeds.size(-2), :]
533 - ).to(embeds.dtype)
534 - elif self.pos_embed_type == WhisperPosEmbedType.NOPE:
535 - embeds = embeds.transpose(-1, -2).to(embeds.dtype)
536 - else:
537 - raise ValueError(f"Unknown pos_embed_type: {self.pos_embed_type}")
503 + embeds = embeds.transpose(-1, -2)
504 + embeds = (embeds + self.embed_positions.weight[: embeds.size(-2), :]).to(
505 + embeds.dtype
506 + )
507
508 hidden_states.append(embeds)
509 input_is_batched = embeds.ndim > 2
510 # Input to MHA must be B x T x D
542 - if input_is_batched or self.is_causal:
511 + if input_is_batched:
512 # Models using WhisperEncoder may handle batching internally.
544 - # If WhisperEncoder is causal, sequences
545 - # are not padded to have identical seq length (T)
546 - # => concat over feature dim
513 hidden_states = torch.cat(hidden_states)
514 else:
515 hidden_states = torch.stack(hidden_states, dim=0)
516
551 - return hidden_states
552 -
553 - def forward_layers(self, hidden_states: torch.Tensor) -> torch.Tensor:
517 for encoder_layer in self.layers:
518 hidden_states = encoder_layer(hidden_states)
519
520 hidden_states = self.layer_norm(hidden_states)
521 return hidden_states
522
560 - def forward(self, input_features: torch.Tensor | list[torch.Tensor]):
561 - hidden_states = self.forward_conv(input_features)
562 - return self.forward_layers(hidden_states)
563 -
523
524 @support_torch_compile(dynamic_arg_dims={"input_ids": 0, "positions": -1})
525 class WhisperDecoder(nn.Module):
@@ -978,19 +937,19 @@ class WhisperForConditionalGeneration(
937 loader = AutoWeightsLoader(self, skip_prefixes=["proj_out."])
938
939 # add fake zeros bias for k_proj to state_dict
981 - weights = _create_fake_bias_for_k_proj(weights)
940 + weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
941 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
942
943
944 def _create_fake_bias_for_k_proj(
986 - weights: Iterable[tuple[str, torch.Tensor]],
945 + weights: Iterable[tuple[str, torch.Tensor]], fake_bias_key_name: str
946 ) -> Iterable[tuple[str, torch.Tensor]]:
947 """
948 Create full zeros bias for k_proj weight in self-attn and x-attn layers.
949 So that the bias for k_proj in qkv_proj can be initialized with zeros.
950 """
951 for name, weight in weights:
993 - if name.endswith(".k_proj.weight"):
952 + if name.endswith(fake_bias_key_name):
953 bias = torch.zeros(weight.size(0))
954 bias_name = name.replace("weight", "bias")
955 yield from [(name, weight), (bias_name, bias)]
vllm/model_executor/models/whisper_causal.py new
+465
@@ -0,0 +1,465 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import copy
4 +import functools
5 +import math
6 +from dataclasses import replace
7 +from functools import partial
8 +
9 +import torch
10 +import torch.nn.functional as F
11 +from torch import nn
12 +
13 +from vllm.attention.layer import Attention
14 +from vllm.config import CacheConfig, VllmConfig
15 +from vllm.distributed import get_tensor_model_parallel_world_size
16 +from vllm.model_executor.layers.layernorm import RMSNorm
17 +from vllm.model_executor.layers.linear import (
18 + QKVParallelLinear,
19 + RowParallelLinear,
20 +)
21 +from vllm.model_executor.layers.quantization import QuantizationConfig
22 +from vllm.model_executor.layers.rotary_embedding import get_rope
23 +from vllm.model_executor.models.mistral import MistralMLP
24 +from vllm.model_executor.models.whisper import WhisperPosEmbedType
25 +from vllm.v1.attention.backend import (
26 + AttentionBackend,
27 + AttentionMetadata,
28 + AttentionType,
29 + CommonAttentionMetadata,
30 + subclass_attention_backend_with_overrides,
31 +)
32 +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
33 +from vllm.v1.attention.selector import get_attn_backend
34 +from vllm.v1.kv_cache_interface import AttentionSpec
35 +
36 +from .utils import make_layers
37 +
38 +CausalRMSNorm = partial(RMSNorm, eps=1e-5)
39 +
40 +
41 +def _pad1d(
42 + x: torch.Tensor,
43 + paddings: tuple[int, int],
44 + mode: str = "constant",
45 + value: float = 0.0,
46 +) -> torch.Tensor:
47 + """Tiny wrapper around F.pad, just to allow for
48 + reflect padding on small input.
49 + If this is the case, we insert extra 0 padding
50 + to the right before the reflection happen.
51 + """
52 + length = x.shape[-1]
53 + padding_left, padding_right = paddings
54 + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
55 + if mode == "reflect":
56 + max_pad = max(padding_left, padding_right)
57 + extra_pad = 0
58 + if length <= max_pad:
59 + extra_pad = max_pad - length + 1
60 + x = F.pad(x, (0, extra_pad))
61 + padded = F.pad(x, paddings, mode, value)
62 + end = padded.shape[-1] - extra_pad
63 + return padded[..., :end]
64 + else:
65 + return F.pad(x, paddings, mode, value)
66 +
67 +
68 +class WhisperCausalConv1d(nn.Conv1d):
69 + def __init__(
70 + self,
71 + in_channels: int,
72 + out_channels: int,
73 + kernel_size: int,
74 + stride: int = 1,
75 + padding: int = 0,
76 + bias: bool = True,
77 + ) -> None:
78 + super().__init__(
79 + in_channels,
80 + out_channels,
81 + kernel_size,
82 + stride=stride,
83 + padding=padding,
84 + bias=bias,
85 + )
86 + self._stride = self.stride[0]
87 + self._effective_kernel_size = (kernel_size - 1) * self.dilation[0] + 1
88 + self._padding_total = self._effective_kernel_size - self._stride
89 +
90 + def forward(self, x: torch.Tensor) -> torch.Tensor:
91 + n_frames = (
92 + x.shape[-1] - self._effective_kernel_size + self._padding_total
93 + ) / self._stride + 1
94 + target_length = (math.ceil(n_frames) - 1) * self._stride + (
95 + self._effective_kernel_size - self._padding_total
96 + )
97 + extra_padding = target_length - x.shape[-1]
98 + x = _pad1d(x, (self._padding_total, extra_padding), mode="constant")
99 + return super().forward(x)
100 +
101 +
102 +@functools.lru_cache
103 +def create_whisper_attention_backend_with_block_pooling(
104 + underlying_attn_backend: AttentionBackend, block_pool_size: int
105 +) -> type[AttentionBackend]:
106 + prefix = "WhisperCausalAttentionWithBlockPooling_"
107 + underlying_builder = underlying_attn_backend.get_builder_cls()
108 +
109 + class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore
110 + def __init__(
111 + self,
112 + kv_cache_spec: AttentionSpec,
113 + layer_names: list[str],
114 + vllm_config: VllmConfig,
115 + device: torch.device,
116 + ):
117 + assert kv_cache_spec.num_kv_heads % block_pool_size == 0
118 + kv_cache_spec = replace(
119 + kv_cache_spec,
120 + block_size=kv_cache_spec.block_size * block_pool_size,
121 + num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
122 + )
123 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
124 +
125 + def build(
126 + self,
127 + common_prefix_len: int,
128 + common_attn_metadata: CommonAttentionMetadata,
129 + fast_build: bool = False,
130 + ) -> AttentionMetadata:
131 + new_common_attn_metadata = copy.deepcopy(common_attn_metadata)
132 + new_common_attn_metadata.query_start_loc *= block_pool_size
133 + new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
134 + new_common_attn_metadata.seq_lens *= block_pool_size
135 + new_common_attn_metadata._seq_lens_cpu *= block_pool_size
136 + new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
137 + new_common_attn_metadata.num_actual_tokens *= block_pool_size
138 + new_common_attn_metadata.max_query_len *= block_pool_size
139 + new_common_attn_metadata.max_seq_len *= block_pool_size
140 + original_slot_mapping = common_attn_metadata.slot_mapping
141 + common_prefix_len *= block_pool_size
142 + new_common_attn_metadata.slot_mapping = (
143 + (
144 + original_slot_mapping.unsqueeze(1) * block_pool_size
145 + + torch.arange(block_pool_size, device=original_slot_mapping.device)
146 + )
147 + .flatten()
148 + .clamp(min=-1)
149 + )
150 + return super().build(
151 + common_prefix_len, new_common_attn_metadata, fast_build
152 + )
153 +
154 + if not issubclass(underlying_attn_backend, FlashAttentionBackend):
155 + raise NotImplementedError(
156 + f"{underlying_attn_backend} is not yet supported."
157 + "Contributions to support more backends are much "
158 + "appreciated."
159 + )
160 +
161 + attn_backend = subclass_attention_backend_with_overrides(
162 + name_prefix=prefix,
163 + attention_backend_cls=underlying_attn_backend,
164 + overrides={
165 + "get_builder_cls": lambda: WhisperCausalAttentionWithBlockPoolingBuilder,
166 + "get_kv_cache_shape": lambda num_blocks,
167 + block_size,
168 + num_kv_heads,
169 + head_size,
170 + cache_dtype_str: (
171 + 2,
172 + num_blocks,
173 + # we stretch each block by `block_pool_size`
174 + block_size * block_pool_size,
175 + num_kv_heads // block_pool_size,
176 + head_size,
177 + ), # TODO: generalize to other backends
178 + },
179 + )
180 +
181 + return attn_backend
182 +
183 +
184 +class WhisperCausalAttentionWithBlockPooling(Attention):
185 + """Attention layer with block pooling."""
186 +
187 + def __init__(
188 + self,
189 + num_heads: int,
190 + head_size: int,
191 + scale: float,
192 + num_kv_heads: int | None = None,
193 + alibi_slopes: list[float] | None = None,
194 + cache_config: CacheConfig | None = None,
195 + quant_config: QuantizationConfig | None = None,
196 + logits_soft_cap: float | None = None,
197 + per_layer_sliding_window: int | None = None,
198 + prefix: str = "",
199 + attn_type: str = AttentionType.DECODER,
200 + kv_sharing_target_layer_name: str | None = None,
201 + block_pool_size: int = 1,
202 + attn_backend: type[AttentionBackend] | None = None,
203 + **extra_impl_args,
204 + ) -> None:
205 + self.block_pool_size = block_pool_size
206 + dtype = torch.get_default_dtype()
207 +
208 + if cache_config is not None:
209 + kv_cache_dtype = cache_config.cache_dtype
210 + block_size = cache_config.block_size
211 + else:
212 + kv_cache_dtype = "auto"
213 + block_size = 16
214 +
215 + underlying_attn_backend = get_attn_backend(
216 + head_size,
217 + dtype,
218 + kv_cache_dtype,
219 + block_size,
220 + attn_type=attn_type,
221 + )
222 + attn_backend = create_whisper_attention_backend_with_block_pooling(
223 + underlying_attn_backend, block_pool_size
224 + )
225 +
226 + super().__init__(
227 + num_heads=num_heads,
228 + head_size=head_size,
229 + scale=scale,
230 + num_kv_heads=num_kv_heads,
231 + alibi_slopes=alibi_slopes,
232 + cache_config=cache_config,
233 + quant_config=quant_config,
234 + logits_soft_cap=logits_soft_cap,
235 + per_layer_sliding_window=per_layer_sliding_window,
236 + prefix=prefix,
237 + attn_type=attn_type,
238 + kv_sharing_target_layer_name=kv_sharing_target_layer_name,
239 + attn_backend=attn_backend,
240 + **extra_impl_args,
241 + )
242 +
243 + def get_kv_cache_spec(self, vllm_config: VllmConfig):
244 + kv_cache_spec = super().get_kv_cache_spec(vllm_config)
245 + assert isinstance(kv_cache_spec, AttentionSpec)
246 + kv_cache_spec = replace(
247 + kv_cache_spec,
248 + num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
249 + )
250 + return kv_cache_spec
251 +
252 +
253 +class WhisperCausalAttention(nn.Module):
254 + def __init__(
255 + self,
256 + embed_dim: int,
257 + num_heads: int,
258 + head_dim: int,
259 + max_position_embeddings: int,
260 + bias: bool = True,
261 + attn_type: AttentionType = AttentionType.DECODER,
262 + per_layer_sliding_window: int | None = None,
263 + block_pool_size: int = 1,
264 + cache_config: CacheConfig | None = None,
265 + quant_config: QuantizationConfig | None = None,
266 + prefix: str = "",
267 + ):
268 + super().__init__()
269 + self.embed_dim = embed_dim
270 + tp_size = get_tensor_model_parallel_world_size()
271 + self.total_num_heads = num_heads
272 + assert self.total_num_heads % tp_size == 0
273 + self.num_heads = self.total_num_heads // tp_size
274 + if self.total_num_heads >= tp_size:
275 + # Number of heads is greater than TP size, so we partition
276 + # the KV heads across multiple tensor parallel GPUs.
277 + assert self.total_num_heads % tp_size == 0
278 + else:
279 + # Number of heads is less than TP size, so we replicate
280 + # the KV heads across multiple tensor parallel GPUs.
281 + assert tp_size % self.total_num_heads == 0
282 + self.num_kv_heads = max(1, self.total_num_heads // tp_size)
283 + self.head_dim = head_dim
284 + self.q_size = self.num_heads * self.head_dim
285 + self.kv_size = self.num_kv_heads * self.head_dim
286 + self.attn_type = attn_type
287 +
288 + self.scaling = self.head_dim**-0.5
289 +
290 + self._init_qkv(embed_dim, bias, quant_config, prefix=prefix)
291 + self.out_proj = RowParallelLinear(
292 + input_size=self.total_num_heads * self.head_dim,
293 + output_size=embed_dim,
294 + bias=bias,
295 + quant_config=quant_config,
296 + prefix=f"{prefix}.out_proj",
297 + )
298 + assert block_pool_size > 1, (
299 + f"Causal attention only supports block_pool_size>1, not {block_pool_size}."
300 + )
301 + self.attn = WhisperCausalAttentionWithBlockPooling(
302 + self.num_heads,
303 + self.head_dim,
304 + self.scaling,
305 + num_kv_heads=self.num_kv_heads,
306 + cache_config=cache_config,
307 + quant_config=quant_config,
308 + prefix=f"{prefix}.attn",
309 + attn_type=AttentionType.DECODER,
310 + per_layer_sliding_window=per_layer_sliding_window,
311 + block_pool_size=block_pool_size,
312 + )
313 +
314 + assert per_layer_sliding_window is not None, (
315 + "rope can only used in combination with a sliding window"
316 + )
317 + self._init_rotary_emb(max_position_embeddings)
318 +
319 + def _init_rotary_emb(self, max_position_embeddings: int) -> None:
320 + self.rotary_emb = get_rope(
321 + self.head_dim,
322 + max_position=max_position_embeddings,
323 + is_neox_style=False,
324 + rope_parameters={"rope_theta": 1e6},
325 + )
326 +
327 + def _init_qkv(
328 + self,
329 + embed_dim: int,
330 + bias: bool = True,
331 + quant_config: QuantizationConfig | None = None,
332 + prefix: str = "",
333 + ) -> None:
334 + self.qkv_proj = QKVParallelLinear(
335 + hidden_size=embed_dim,
336 + head_size=self.head_dim,
337 + total_num_heads=self.total_num_heads,
338 + total_num_kv_heads=self.total_num_heads,
339 + bias=bias,
340 + quant_config=quant_config,
341 + prefix=f"{prefix}.qkv_proj",
342 + )
343 +
344 + def forward(
345 + self,
346 + hidden_states: torch.Tensor,
347 + positions: torch.Tensor | None = None,
348 + ):
349 + qkv, _ = self.qkv_proj(hidden_states)
350 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
351 +
352 + assert positions is not None
353 + q, k = self.rotary_emb(positions, q, k)
354 +
355 + attn_output = self.attn(q, k, v)
356 +
357 + output, _ = self.out_proj(attn_output)
358 +
359 + return output
360 +
361 +
362 +class WhisperCausalEncoderLayer(nn.Module):
363 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
364 + super().__init__()
365 + config = vllm_config.model_config.hf_config
366 + sliding_window = getattr(config, "sliding_window", None)
367 + block_pool_size = config.block_pool_size
368 + assert block_pool_size > 1
369 +
370 + cache_config = vllm_config.cache_config
371 + quant_config = vllm_config.quant_config
372 +
373 + self.embed_dim = config.d_model
374 + self.head_dim = self.embed_dim // config.encoder_attention_heads
375 + self.self_attn = WhisperCausalAttention(
376 + embed_dim=self.embed_dim,
377 + num_heads=config.encoder_attention_heads,
378 + head_dim=config.encoder_head_dim,
379 + max_position_embeddings=config.max_position_embeddings,
380 + block_pool_size=block_pool_size,
381 + per_layer_sliding_window=sliding_window,
382 + cache_config=cache_config,
383 + quant_config=quant_config,
384 + prefix=f"{prefix}.self_attn",
385 + )
386 + self.self_attn_layer_norm = CausalRMSNorm(self.embed_dim)
387 +
388 + self.mlp = MistralMLP(
389 + hidden_size=config.d_model,
390 + intermediate_size=config.encoder_ffn_dim,
391 + hidden_act="silu",
392 + quant_config=quant_config,
393 + bias=True,
394 + gate_up_proj_bias=False,
395 + prefix=f"{prefix}.mlp",
396 + )
397 + self.final_layer_norm = CausalRMSNorm(self.embed_dim)
398 +
399 + def forward(
400 + self,
401 + hidden_states: torch.Tensor,
402 + positions: torch.Tensor | None = None,
403 + ):
404 + residual = hidden_states
405 + hidden_states = self.self_attn_layer_norm(hidden_states)
406 + hidden_states = self.self_attn(hidden_states=hidden_states, positions=positions)
407 + hidden_states = residual + hidden_states
408 + residual = hidden_states
409 + hidden_states = self.final_layer_norm(hidden_states)
410 + hidden_states = self.mlp(hidden_states)
411 + hidden_states = residual + hidden_states
412 +
413 + return hidden_states
414 +
415 +
416 +class WhisperCausalEncoder(nn.Module):
417 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
418 + super().__init__()
419 + config = vllm_config.model_config.hf_config
420 + embed_dim = config.d_model
421 +
422 + assert WhisperPosEmbedType(config.pos_embed) == WhisperPosEmbedType.ROPE
423 + assert config.is_causal
424 +
425 + self.num_mel_bins = config.num_mel_bins
426 + self.max_source_positions = config.max_source_positions
427 + self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
428 +
429 + self.conv1 = WhisperCausalConv1d(self.num_mel_bins, embed_dim, kernel_size=3)
430 + self.conv2 = WhisperCausalConv1d(embed_dim, embed_dim, stride=2, kernel_size=3)
431 +
432 + self.total_stride = self.conv1.stride[0] * self.conv2.stride[0]
433 + self.start_layer, self.end_layer, self.layers = make_layers(
434 + config.encoder_layers,
435 + lambda prefix: WhisperCausalEncoderLayer(
436 + vllm_config=vllm_config, prefix=f"{prefix}.layers"
437 + ),
438 + prefix=f"{prefix}.layers",
439 + )
440 + self.layer_norm = CausalRMSNorm(config.d_model)
441 +
442 + def forward_conv(
443 + self, input_features: torch.Tensor | list[torch.Tensor]
444 + ) -> torch.Tensor:
445 + hidden_states = []
446 + for features in input_features:
447 + embeds = nn.functional.gelu(self.conv1(features))
448 + embeds = nn.functional.gelu(self.conv2(embeds))
449 +
450 + embeds = embeds.transpose(-1, -2).to(embeds.dtype)
451 +
452 + hidden_states.append(embeds)
453 +
454 + hidden_states = torch.cat(hidden_states)
455 +
456 + return hidden_states
457 +
458 + def forward(
459 + self, hidden_states: torch.Tensor, positions: torch.Tensor
460 + ) -> torch.Tensor:
461 + for encoder_layer in self.layers:
462 + hidden_states = encoder_layer(hidden_states, positions)
463 +
464 + hidden_states = self.layer_norm(hidden_states)
465 + return hidden_states
vllm/model_executor/models/whisper_utils.py
-233
@@ -1,27 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import copy
4 -import functools
5 -import math
6 -from dataclasses import replace
3
8 -import torch
9 -import torch.nn.functional as F
10 -from torch import nn
11 -
12 -from vllm.attention.layer import Attention
13 -from vllm.config import CacheConfig, VllmConfig
14 -from vllm.model_executor.layers.quantization import QuantizationConfig
15 -from vllm.v1.attention.backend import (
16 - AttentionBackend,
17 - AttentionMetadata,
18 - AttentionType,
19 - CommonAttentionMetadata,
20 - subclass_attention_backend_with_overrides,
21 -)
22 -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
23 -from vllm.v1.attention.selector import get_attn_backend
24 -from vllm.v1.kv_cache_interface import AttentionSpec
4
5 # From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
6 ISO639_1_SUPPORTED_LANGS = {
@@ -83,215 +62,3 @@ ISO639_1_SUPPORTED_LANGS = {
62 "vi": "Vietnamese",
63 "cy": "Welsh",
64 }
86 -
87 -
88 -def _pad1d(
89 - x: torch.Tensor,
90 - paddings: tuple[int, int],
91 - mode: str = "constant",
92 - value: float = 0.0,
93 -) -> torch.Tensor:
94 - """Tiny wrapper around F.pad, just to allow for
95 - reflect padding on small input.
96 - If this is the case, we insert extra 0 padding
97 - to the right before the reflection happen.
98 - """
99 - length = x.shape[-1]
100 - padding_left, padding_right = paddings
101 - assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
102 - if mode == "reflect":
103 - max_pad = max(padding_left, padding_right)
104 - extra_pad = 0
105 - if length <= max_pad:
106 - extra_pad = max_pad - length + 1
107 - x = F.pad(x, (0, extra_pad))
108 - padded = F.pad(x, paddings, mode, value)
109 - end = padded.shape[-1] - extra_pad
110 - return padded[..., :end]
111 - else:
112 - return F.pad(x, paddings, mode, value)
113 -
114 -
115 -class WhisperCausalConv1d(nn.Conv1d):
116 - def __init__(
117 - self,
118 - in_channels: int,
119 - out_channels: int,
120 - kernel_size: int,
121 - stride: int = 1,
122 - padding: int = 0,
123 - bias: bool = True,
124 - ) -> None:
125 - super().__init__(
126 - in_channels,
127 - out_channels,
128 - kernel_size,
129 - stride=stride,
130 - padding=padding,
131 - bias=bias,
132 - )
133 - self._stride = self.stride[0]
134 - self._effective_kernel_size = (kernel_size - 1) * self.dilation[0] + 1
135 - self._padding_total = self._effective_kernel_size - self._stride
136 -
137 - def forward(self, x: torch.Tensor) -> torch.Tensor:
138 - n_frames = (
139 - x.shape[-1] - self._effective_kernel_size + self._padding_total
140 - ) / self._stride + 1
141 - target_length = (math.ceil(n_frames) - 1) * self._stride + (
142 - self._effective_kernel_size - self._padding_total
143 - )
144 - extra_padding = target_length - x.shape[-1]
145 - x = _pad1d(x, (self._padding_total, extra_padding), mode="constant")
146 - return super().forward(x)
147 -
148 -
149 -@functools.lru_cache
150 -def create_whisper_attention_backend_with_block_pooling(
151 - underlying_attn_backend: AttentionBackend, block_pool_size: int
152 -) -> type[AttentionBackend]:
153 - prefix = "WhisperAttentionWithBlockPooling_"
154 - underlying_builder = underlying_attn_backend.get_builder_cls()
155 -
156 - class WhisperAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore
157 - def __init__(
158 - self,
159 - kv_cache_spec: AttentionSpec,
160 - layer_names: list[str],
161 - vllm_config: VllmConfig,
162 - device: torch.device,
163 - ):
164 - assert kv_cache_spec.num_kv_heads % block_pool_size == 0
165 - kv_cache_spec = replace(
166 - kv_cache_spec,
167 - block_size=kv_cache_spec.block_size * block_pool_size,
168 - num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
169 - )
170 - super().__init__(kv_cache_spec, layer_names, vllm_config, device)
171 -
172 - def build(
173 - self,
174 - common_prefix_len: int,
175 - common_attn_metadata: CommonAttentionMetadata,
176 - fast_build: bool = False,
177 - ) -> AttentionMetadata:
178 - new_common_attn_metadata = copy.deepcopy(common_attn_metadata)
179 - new_common_attn_metadata.query_start_loc *= block_pool_size
180 - new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
181 - new_common_attn_metadata.seq_lens *= block_pool_size
182 - new_common_attn_metadata._seq_lens_cpu *= block_pool_size
183 - new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
184 - new_common_attn_metadata.num_actual_tokens *= block_pool_size
185 - new_common_attn_metadata.max_query_len *= block_pool_size
186 - new_common_attn_metadata.max_seq_len *= block_pool_size
187 - original_slot_mapping = common_attn_metadata.slot_mapping
188 - common_prefix_len *= block_pool_size
189 - new_common_attn_metadata.slot_mapping = (
190 - (
191 - original_slot_mapping.unsqueeze(1) * block_pool_size
192 - + torch.arange(block_pool_size, device=original_slot_mapping.device)
193 - )
194 - .flatten()
195 - .clamp(min=-1)
196 - )
197 - return super().build(
198 - common_prefix_len, new_common_attn_metadata, fast_build
199 - )
200 -
201 - if not issubclass(underlying_attn_backend, FlashAttentionBackend):
202 - raise NotImplementedError(
203 - f"{underlying_attn_backend} is not yet supported."
204 - "Contributions to support more backends are much "
205 - "appreciated."
206 - )
207 -
208 - attn_backend = subclass_attention_backend_with_overrides(
209 - name_prefix=prefix,
210 - attention_backend_cls=underlying_attn_backend,
211 - overrides={
212 - "get_builder_cls": lambda: WhisperAttentionWithBlockPoolingBuilder,
213 - "get_kv_cache_shape": lambda num_blocks,
214 - block_size,
215 - num_kv_heads,
216 - head_size,
217 - cache_dtype_str: (
218 - 2,
219 - num_blocks,
220 - # we stretch each block by `block_pool_size`
221 - block_size * block_pool_size,
222 - num_kv_heads // block_pool_size,
223 - head_size,
224 - ), # TODO: generalize to other backends
225 - },
226 - )
227 -
228 - return attn_backend
229 -
230 -
231 -class WhisperAttentionWithBlockPooling(Attention):
232 - """Attention layer with block pooling."""
233 -
234 - def __init__(
235 - self,
236 - num_heads: int,
237 - head_size: int,
238 - scale: float,
239 - num_kv_heads: int | None = None,
240 - alibi_slopes: list[float] | None = None,
241 - cache_config: CacheConfig | None = None,
242 - quant_config: QuantizationConfig | None = None,
243 - logits_soft_cap: float | None = None,
244 - per_layer_sliding_window: int | None = None,
245 - prefix: str = "",
246 - attn_type: str = AttentionType.DECODER,
247 - kv_sharing_target_layer_name: str | None = None,
248 - block_pool_size: int = 1,
249 - attn_backend: type[AttentionBackend] | None = None,
250 - **extra_impl_args,
251 - ) -> None:
252 - self.block_pool_size = block_pool_size
253 - dtype = torch.get_default_dtype()
254 -
255 - if cache_config is not None:
256 - kv_cache_dtype = cache_config.cache_dtype
257 - block_size = cache_config.block_size
258 - else:
259 - kv_cache_dtype = "auto"
260 - block_size = 16
261 -
262 - underlying_attn_backend = get_attn_backend(
263 - head_size,
264 - dtype,
265 - kv_cache_dtype,
266 - block_size,
267 - attn_type=attn_type,
268 - )
269 - attn_backend = create_whisper_attention_backend_with_block_pooling(
270 - underlying_attn_backend, block_pool_size
271 - )
272 -
273 - super().__init__(
274 - num_heads=num_heads,
275 - head_size=head_size,
276 - scale=scale,
277 - num_kv_heads=num_kv_heads,
278 - alibi_slopes=alibi_slopes,
279 - cache_config=cache_config,
280 - quant_config=quant_config,
281 - logits_soft_cap=logits_soft_cap,
282 - per_layer_sliding_window=per_layer_sliding_window,
283 - prefix=prefix,
284 - attn_type=attn_type,
285 - kv_sharing_target_layer_name=kv_sharing_target_layer_name,
286 - attn_backend=attn_backend,
287 - **extra_impl_args,
288 - )
289 -
290 - def get_kv_cache_spec(self, vllm_config: VllmConfig):
291 - kv_cache_spec = super().get_kv_cache_spec(vllm_config)
292 - assert isinstance(kv_cache_spec, AttentionSpec)
293 - kv_cache_spec = replace(
294 - kv_cache_spec,
295 - num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
296 - )
297 - return kv_cache_spec
vllm/transformers_utils/configs/mistral.py
+3
@@ -224,6 +224,7 @@ def _remap_mistral_audio_args(config: dict) -> dict:
224 encoder_layers=encoder_args["n_layers"],
225 encoder_ffn_dim=encoder_args["hidden_dim"],
226 encoder_attention_heads=encoder_args["n_heads"],
227 + encoder_head_dim=encoder_args["head_dim"],
228 vocab_size=encoder_args["vocab_size"],
229 max_source_positions=encoder_args["max_source_positions"],
230 is_encoder_decoder=False, # Override WhisperConfig default
@@ -231,6 +232,8 @@ def _remap_mistral_audio_args(config: dict) -> dict:
232 sliding_window=sliding_window,
233 block_pool_size=block_pool_size,
234 pos_embed=encoder_args.get("pos_embed", "sinusoidal"),
235 + # only needed for RoPE
236 + max_position_embeddings=block_pool_size * config["max_position_embeddings"],
237 ),
238 }
239 if quant_config: