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, 14:45 UTC
[Bugfix] Fix Voxtral streaming slot_mapping (#33073)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jan 26, 2026 at 19:40 UTC
19ab0f7ce56af548d54632cfc278aefa894d797e
1 file changed
+40
vllm/model_executor/models/whisper_causal.py
+40
@@ -105,6 +105,7 @@ def create_whisper_attention_backend_with_block_pooling(
105
) -> type[AttentionBackend]:
106
prefix = "WhisperCausalAttentionWithBlockPooling_"
107
underlying_builder = underlying_attn_backend.get_builder_cls()
108
+ underlying_impl = underlying_attn_backend.get_impl_cls()
109
110
class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore
111
def __init__(
@@ -151,6 +152,43 @@ def create_whisper_attention_backend_with_block_pooling(
152
common_prefix_len, new_common_attn_metadata, fast_build
153
)
154
155
+ # NOTE: We need a custom impl so we can use the transformed slot_mapping
156
+ # computed by `WhisperCausalAttentionWithBlockPoolingBuilder` instead of
157
+ # the one from `forward_context.slot_mapping` (gpu_model_runner).
158
+ # This follows the same pattern as CrossAttentionImpl.
159
+ class WhisperCausalAttentionWithBlockPoolingImpl(underlying_impl): # type: ignore[valid-type,misc]
160
+ def forward(
161
+ self,
162
+ layer: torch.nn.Module,
163
+ query: torch.Tensor,
164
+ key: torch.Tensor,
165
+ value: torch.Tensor,
166
+ kv_cache: torch.Tensor,
167
+ attn_metadata: AttentionMetadata,
168
+ output: torch.Tensor | None = None,
169
+ output_scale: torch.Tensor | None = None,
170
+ output_block_scale: torch.Tensor | None = None,
171
+ ) -> torch.Tensor:
172
+ if (
173
+ not underlying_attn_backend.forward_includes_kv_cache_update
174
+ and attn_metadata is not None
175
+ ):
176
+ self.do_kv_cache_update(
177
+ layer, key, value, kv_cache, attn_metadata.slot_mapping
178
+ )
179
+
180
+ return super().forward(
181
+ layer,
182
+ query,
183
+ key,
184
+ value,
185
+ kv_cache,
186
+ attn_metadata,
187
+ output,
188
+ output_scale,
189
+ output_block_scale,
190
+ )
191
+
192
if not issubclass(underlying_attn_backend, FlashAttentionBackend):
193
raise NotImplementedError(
194
f"{underlying_attn_backend} is not yet supported."
@@ -163,6 +201,7 @@ def create_whisper_attention_backend_with_block_pooling(
201
attention_backend_cls=underlying_attn_backend,
202
overrides={
203
"get_builder_cls": lambda: WhisperCausalAttentionWithBlockPoolingBuilder,
204
+ "get_impl_cls": lambda: WhisperCausalAttentionWithBlockPoolingImpl,
205
"get_kv_cache_shape": lambda num_blocks,
206
block_size,
207
num_kv_heads,
@@ -175,6 +214,7 @@ def create_whisper_attention_backend_with_block_pooling(
214
num_kv_heads // block_pool_size,
215
head_size,
216
), # TODO: generalize to other backends
217
+ "forward_includes_kv_cache_update": True,
218
},
219
)
220